Merge pull request #12698 from Akegarasu/fix-ssrf-in-api

fix potential ssrf attack in #12663
This commit is contained in:
AUTOMATIC1111 2023-08-21 07:19:48 +03:00 committed by GitHub
commit a7f18b2297
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -4,6 +4,8 @@ import os
import time import time
import datetime import datetime
import uvicorn import uvicorn
import ipaddress
import requests
import gradio as gr import gradio as gr
from threading import Lock from threading import Lock
from io import BytesIO from io import BytesIO
@ -56,8 +58,27 @@ def setUpscalers(req: dict):
def decode_base64_to_image(encoding): def decode_base64_to_image(encoding):
def verify_url(url):
import socket
from urllib.parse import urlparse
try:
parsed_url = urlparse(url)
domain_name = parsed_url.netloc
host = socket.gethostbyname_ex(domain_name)
for ip in host[2]:
ip_addr = ipaddress.ip_address(ip)
# https://docs.python.org/3/library/ipaddress.html#ipaddress.IPv4Address.is_global
if not ip_addr.is_global:
return False
except Exception:
return False
return True
if encoding.startswith("http://") or encoding.startswith("https://"): if encoding.startswith("http://") or encoding.startswith("https://"):
import requests if not verify_url(encoding):
raise HTTPException(status_code=500, detail="Invalid image url")
response = requests.get(encoding, timeout=30, headers={'user-agent':'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/62.0.3202.94 Safari/537.36'}) response = requests.get(encoding, timeout=30, headers={'user-agent':'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/62.0.3202.94 Safari/537.36'})
try: try:
image = Image.open(BytesIO(response.content)) image = Image.open(BytesIO(response.content))