From 390f75fd3c1ade1a8150d023e288484d40cae79d Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Fri, 26 Jun 2015 09:08:16 +0500
Subject: [PATCH] Handle Transfer-Encoding: chunked

---
 output_http.go      | 16 +++++++++++++++-
 output_http_test.go | 38 +++++++++++++++++++++++++++++++++++++-
 test_input.go       |  6 +++++-
 3 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/output_http.go b/output_http.go
index 4295517..e254ff4 100644
--- a/output_http.go
+++ b/output_http.go
@@ -11,6 +11,7 @@ import (
 	"sync/atomic"
 	"time"
 	"io/ioutil"
+	"net/http/httputil"
 )
 
 type RedirectNotAllowed struct{}
@@ -29,14 +30,27 @@ func customCheckRedirect(req *http.Request, via []*http.Request) error {
 
 // ParseRequest in []byte returns a http request or an error
 func ParseRequest(data []byte) (request *http.Request, err error) {
+	var body []byte
+
+	// Test if request have Transfer-Encoding: chunked
+	isChunked := bytes.Contains(data, []byte(": chunked\r\n"));
+
 	buf := bytes.NewBuffer(data)
 	reader := bufio.NewReader(buf)
 
+	// ReadRequest does not read POST bodies, we have to do it by ourseves
 	request, err = http.ReadRequest(reader)
 
 	if request.Method == "POST" {
-		body, _ := ioutil.ReadAll(reader)
+		// This works, because ReadRequest method modify buffer and strips all headers, leaving only body
+		if isChunked {
+			body, _ = ioutil.ReadAll(httputil.NewChunkedReader(reader))
+		} else {
+			body, _ = ioutil.ReadAll(reader)
+		}
+
 		bodyBuf := bytes.NewBuffer(body)
+
 		request.Body = ioutil.NopCloser(bodyBuf)
 		request.ContentLength = int64(bodyBuf.Len())
 	}
diff --git a/output_http_test.go b/output_http_test.go
index eec343a..5aef0e2 100644
--- a/output_http_test.go
+++ b/output_http_test.go
@@ -67,7 +67,7 @@ func TestHTTPOutput(t *testing.T) {
 			defer req.Body.Close()
 			body, _ := ioutil.ReadAll(req.Body)
 
-			if string(body) != "a=1&b=2\r\n\r\n" {
+			if string(body) != "a=1&b=2" {
 				buf, _ := httputil.DumpRequest(req, true)
 				t.Error("Wrong POST body:", string(buf))
 			}
@@ -95,6 +95,42 @@ func TestHTTPOutput(t *testing.T) {
 	close(quit)
 }
 
+func TestHTTPOutputChunkedEncoding(t *testing.T) {
+	wg := new(sync.WaitGroup)
+	quit := make(chan int)
+
+	input := NewTestInput()
+
+	headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}}
+	methods := HTTPMethods{"GET", "PUT", "POST"}
+
+	listener := startHTTP(func(req *http.Request) {
+		defer req.Body.Close()
+		body, _ := ioutil.ReadAll(req.Body)
+
+		if string(body) != "Wikipedia in\r\n\r\nchunks." {
+			buf, _ := httputil.DumpRequest(req, true)
+			t.Error("Wrong POST body:", buf, body, []byte("Wikipedia in\r\n\r\nchunks."))
+		}
+
+		wg.Done()
+	})
+
+	output := NewHTTPOutput(listener.Addr().String(), headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{})
+
+	Plugins.Inputs = []io.Reader{input}
+	Plugins.Outputs = []io.Writer{output}
+
+	go Start(quit)
+
+	wg.Add(1)
+	input.EmitChunkedPOST()
+
+	wg.Wait()
+
+	close(quit)
+}
+
 func BenchmarkHTTPOutput(b *testing.B) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
diff --git a/test_input.go b/test_input.go
index 9982551..694736c 100644
--- a/test_input.go
+++ b/test_input.go
@@ -28,7 +28,11 @@ func (i *TestInput) EmitGET() {
 }
 
 func (i *TestInput) EmitPOST() {
-	i.data <- []byte("POST /pub/WWW/ HTTP/1.1\nHost: www.w3.org\r\n\r\na=1&b=2\r\n\r\n")
+	i.data <- []byte("POST /pub/WWW/ HTTP/1.1\nHost: www.w3.org\r\n\r\na=1&b=2")
+}
+
+func (i *TestInput) EmitChunkedPOST() {
+	i.data <- []byte("POST /pub/WWW/ HTTP/1.1\nHost: www.w3.org\nTransfer-Encoding: chunked\r\n\r\n4\r\nWiki\r\n5\r\npedia\r\ne\r\n in\r\n\r\nchunks.\r\n0\r\n\r\n")
 }
 
 func (i *TestInput) EmitFile() {