mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-18 23:38:29 +08:00
Amp-Thread-ID: https://ampcode.com/threads/T-019ff807-a342-7628-863e-847d1dad3d5d Co-authored-by: Amp <amp@ampcode.com>
79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
import ipaddress
|
|
import logging
|
|
import socket
|
|
import urllib
|
|
|
|
from aiohttp import web
|
|
|
|
|
|
def is_loopback(host):
|
|
if host is None:
|
|
return False
|
|
try:
|
|
if ipaddress.ip_address(host).is_loopback:
|
|
return True
|
|
else:
|
|
return False
|
|
except:
|
|
pass
|
|
|
|
loopback = False
|
|
for family in (socket.AF_INET, socket.AF_INET6):
|
|
try:
|
|
r = socket.getaddrinfo(host, None, family, socket.SOCK_STREAM)
|
|
for family, _, _, _, sockaddr in r:
|
|
if not ipaddress.ip_address(sockaddr[0]).is_loopback:
|
|
return loopback
|
|
else:
|
|
loopback = True
|
|
except socket.gaierror:
|
|
pass
|
|
|
|
return loopback
|
|
|
|
|
|
def create_origin_only_middleware():
|
|
@web.middleware
|
|
async def origin_only_middleware(request: web.Request, handler):
|
|
if request.headers.get('Sec-Fetch-Site') == 'cross-site':
|
|
is_top_level_navigation = (
|
|
request.method == 'GET'
|
|
and request.path == '/'
|
|
and request.headers.get('Sec-Fetch-Mode') == 'navigate'
|
|
and request.headers.get('Sec-Fetch-Dest') == 'document'
|
|
)
|
|
if not is_top_level_navigation:
|
|
return web.Response(status=403)
|
|
#this code is used to prevent the case where a random website can queue comfy workflows by making a POST to 127.0.0.1 which browsers don't prevent for some dumb reason.
|
|
#in that case the Host and Origin hostnames won't match
|
|
#I know the proper fix would be to add a cookie but this should take care of the problem in the meantime
|
|
if 'Host' in request.headers and 'Origin' in request.headers:
|
|
host = request.headers['Host']
|
|
origin = request.headers['Origin']
|
|
host_domain = host.lower()
|
|
parsed = urllib.parse.urlparse(origin)
|
|
origin_domain = parsed.netloc.lower()
|
|
host_domain_parsed = urllib.parse.urlsplit('//' + host_domain)
|
|
|
|
#limit the check to when the host domain is localhost, this makes it slightly less safe but should still prevent the exploit
|
|
loopback = is_loopback(host_domain_parsed.hostname)
|
|
|
|
if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers, same for host
|
|
host_domain = host_domain_parsed.hostname
|
|
if host_domain_parsed.port is None:
|
|
origin_domain = parsed.hostname
|
|
|
|
if loopback and host_domain is not None and origin_domain is not None and len(host_domain) > 0 and len(origin_domain) > 0:
|
|
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)
|
|
|
|
if request.method == "OPTIONS":
|
|
response = web.Response()
|
|
else:
|
|
response = await handler(request)
|
|
|
|
return response
|
|
|
|
return origin_only_middleware
|