lib/secure_headers/middleware.rb in secure_headers-3.0.3 vs lib/secure_headers/middleware.rb in secure_headers-3.1.0

@@ -1,15 +1,38 @@
 module SecureHeaders
   class Middleware
+    SECURE_COOKIE_REGEXP = /;\s*secure\s*(;|$)/i.freeze
+
     def initialize(app)
       @app = app
     end
 
     # merges the hash of headers into the current header set.
     def call(env)
       req = Rack::Request.new(env)
       status, headers, response = @app.call(env)
+
+      config = SecureHeaders.config_for(req)
+      flag_cookies_as_secure!(headers) if config.secure_cookies
       headers.merge!(SecureHeaders.header_hash_for(req))
       [status, headers, response]
+    end
+
+    private
+
+    # inspired by https://github.com/tobmatth/rack-ssl-enforcer/blob/6c014/lib/rack/ssl-enforcer.rb#L183-L194
+    def flag_cookies_as_secure!(headers)
+      if cookies = headers['Set-Cookie']
+        # Support Rails 2.3 / Rack 1.1 arrays as headers
+        cookies = cookies.split("\n") unless cookies.is_a?(Array)
+
+        headers['Set-Cookie'] = cookies.map do |cookie|
+          if cookie !~ SECURE_COOKIE_REGEXP
+            "#{cookie}; secure"
+          else
+            cookie
+          end
+        end.join("\n")
+      end
     end
   end
 end