Don't buffer response bodies in session store by creating an unnecessary response wrapper. Extracted set and delete cookie helpers into Utils so they can be used outside Response.

4 files changed, 56 insertions, 45 deletions

diff --git a/lib/rack/response.rb b/lib/rack/response.rb
index 28b4d830..d1f6a123 100644
--- a/lib/rack/response.rb
+++ b/lib/rack/response.rb
@@ -54,45 +54,11 @@ module Rack
     end
 
     def set_cookie(key, value)
-      case value
-      when Hash
-        domain  = "; domain="  + value[:domain]    if value[:domain]
-        path    = "; path="    + value[:path]      if value[:path]
-        # According to RFC 2109, we need dashes here.
-        # N.B.: cgi.rb uses spaces...
-        expires = "; expires=" + value[:expires].clone.gmtime.
-          strftime("%a, %d-%b-%Y %H:%M:%S GMT")    if value[:expires]
-        secure = "; secure"  if value[:secure]
-        httponly = "; HttpOnly" if value[:httponly]
-        value = value[:value]
-      end
-      value = [value]  unless Array === value
-      cookie = Utils.escape(key) + "=" +
-        value.map { |v| Utils.escape v }.join("&") +
-        "#{domain}#{path}#{expires}#{secure}#{httponly}"
-
-      case self["Set-Cookie"]
-      when Array
-        self["Set-Cookie"] << cookie
-      when String
-        self["Set-Cookie"] = [self["Set-Cookie"], cookie]
-      when nil
-        self["Set-Cookie"] = cookie
-      end
+      Utils.set_cookie_header!(header, key, value)
     end
 
     def delete_cookie(key, value={})
-      unless Array === self["Set-Cookie"]
-        self["Set-Cookie"] = [self["Set-Cookie"]].compact
-      end
-
-      self["Set-Cookie"].reject! { |cookie|
-        cookie =~ /\A#{Utils.escape(key)}=/
-      }
-
-      set_cookie(key,
-                 {:value => '', :path => nil, :domain => nil,
-                   :expires => Time.at(0) }.merge(value))
+      Utils.delete_cookie_header!(header, key, value)
     end
 
     def redirect(target, status=302)
diff --git a/lib/rack/session/abstract/id.rb b/lib/rack/session/abstract/id.rb
index 218144c1..98746705 100644
--- a/lib/rack/session/abstract/id.rb
+++ b/lib/rack/session/abstract/id.rb
@@ -107,18 +107,16 @@ module Rack
 
           if not session_id = set_session(env, session_id, session, options)
             env["rack.errors"].puts("Warning! #{self.class.name} failed to save session. Content dropped.")
-            [status, headers, body]
           elsif options[:defer] and not options[:renew]
             env["rack.errors"].puts("Defering cookie for #{session_id}") if $VERBOSE
-            [status, headers, body]
           else
             cookie = Hash.new
             cookie[:value] = session_id
             cookie[:expires] = Time.now + options[:expire_after] unless options[:expire_after].nil?
-            response = Rack::Response.new(body, status, headers)
-            response.set_cookie(@key, cookie.merge(options))
-            response.to_a
+            Utils.set_cookie_header!(headers, @key, cookie.merge(options))
           end
+
+          [status, headers, body]
         end
 
         # All thread safety and session retrival proceedures should occur here.
diff --git a/lib/rack/session/cookie.rb b/lib/rack/session/cookie.rb
index eace9bd0..240e6c8d 100644
--- a/lib/rack/session/cookie.rb
+++ b/lib/rack/session/cookie.rb
@@ -70,16 +70,15 @@ module Rack
 
         if session_data.size > (4096 - @key.size)
           env["rack.errors"].puts("Warning! Rack::Session::Cookie data size exceeds 4K. Content dropped.")
-          [status, headers, body]
         else
           options = env["rack.session.options"]
           cookie = Hash.new
           cookie[:value] = session_data
           cookie[:expires] = Time.now + options[:expire_after] unless options[:expire_after].nil?
-          response = Rack::Response.new(body, status, headers)
-          response.set_cookie(@key, cookie.merge(options))
-          response.to_a
+          Utils.set_cookie_header!(headers, @key, cookie.merge(options))
         end
+
+        [status, headers, body]
       end
 
       def generate_hmac(data)
diff --git a/lib/rack/utils.rb b/lib/rack/utils.rb
index a70dcf64..55f01f3b 100644
--- a/lib/rack/utils.rb
+++ b/lib/rack/utils.rb
@@ -168,6 +168,54 @@ module Rack
     end
     module_function :select_best_encoding
 
+    def set_cookie_header!(header, key, value)
+      case value
+      when Hash
+        domain  = "; domain="  + value[:domain] if value[:domain]
+        path    = "; path="    + value[:path]   if value[:path]
+        # According to RFC 2109, we need dashes here.
+        # N.B.: cgi.rb uses spaces...
+        expires = "; expires=" + value[:expires].clone.gmtime.
+          strftime("%a, %d-%b-%Y %H:%M:%S GMT") if value[:expires]
+        secure = "; secure"  if value[:secure]
+        httponly = "; HttpOnly" if value[:httponly]
+        value = value[:value]
+      end
+      value = [value] unless Array === value
+      cookie = escape(key) + "=" +
+        value.map { |v| escape v }.join("&") +
+        "#{domain}#{path}#{expires}#{secure}#{httponly}"
+
+      case header["Set-Cookie"]
+      when Array
+        header["Set-Cookie"] << cookie
+      when String
+        header["Set-Cookie"] = [header["Set-Cookie"], cookie]
+      when nil
+        header["Set-Cookie"] = cookie
+      end
+
+      nil
+    end
+    module_function :set_cookie_header!
+
+    def delete_cookie_header!(header, key, value = {})
+      unless Array === header["Set-Cookie"]
+        header["Set-Cookie"] = [header["Set-Cookie"]].compact
+      end
+
+      header["Set-Cookie"].reject! { |cookie|
+        cookie =~ /\A#{escape(key)}=/
+      }
+
+      set_cookie_header!(header, key,
+                 {:value => '', :path => nil, :domain => nil,
+                   :expires => Time.at(0) }.merge(value))
+
+      nil
+    end
+    module_function :delete_cookie_header!
+
     # Return the bytesize of String; uses String#length under Ruby 1.8 and
     # String#bytesize under 1.9.
     if ''.respond_to?(:bytesize)