From 59d58b1158141ed5de65adabf4ef42bcbb3ee3da Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Fri, 13 Dec 2024 01:56:43 -0800 Subject: [PATCH] [Security] Fix potential XSS on /view (#6034) --- server.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/server.py b/server.py index a8f0d767..ddd71e06 100644 --- a/server.py +++ b/server.py @@ -460,7 +460,21 @@ class PromptServer(): return web.Response(body=alpha_buffer.read(), content_type='image/png', headers={"Content-Disposition": f"filename=\"{filename}\""}) else: - return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""}) + # Get content type from mimetype, defaulting to 'application/octet-stream' + content_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream' + + # For security, force certain extensions to download instead of display + file_extension = os.path.splitext(filename)[1].lower() + if file_extension in {'.html', '.htm', '.js', '.css'}: + content_type = 'application/octet-stream' # Forces download + + return web.FileResponse( + file, + headers={ + "Content-Disposition": f"filename=\"{filename}\"", + "Content-Type": content_type + } + ) return web.Response(status=404)