Added http client buffer size (#249)

http_client.go

@@ -3,7 +3,6 @@ package main
 import (
 	"bytes"
 	"crypto/tls"
-	"github.com/buger/gor/proto"
 	"io"
 	"log"
 	"net"
@@ -14,6 +13,8 @@ import (
 	"sync"
 	"syscall"
 	"time"
+
+	"github.com/buger/gor/proto"
 )
 
 var httpMu sync.Mutex

output_http.go

@@ -1,10 +1,11 @@
 package main
 
 import (
-	"github.com/buger/gor/proto"
 	"io"
 	"sync/atomic"
 	"time"
+
+	"github.com/buger/gor/proto"
 )
 
 const initialDynamicWorkers = 10
@@ -26,6 +27,7 @@ type HTTPOutputConfig struct {
 
 	Timeout      time.Duration
 	OriginalHost bool
+	BufferSize   int
 
 	Debug bool
 
@@ -109,10 +111,11 @@ func (o *HTTPOutput) workerMaster() {
 
 func (o *HTTPOutput) startWorker() {
 	client := NewHTTPClient(o.address, &HTTPClientConfig{
-		FollowRedirects: o.config.redirectLimit,
-		Debug:           o.config.Debug,
-		OriginalHost:    o.config.OriginalHost,
-		Timeout:         o.config.Timeout,
+		FollowRedirects:    o.config.redirectLimit,
+		Debug:              o.config.Debug,
+		OriginalHost:       o.config.OriginalHost,
+		Timeout:            o.config.Timeout,
+		ResponseBufferSize: o.config.BufferSize,
 	})
 
 	deathCount := 0

settings.go

@@ -90,6 +90,7 @@ func init() {
 	flag.Var(&Settings.inputHTTP, "input-http", "Read requests from HTTP, should be explicitly sent from your application:\n\t# Listen for http on 9000\n\tgor --input-http :9000 --output-http staging.com")
 
 	flag.Var(&Settings.outputHTTP, "output-http", "Forwards incoming requests to given http address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --output-http http://staging.com")
+	flag.IntVar(&Settings.outputHTTPConfig.BufferSize, "output-http-buffer", 0, "HTTP response buffer size, all data after this size will be discarded.")
 	flag.IntVar(&Settings.outputHTTPConfig.workers, "output-http-workers", 0, "Gor uses dynamic worker scaling by default.  Enter a number to run a set number of workers.")
 	flag.IntVar(&Settings.outputHTTPConfig.redirectLimit, "output-http-redirects", 0, "Enable how often redirects should be followed.")
 	flag.DurationVar(&Settings.outputHTTPConfig.Timeout, "output-http-timeout", 0, "Specify HTTP request/response timeout. By default 5s. Example: --output-http-timeout 30s")