From d5079d069b03c825eda20043cdd66091b6c0e22c Mon Sep 17 00:00:00 2001 From: Mike Guo Date: Wed, 14 Apr 2021 15:31:47 +0800 Subject: [PATCH] fix security warning like X-Content-Type-Options is required to be nosniff --- tb_plugin/torch_tb_profiler/plugin.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tb_plugin/torch_tb_profiler/plugin.py b/tb_plugin/torch_tb_profiler/plugin.py index 95a479448..3a5231e93 100644 --- a/tb_plugin/torch_tb_profiler/plugin.py +++ b/tb_plugin/torch_tb_profiler/plugin.py @@ -29,6 +29,7 @@ class TorchProfilerPlugin(base_plugin.TBPlugin): """TensorBoard plugin for Torch Profiler.""" plugin_name = consts.PLUGIN_NAME + headers = [('X-Content-Type-Options', 'nosniff')] def __init__(self, context): """Instantiates TorchProfilerPlugin. @@ -237,9 +238,10 @@ def trace_route(self, request): if profile.trace_file_path.endswith('.gz'): headers = [] headers.append(('Content-Encoding', 'gzip')) + headers.extend(TorchProfilerPlugin.headers) return werkzeug.Response(raw_data, content_type="application/json", headers=headers) else: - return werkzeug.Response(raw_data, content_type="application/json") + return werkzeug.Response(raw_data, content_type="application/json", headers=TorchProfilerPlugin.headers) @wrappers.Request.application def static_file_route(self, request): @@ -258,15 +260,15 @@ def static_file_route(self, request): with open(filepath, 'rb') as infile: contents = infile.read() except IOError: - return werkzeug.Response('404 Not Found', 'text/plain', code=404) + return werkzeug.Response('404 Not Found', 'text/plain', code=404, headers=TorchProfilerPlugin.headers) return werkzeug.Response( - contents, content_type=mimetype + contents, content_type=mimetype, headers=TorchProfilerPlugin.headers ) @staticmethod def respond_as_json(obj): content = json.dumps(obj) - return werkzeug.Response(content, content_type="application/json") + return werkzeug.Response(content, content_type="application/json", headers=TorchProfilerPlugin.headers) def _load_run(queue, name, run_dir):