From 65cf3fd59b489ff94f67f0610b4ca4b73e731f38 Mon Sep 17 00:00:00 2001
From: Konstantin Haase <konstantin.mailinglists@googlemail.com>
Date: Fri, 1 Mar 2013 15:36:05 +1100
Subject: [PATCH] rework protection headers, fixes #40

---
 .../lib/rack/protection/frame_options.rb           | 14 ++++++++++----
 rack-protection/lib/rack/protection/xss_header.rb  | 12 ++----------
 rack-protection/spec/xss_header_spec.rb            |  6 ++++++
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/rack-protection/lib/rack/protection/frame_options.rb b/rack-protection/lib/rack/protection/frame_options.rb
index 2939748b..bce75c47 100644
--- a/rack-protection/lib/rack/protection/frame_options.rb
+++ b/rack-protection/lib/rack/protection/frame_options.rb
@@ -16,16 +16,22 @@ module Rack
     # frame_options:: Defines who should be allowed to embed the page in a
     #                 frame. Use :deny to forbid any embedding, :sameorigin
     #                 to allow embedding from the same origin (default).
-    class FrameOptions < XSSHeader
+    class FrameOptions < Base
       default_options :frame_options => :sameorigin
 
-      def header
-        @header ||= begin
+      def frame_options
+        @frame_options ||= begin
           frame_options = options[:frame_options]
           frame_options = options[:frame_options].to_s.upcase unless frame_options.respond_to? :to_str
-          { 'X-Frame-Options' => frame_options.to_str }
+          frame_options.to_str
         end
       end
+
+      def call(env)
+        status, headers, body        = @app.call(env)
+        headers['X-Frame-Options'] ||= frame_options if html? headers
+        [status, headers, body]
+      end
     end
   end
 end
diff --git a/rack-protection/lib/rack/protection/xss_header.rb b/rack-protection/lib/rack/protection/xss_header.rb
index 54fd9840..6bb14861 100644
--- a/rack-protection/lib/rack/protection/xss_header.rb
+++ b/rack-protection/lib/rack/protection/xss_header.rb
@@ -14,18 +14,10 @@ module Rack
     class XSSHeader < Base
       default_options :xss_mode => :block, :nosniff => true
 
-      def header
-        headers = {
-          'X-XSS-Protection' => "1; mode=#{options[:xss_mode]}",
-          'X-Content-Type-Options' => "nosniff"
-        }
-        headers.delete("X-Content-Type-Options") unless options[:nosniff]
-        headers
-      end
-
       def call(env)
         status, headers, body = @app.call(env)
-        headers = header.merge(headers) if options[:nosniff] and html?(headers)
+        headers['X-XSS-Protection']       ||= "1; mode=#{options[:xss_mode]}" if html? headers
+        headers['X-Content-Type-Options'] ||= 'nosniff'                       if options[:nosniff]
         [status, headers, body]
       end
     end
diff --git a/rack-protection/spec/xss_header_spec.rb b/rack-protection/spec/xss_header_spec.rb
index cf62e555..05c94697 100644
--- a/rack-protection/spec/xss_header_spec.rb
+++ b/rack-protection/spec/xss_header_spec.rb
@@ -34,6 +34,12 @@ describe Rack::Protection::XSSHeader do
     get('/', {}, 'wants' => 'text/html').header["X-Content-Type-Options"].should == "nosniff"
   end
 
+
+  it 'should set the X-Content-Type-Options for other content types' do
+    get('/', {}, 'wants' => 'application/foo').header["X-Content-Type-Options"].should == "nosniff"
+  end
+
+
   it 'should allow changing the nosniff-mode off' do
     mock_app do
       use Rack::Protection::XSSHeader, :nosniff => false