vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

prevent prompt leakage during usage of check endpoint when genkey is provided in multiuser mode

Browse source
This commit is contained in:
Concedo 2023-12-24 17:08:43 +08:00
parent 71a5afaab5
commit bc24c9334c
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
koboldcpp.py
View file

@ -716,7 +716,7 @@ Enter Prompt:<br>
self.wfile.write(finalhtml) self.wfile.write(finalhtml)
def do_GET(self): def do_GET(self):
global maxctx, maxhordelen, friendlymodelname, KcppVersion, totalgens, preloaded_story, exitcounter global maxctx, maxhordelen, friendlymodelname, KcppVersion, totalgens, preloaded_story, exitcounter, currentusergenkey
self.path = self.path.rstrip('/') self.path = self.path.rstrip('/')
response_body = None response_body = None
content_type = 'application/json' content_type = 'application/json'
@ -766,7 +766,7 @@ Enter Prompt:<br>
elif self.path.endswith('/api/extra/generate/check'): elif self.path.endswith('/api/extra/generate/check'):
pendtxtStr = "" pendtxtStr = ""
if requestsinqueue==0 and totalgens>0: if requestsinqueue==0 and totalgens>0 and (not args.multiuser or currentusergenkey==""):
pendtxt = handle.get_pending_output() pendtxt = handle.get_pending_output()
pendtxtStr = ctypes.string_at(pendtxt).decode("UTF-8","ignore") pendtxtStr = ctypes.string_at(pendtxt).decode("UTF-8","ignore")
response_body = (json.dumps({"results": [{"text": pendtxtStr}]}).encode()) response_body = (json.dumps({"results": [{"text": pendtxtStr}]}).encode())
@ -861,7 +861,7 @@ Enter Prompt:<br>
multiuserkey = "" multiuserkey = ""
if totalgens>0: if totalgens>0:
if (multiuserkey=="" and requestsinqueue==0) or (multiuserkey!="" and multiuserkey==currentusergenkey): if (multiuserkey=="" and multiuserkey==currentusergenkey) or (multiuserkey=="" and requestsinqueue==0 and not args.multiuser) or (multiuserkey!="" and multiuserkey==currentusergenkey): #avoid leaking prompts in multiuser
pendtxt = handle.get_pending_output() pendtxt = handle.get_pending_output()
pendtxtStr = ctypes.string_at(pendtxt).decode("UTF-8","ignore") pendtxtStr = ctypes.string_at(pendtxt).decode("UTF-8","ignore")
response_body = (json.dumps({"results": [{"text": pendtxtStr}]}).encode()) response_body = (json.dumps({"results": [{"text": pendtxtStr}]}).encode())