Merge pull request whitequark#26 from sj26/adjust-rewritten-content-l…

…ength

Make sure Content-Length is adjusted

lib/rack/utf8_sanitizer.rb

@@ -62,7 +62,13 @@ def sanitize_rack_input(env)
       content_type &&= content_type.downcase
       return unless SANITIZABLE_CONTENT_TYPES.any? {|type| content_type == type }
       uri_encoded = URI_ENCODED_CONTENT_TYPES.any? {|type| content_type == type}
-      env['rack.input'] &&= sanitize_io(env['rack.input'], uri_encoded)
+
+      if env["rack.input"]
+        sanitized_input = sanitize_io(env['rack.input'], uri_encoded)
+
+        env['rack.input'] = sanitized_input
+        env['CONTENT_LENGTH'] &&= sanitized_input.size.to_s
+      end
     end
 
     # Modeled after Rack::RewindableInput
@@ -72,18 +78,28 @@ def initialize(original_io, sanitized_io)
         @original_io = original_io
         @sanitized_io = sanitized_io
       end
+
       def gets
         @sanitized_io.gets
       end
+
       def read(*args)
         @sanitized_io.read(*args)
       end
+
       def each(&block)
         @sanitized_io.each(&block)
       end
+
       def rewind
         @sanitized_io.rewind
       end
+
+      def size
+        # StringIO#size is bytesize
+        @sanitized_io.size
+      end
+
       def close
         @sanitized_io.close
         @original_io.close if @original_io.respond_to?(:close)

test/test_utf8_sanitizer.rb

@@ -178,17 +178,20 @@ def request_env
          "rack.input" => @rack_input,
       }
     end
+
     def sanitize_form_data(request_env = request_env)
-      @uri_input   = "http://bar/foo+%2F%3A+bar+%D0%BB%D0%BE%D0%BB".force_encoding('UTF-8')
-      env = @app.(request_env)
-      sanitized_input = env['rack.input'].read
+      @uri_input = "http://bar/foo+%2F%3A+bar+%D0%BB%D0%BE%D0%BB".force_encoding('UTF-8')
+      @response_env = @app.(request_env)
+      sanitized_input = @response_env['rack.input'].read
+
       yield sanitized_input if block_given?
-      env['rack.input'].rewind
+
+      @response_env['rack.input'].rewind
       behaves_like :does_sanitize_plain
       behaves_like :does_sanitize_uri
       behaves_like :identity_plain
       behaves_like :identity_uri
-      env['rack.input'].close
+      @response_env['rack.input'].close
     end
 
     it "sanitizes StringIO rack.input" do
@@ -288,5 +291,15 @@ def sanitize_form_data(request_env = request_env)
       end
     end
 
+    it "adjusts content-length when replacing input" do
+      input =  "foo=bla&quux=bar\xED"
+      @rack_input = StringIO.new input
+
+      env = request_env.update("CONTENT_LENGTH" => input.bytesize)
+      sanitize_form_data(env) do |sanitized_input|
+        sanitized_input.bytesize.should != input.bytesize
+        @response_env["CONTENT_LENGTH"].should == sanitized_input.bytesize.to_s
+      end
+    end
   end
 end