proxy_pass: preliminary support for passing trailers

Rack apps may (through a round-about way) send HTTP trailers
to HTTP/1.1 clients, and we need a way to forward those responses
through without losing the trailers.

2 files changed, 62 insertions, 20 deletions

diff --git a/lib/yahns/proxy_http_response.rb b/lib/yahns/proxy_http_response.rb
index 8f4790e..fe0a37b 100644
--- a/lib/yahns/proxy_http_response.rb
+++ b/lib/yahns/proxy_http_response.rb
@@ -120,6 +120,7 @@ module Yahns::HttpResponse # :nodoc:
         end until len == 0
 
       else # nasty chunked body
+        req_res.proxy_trailers = nil # define to avoid warnings for now
 
         # Only HTTP/1.1 supports chunked responses, we must translate
         # otherwise.  Otherwise, we must drop the connection to signal
@@ -130,9 +131,7 @@ module Yahns::HttpResponse # :nodoc:
         when String
           kcar.filter_body(buf, tmp)
           unless buf.empty?
-            tmp = rechunk ? [ "#{buf.size.to_s(16)}\r\n", buf, "\r\n".freeze ]
-                          : buf
-            wbuf = proxy_write(wbuf, tmp, alive)
+            wbuf = proxy_write(wbuf, rechunk ? chunk_out(buf) : buf, alive)
           end
         when nil # premature EOF
           return proxy_err_response(nil, req_res, nil, wbuf)
@@ -142,8 +141,25 @@ module Yahns::HttpResponse # :nodoc:
           return :wait_readable # self remains in :ignore, wait on upstream
         end until kcar.body_eof?
 
-        # TODO: Trailer support
-        wbuf = proxy_write(wbuf, "0\r\n\r\n".freeze, alive) if rechunk
+        buf = tmp
+        req_res.proxy_trailers = [ buf, tlr = [] ]
+        rbuf = Thread.current[:yahns_rbuf] = ''
+        if rechunk
+          until kcar.trailers(tlr, buf)
+            case rv = req_res.kgio_tryread(0x2000, rbuf)
+            when String
+              buf << rv
+            when :wait_readable
+              # for ensure:
+              wbuf ||= Yahns::Wbuf.new(nil, alive, k.output_buffer_tmpdir,
+                                       false)
+              return :wait_readable
+            when nil # premature EOF
+              return proxy_err_response(nil, req_res, nil, wbuf)
+            end # no loop here
+          end
+          wbuf = proxy_write(wbuf, trailer_out(tlr), alive)
+        end
       end
     end
 
@@ -178,22 +194,36 @@ module Yahns::HttpResponse # :nodoc:
       # Rack/Rails can always send streaming responses.
       rechunk = @hs.env['HTTP_VERSION'] == 'HTTP/1.1'.freeze
       buf = ''
-      case tmp = req_res.kgio_tryread(0x2000, rbuf)
-      when String
-        kcar.filter_body(buf, tmp)
-        unless buf.empty?
-          tmp = rechunk ? [ "#{buf.size.to_s(16)}\r\n", buf, "\r\n".freeze ]
-                        : buf
-          wbuf.wbuf_write(self, tmp)
-        end
-      when nil # premature EOF
-        return proxy_err_response(nil, req_res, nil, wbuf)
-      when :wait_readable
-        return :wait_readable # self remains in :ignore, wait on upstream
-      end until kcar.body_eof?
 
-      # TODO: Trailer support
-      wbuf.wbuf_write(self, "0\r\n\r\n".freeze) if rechunk
+      unless req_res.proxy_trailers
+        # are we done dechunking the main body, yet?
+        case tmp = req_res.kgio_tryread(0x2000, rbuf)
+        when String
+          kcar.filter_body(buf, tmp)
+          buf.empty? or wbuf.wbuf_write(self, rechunk ? chunk_out(buf) : buf)
+        when nil # premature EOF
+          return proxy_err_response(nil, req_res, nil, wbuf)
+        when :wait_readable
+          return :wait_readable # self remains in :ignore, wait on upstream
+        end until kcar.body_eof?
+        req_res.proxy_trailers = [ tmp, [] ] # onto trailers!
+        rbuf = Thread.current[:yahns_rbuf] = ''
+      end
+
+      buf, tlr = *req_res.proxy_trailers
+      if rechunk
+        until kcar.trailers(tlr, buf)
+          case rv = req_res.kgio_tryread(0x2000, rbuf)
+          when String
+            buf << rv
+          when :wait_readable
+            return :wait_readable
+          when nil # premature EOF
+            return proxy_err_response(nil, req_res, nil, wbuf)
+          end # no loop here
+        end
+        wbuf.wbuf_write(self, trailer_out(tlr))
+      end
     end
 
     busy = wbuf.busy and return proxy_busy_mod_blocked(wbuf, busy)
@@ -226,4 +256,15 @@ module Yahns::HttpResponse # :nodoc:
     # no touching self after queue_mod
     :ignore
   end
+
+  # n.b.: we can use String#size for optimized dispatch under YARV instead
+  # of String#bytesize because all the IO read methods return a binary
+  # string when given a maximum read length
+  def chunk_out(buf)
+    [ "#{buf.size.to_s(16)}\r\n", buf, "\r\n".freeze ]
+  end
+
+  def trailer_out(tlr)
+    "0\r\n#{tlr.map! do |k,v| "#{k}: #{v}\r\n" end.join}\r\n"
+  end
 end
diff --git a/lib/yahns/proxy_pass.rb b/lib/yahns/proxy_pass.rb
index ec1fbcf..3812dbf 100644
--- a/lib/yahns/proxy_pass.rb
+++ b/lib/yahns/proxy_pass.rb
@@ -12,6 +12,7 @@ require_relative 'proxy_http_response'
 class Yahns::ProxyPass # :nodoc:
   class ReqRes < Kgio::Socket
     attr_writer :resbuf
+    attr_accessor :proxy_trailers
 
     def req_start(c, req, input, chunked)
       @hdr = @resbuf = nil