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)