[PATCH] wbuf: fix writev calls for vectors

[PATCH] wbuf: fix writev calls for vectors

[Eric Wong]

kgio_writev returns nil on success instead of the number of bytes
written, so we must manually calculate the number of bytes written
intead :x

This is triggerable when sending giant chunked responses.
---
 lib/yahns/wbuf.rb       |  7 ++++++-
 test/test_proxy_pass.rb | 27 +++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/lib/yahns/wbuf.rb b/lib/yahns/wbuf.rb
index 42776cf..9888a8c 100644
--- a/lib/yahns/wbuf.rb
+++ b/lib/yahns/wbuf.rb
@@ -41,6 +41,11 @@ class Yahns::Wbuf # :nodoc:
     @busy = busy # may be false
   end
 
+  def wbuf_writev(buf)
+    @tmpio.kgio_writev(buf)
+    buf.inject(0) { |n, s| n += s.size }
+  end
+
   def wbuf_write(c, buf)
     # try to bypass the VFS layer and write directly to the socket
     # if we're all caught up
@@ -54,7 +59,7 @@ class Yahns::Wbuf # :nodoc:
     end until @busy
 
     @tmpio ||= Yahns::TmpIO.new(@tmpdir)
-    @sf_count += String === buf ? @tmpio.write(buf) : @tmpio.kgio_writev(buf)
+    @sf_count += String === buf ? @tmpio.write(buf) : wbuf_writev(buf)
 
     # we spent some time copying to the FS, try to write to
     # the socket again in case some space opened up...
diff --git a/test/test_proxy_pass.rb b/test/test_proxy_pass.rb
index f5ae12f..e4e97d1 100644
--- a/test/test_proxy_pass.rb
+++ b/test/test_proxy_pass.rb
@@ -2,6 +2,7 @@
 # License: GPLv3 or later (https://www.gnu.org/licenses/gpl-3.0.txt)
 require_relative 'server_helper'
 require 'json'
+require 'digest'
 
 class TestProxyPass < Testcase
   ENV["N"].to_i > 1 and parallelize_me!
@@ -17,6 +18,8 @@ class TestProxyPass < Testcase
   BIG_HEADER = [%w(Content-Type text/plain), %W(Content-Length #{OMFG.size})]
   3000.times { |i| BIG_HEADER << %W(X-#{i} BIG-HEADER!!!!!!!!!!!!!!) }
   BIG_HEADER.freeze
+  STR4 = 'abcd' * (256 * 1024)
+  NCHUNK = 50
 
   class ProxiedApp
     def call(env)
@@ -27,6 +30,19 @@ class TestProxyPass < Testcase
         when '/giant-body'
           h = [ %W(Content-Length #{OMFG.size}), %w(Content-Type text/plain) ]
           [ 200, h, [ OMFG ] ]
+        when '/giant-chunky-body'
+          h = [ %w(content-type text/pain), %w(transfer-encoding chunked) ]
+          chunky = Object.new
+          def chunky.each
+            head = STR4.size.to_s(16) << "\r\n"
+            NCHUNK.times do
+              yield head
+              yield STR4
+              yield "\r\n".freeze
+            end
+            yield "0\r\n\r\n"
+          end
+          [ 200, h, chunky ]
         when '/big-headers'
           [ 200, BIG_HEADER, [ OMFG ] ]
         when '/oversize-headers'
@@ -281,6 +297,17 @@ class TestProxyPass < Testcase
       assert_equal 200, res.code.to_i
       assert_equal OMFG, res.body
 
+      # giant chunky body
+      sha1 = Digest::SHA1.new
+      http.request(Net::HTTP::Get.new('/giant-chunky-body')) do |response|
+        response.read_body do |chunk|
+          sha1.update(chunk)
+        end
+      end
+      check = Digest::SHA1.new
+      NCHUNK.times { check.update(STR4) }
+      assert_equal check.hexdigest, sha1.hexdigest
+
       # giant PUT content-length
       req = Net::HTTP::Put.new('/')
       req.body_stream = StringIO.new(OMFG)
-- 
EW