protocol.rb in ftw-0.0.14

- old
+ new

@@ -1,11 +1,14 @@
 require "ftw/namespace"
+require "ftw/crlf"
 require "cabin"
 require "logger"
 
 # This module provides web protocol handling as a mixin.
 module FTW::Protocol
+  include FTW::CRLF
+
   # Read an HTTP message from a given connection
   #
   # This method blocks until a full http message header has been consumed
   # (request *or* response)
   #
@@ -54,7 +57,106 @@
       parser.headers.each { |field, value| response.headers.add(field, value) }
       return response
     end
   end # def read_http_message
 
-  public(:read_http_message)
+  def write_http_body(body, io, chunked=false)
+    if chunked
+      write_http_body_chunked(body, io)
+    else
+      write_http_body_normal(body, io)
+    end
+  end # def write_http_body
+
+  # Encode the given text as in 'chunked' encoding.
+  def encode_chunked(text)
+    return sprintf("%x%s%s%s", text.size, CRLF, text, CRLF)
+  end # def encode_chunked
+
+  def write_http_body_chunked(body, io)
+    if body.is_a?(String)
+      io.write(encode_chunked(body))
+    elsif body.respond_to?(:sysread)
+      true while io.write(encode_chunked(body.sysread(16384)))
+    elsif body.respond_to?(:read)
+      true while io.write(encode_chunked(body.read(16384)))
+    elsif body.respond_to?(:each)
+      body.each { |s| io.write(encode_chunked(s)) }
+    end
+
+    # The terminating chunk is an empty one.
+    io.write(encode_chunked(""))
+  end # def write_http_body_chunked
+
+  def write_http_body_normal(body, io)
+    if body.is_a?(String)
+      io.write(body)
+    elsif body.respond_to?(:read)
+      true while io.write(body.read(16384))
+    elsif body.respond_to?(:each)
+      body.each { |s| io.write(s) }
+    end
+  end # def write_http_body_normal
+
+  # Read the body of this message. The block is called with chunks of the
+  # response as they are read in.
+  #
+  # This method is generally only called by http clients, not servers.
+  def read_http_body(&block)
+    if @body.respond_to?(:read)
+      if headers.include?("Content-Length") and headers["Content-Length"].to_i > 0
+        @logger.debug("Reading body with Content-Length")
+        read_http_body_length(headers["Content-Length"].to_i, &block)
+      elsif headers["Transfer-Encoding"] == "chunked"
+        @logger.debug("Reading body with chunked encoding")
+        read_http_body_chunked(&block)
+      end
+
+      # If this is a poolable resource, release it (like a FTW::Connection)
+      @body.release if @body.respond_to?(:release)
+    elsif !@body.nil?
+      yield @body
+    end
+  end # def read_http_body
+
+  # Old api compat
+  alias_method :read_body, :read_http_body
+
+  # Read the length bytes from the body. Yield each chunk read to the block
+  # given. This method is generally only called by http clients, not servers.
+  def read_http_body_length(length, &block)
+    remaining = length
+    while remaining > 0
+      data = @body.read
+      @logger.debug("Read bytes", :length => data.size)
+      if data.size > remaining
+        # Read too much data, only wanted part of this. Push the rest back.
+        yield data[0..remaining]
+        remaining = 0
+        @body.pushback(data[remaining .. -1]) if remaining < 0
+      else
+        yield data
+        remaining -= data.size
+      end
+    end
+  end # def read_http_body_length
+
+  # This is kind of messed, need to fix it.
+  def read_http_body_chunked(&block)
+    parser = HTTP::Parser.new
+
+    # Fake fill-in the response we've already read into the parser.
+    parser << to_s
+    parser << CRLF
+    parser.on_body = block
+    done = false
+    parser.on_message_complete = proc { done = true }
+
+    while !done # will break on special conditions below
+      data = @body.read
+      offset = parser << data
+      if offset != data.length
+        raise "Parser did not consume all data read?"
+      end
+    end
+  end # def read_http_body_chunked
 end # module FTW::Protocol