Skip to content

Commit

Permalink
By default only accept requests where origin header matches the host.
Browse files Browse the repository at this point in the history
Browsers are dumb and let any website do requests to localhost this should
prevent this without breaking things. CORS prevents the javascript from
reading the response but they can still write it.

At the moment this is only enabled when the --enable-cors-header argument
is not used.
  • Loading branch information
comfyanonymous committed Sep 8, 2024
1 parent 9c5fca7 commit 3ab3516
Showing 1 changed file with 23 additions and 0 deletions.
23 changes: 23 additions & 0 deletions server.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,27 @@ async def cors_middleware(request: web.Request, handler):

return cors_middleware

def create_origin_only_middleware():
@web.middleware
async def origin_only_middleware(request: web.Request, handler):
if request.method == "OPTIONS":
response = web.Response()
else:
response = await handler(request)

if 'Host' in request.headers and 'Origin' in request.headers:
host = request.headers['Host']
origin = request.headers['Origin']
host_domain = host.lower()
origin_domain = urllib.parse.urlparse(origin).netloc.lower()
if host_domain != origin_domain:
logging.warning("WARNING: request with non matching host and origin {} != {}, returning 403".format(host_domain, origin_domain))
return web.Response(status=403)

return response

return origin_only_middleware

class PromptServer():
def __init__(self, loop):
PromptServer.instance = self
Expand All @@ -99,6 +120,8 @@ def __init__(self, loop):
middlewares = [cache_control]
if args.enable_cors_header:
middlewares.append(create_cors_middleware(args.enable_cors_header))
else:
middlewares.append(create_origin_only_middleware())

max_upload_size = round(args.max_upload_size * 1024 * 1024)
self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares)
Expand Down

0 comments on commit 3ab3516

Please sign in to comment.