From 2458a601293c86bdfba5be24dc95b353447e6460 Mon Sep 17 00:00:00 2001
From: Kevin McDermott <kevin@heroku.com>
Date: Thu, 30 Nov 2017 16:48:17 +0000
Subject: [PATCH] Add mechanism for additional headers.

---
 lib/pliny/middleware/cors.rb | 14 ++++++++++++--
 spec/middleware/cors_spec.rb | 11 +++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/lib/pliny/middleware/cors.rb b/lib/pliny/middleware/cors.rb
index 445e857b..813bb691 100644
--- a/lib/pliny/middleware/cors.rb
+++ b/lib/pliny/middleware/cors.rb
@@ -8,6 +8,12 @@ class CORS
     EXPOSE_HEADERS =
       %w( Cache-Control Content-Language Content-Type Expires Last-Modified Pragma ).freeze
 
+    @@additional_headers = []
+
+    def self.add_additional_header(header)
+      @@additional_headers << header
+    end
+
     def initialize(app)
       @app = app
     end
@@ -19,7 +25,7 @@ def call(env)
       else
         status, headers, response = @app.call(env)
 
-        # regualar CORS request: append CORS headers to response
+        # regular CORS request: append CORS headers to response
         if cors_request?(env)
           headers.merge!(cors_headers(env))
         end
@@ -32,11 +38,15 @@ def cors_request?(env)
       env.has_key?("HTTP_ORIGIN")
     end
 
+    def allow_headers
+      ALLOW_HEADERS + @@additional_headers
+    end
+
     def cors_headers(env)
       {
         'Access-Control-Allow-Origin'      => env["HTTP_ORIGIN"],
         'Access-Control-Allow-Methods'     => ALLOW_METHODS.join(', '),
-        'Access-Control-Allow-Headers'     => ALLOW_HEADERS.join(', '),
+        'Access-Control-Allow-Headers'     => allow_headers.join(', '),
         'Access-Control-Allow-Credentials' => "true",
         'Access-Control-Max-Age'           => "1728000",
         'Access-Control-Expose-Headers'    => EXPOSE_HEADERS.join(', ')
diff --git a/spec/middleware/cors_spec.rb b/spec/middleware/cors_spec.rb
index 59462f40..3dcd246b 100644
--- a/spec/middleware/cors_spec.rb
+++ b/spec/middleware/cors_spec.rb
@@ -39,4 +39,15 @@ def app
     assert_equal "http://localhost",
       last_response.headers["Access-Control-Allow-Origin"]
   end
+
+  it "allows additional headers to be added to every response" do
+    Pliny::Middleware::CORS.add_additional_header("X-Origin")
+
+    header "Origin", "http://localhost"
+    get "/"
+    assert_equal 200, last_response.status
+    assert_equal "hi", last_response.body
+
+    assert last_response.headers["Access-Control-Allow-Headers"].include?("X-Origin")
+  end
 end