1 files changed, 33 insertions, 35 deletions

diff --git a/lib/unicorn/http_response.rb b/lib/unicorn/http_response.rb
index 7bbb940..15df3f6 100644
--- a/lib/unicorn/http_response.rb
+++ b/lib/unicorn/http_response.rb
@@ -21,52 +21,50 @@ module Unicorn
 
   class HttpResponse
 
-    # headers we allow duplicates for
-    ALLOWED_DUPLICATES = {
-      'Set-Cookie' => true,
-      'Set-Cookie2' => true,
-      'Warning' => true,
-      'WWW-Authenticate' => true,
-    }.freeze
+    # Every standard HTTP code mapped to the appropriate message.
+    CODES = Rack::Utils::HTTP_STATUS_CODES.inject({}) { |hash,(code,msg)|
+      hash[code] = "#{code} #{msg}"
+      hash
+    }
+
+    # Rack does not set/require a Date: header.  We always override the
+    # Connection: and Date: headers no matter what (if anything) our
+    # Rack application sent us.
+    SKIP = { 'connection' => true, 'date' => true, 'status' => true }.freeze
+    EMPTY = ''.freeze # :nodoc
+    OUT = [] # :nodoc
 
     # writes the rack_response to socket as an HTTP response
     def self.write(socket, rack_response)
       status, headers, body = rack_response
+      status = CODES[status.to_i]
+      OUT.clear
 
-      # Rack does not set/require Date, but don't worry about Content-Length
-      # since Rack applications that conform to Rack::Lint enforce that
-      out = [ "#{Const::DATE}: #{Time.now.httpdate}" ]
-      sent = { Const::CONNECTION => true, Const::DATE => true }
-
+      # Don't bother enforcing duplicate supression, it's a Hash most of
+      # the time anyways so just hope our app knows what it's doing
       headers.each do |key, value|
-        if ! sent[key] || ALLOWED_DUPLICATES[key]
-          sent[key] = true
-          out << "#{key}: #{value}"
+        next if SKIP.include?(key.downcase)
+        if value =~ /\n/
+          value.split(/\n/).each { |v| OUT << "#{key}: #{v}\r\n" }
+        else
+          OUT << "#{key}: #{value}\r\n"
         end
       end
 
-      socket_write(socket,
-                   "HTTP/1.1 #{status} #{HTTP_STATUS_CODES[status]}\r\n" \
+      # Rack should enforce Content-Length or chunked transfer encoding,
+      # so don't worry or care about them.
+      # Date is required by HTTP/1.1 as long as our clock can be trusted.
+      # Some broken clients require a "Status" header so we accomodate them
+      socket.write("HTTP/1.1 #{status}\r\n" \
+                   "Date: #{Time.now.httpdate}\r\n" \
+                   "Status: #{status}\r\n" \
                    "Connection: close\r\n" \
-                   "#{out.join("\r\n")}\r\n\r\n")
-      body.each { |chunk| socket_write(socket, chunk) }
+                   "#{OUT.join(EMPTY)}\r\n")
+      body.each { |chunk| socket.write(chunk) }
+      socket.close # flushes and uncorks the socket immediately
+      ensure
+        body.respond_to?(:close) and body.close rescue nil
     end
 
-    private
-
-      # write(2) can return short on slow devices like sockets as well
-      # as fail with EINTR if a signal was caught.
-      def self.socket_write(socket, buffer)
-        loop do
-          begin
-            written = socket.syswrite(buffer)
-            return written if written == buffer.length
-            buffer = buffer[written..-1]
-          rescue Errno::EINTR
-            retry
-          end
-        end
-      end
-
   end
 end