From cc85946f551bc328eefa21012f48f9f3d6d08e9c Mon Sep 17 00:00:00 2001
From: David Baker <dbkr@users.noreply.github.com>
Date: Wed, 19 Dec 2018 06:17:02 +0000
Subject: [PATCH] feat: Support response headers in File protocol handler
 (#16098)

* feat: Support response headers in File protocol handler

* bugfix: Null check headers value & fix tests

* refactor: Use non-deprecated FindKeyOfType
---
 .../browser/net/url_request_async_asar_job.cc | 27 ++++++++++++-------
 atom/browser/net/url_request_async_asar_job.h |  1 +
 docs/api/protocol.md                          |  4 ++-
 spec/api-protocol-spec.js                     | 22 +++++++++++++++
 4 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/atom/browser/net/url_request_async_asar_job.cc b/atom/browser/net/url_request_async_asar_job.cc
index 1139b1c35402..bd6601219f1c 100644
--- a/atom/browser/net/url_request_async_asar_job.cc
+++ b/atom/browser/net/url_request_async_asar_job.cc
@@ -71,14 +71,27 @@ void URLRequestAsyncAsarJob::StartAsync(std::unique_ptr<base::Value> options,
   }
 
   std::string file_path;
+  response_headers_ = new net::HttpResponseHeaders("HTTP/1.1 200 OK");
   if (options->is_dict()) {
-    auto* path_value =
-        options->FindKeyOfType("path", base::Value::Type::STRING);
-    if (path_value)
-      file_path = path_value->GetString();
+    base::DictionaryValue* dict =
+        static_cast<base::DictionaryValue*>(options.get());
+    base::Value* pathValue =
+        dict->FindKeyOfType("path", base::Value::Type::STRING);
+    if (pathValue) {
+      file_path = pathValue->GetString();
+    }
+    base::Value* headersValue =
+        dict->FindKeyOfType("headers", base::Value::Type::DICTIONARY);
+    if (headersValue) {
+      for (const auto& iter : headersValue->DictItems()) {
+        response_headers_->AddHeader(iter.first + ": " +
+                                     iter.second.GetString());
+      }
+    }
   } else if (options->is_string()) {
     file_path = options->GetString();
   }
+  response_headers_->AddHeader(kCORSHeader);
 
   if (file_path.empty()) {
     NotifyStartError(net::URLRequestStatus(net::URLRequestStatus::FAILED,
@@ -103,11 +116,7 @@ void URLRequestAsyncAsarJob::Kill() {
 }
 
 void URLRequestAsyncAsarJob::GetResponseInfo(net::HttpResponseInfo* info) {
-  std::string status("HTTP/1.1 200 OK");
-  auto* headers = new net::HttpResponseHeaders(status);
-
-  headers->AddHeader(kCORSHeader);
-  info->headers = headers;
+  info->headers = response_headers_;
 }
 
 }  // namespace atom
diff --git a/atom/browser/net/url_request_async_asar_job.h b/atom/browser/net/url_request_async_asar_job.h
index f97dcc1570ea..9c9cbb69cbd3 100644
--- a/atom/browser/net/url_request_async_asar_job.h
+++ b/atom/browser/net/url_request_async_asar_job.h
@@ -26,6 +26,7 @@ class URLRequestAsyncAsarJob : public asar::URLRequestAsarJob, public JsAsker {
   void Kill() override;
 
  private:
+  scoped_refptr<net::HttpResponseHeaders> response_headers_;
   base::WeakPtrFactory<URLRequestAsyncAsarJob> weak_factory_;
 
   DISALLOW_COPY_AND_ASSIGN(URLRequestAsyncAsarJob);
diff --git a/docs/api/protocol.md b/docs/api/protocol.md
index 5220066143b7..5829f3d12ce3 100644
--- a/docs/api/protocol.md
+++ b/docs/api/protocol.md
@@ -99,7 +99,9 @@ going to be created with `scheme`. `completion` will be called with
 
 To handle the `request`, the `callback` should be called with either the file's
 path or an object that has a `path` property, e.g. `callback(filePath)` or
-`callback({ path: filePath })`.
+`callback({ path: filePath })`. The object may also have a `headers` property
+which gives a list of strings for the response headers, e.g.
+`callback({ path: filePath, headers: ["Content-Security-Policy: default-src 'none'"]})`.
 
 When `callback` is called with nothing, a number, or an object that has an
 `error` property, the `request` will fail with the `error` number you
diff --git a/spec/api-protocol-spec.js b/spec/api-protocol-spec.js
index 55ec69e92581..d2d8760c33c7 100644
--- a/spec/api-protocol-spec.js
+++ b/spec/api-protocol-spec.js
@@ -320,6 +320,28 @@ describe('protocol module', () => {
       })
     })
 
+    it('sets custom headers', (done) => {
+      const handler = (request, callback) => callback({
+        path: filePath,
+        headers: { 'X-Great-Header': 'sogreat' }
+      })
+      protocol.registerFileProtocol(protocolName, handler, (error) => {
+        if (error) return done(error)
+        $.ajax({
+          url: protocolName + '://fake-host',
+          cache: false,
+          success: (data, status, request) => {
+            assert.strictEqual(data, String(fileContent))
+            assert.strictEqual(request.getResponseHeader('X-Great-Header'), 'sogreat')
+            done()
+          },
+          error: (xhr, errorType, error) => {
+            done(error)
+          }
+        })
+      })
+    })
+
     it('sends object as response', (done) => {
       const handler = (request, callback) => callback({ path: filePath })
       protocol.registerFileProtocol(protocolName, handler, (error) => {