Merge branch 'master' into input-modifier

Conflicts:
	Makefile

.gitignore

@@ -2,3 +2,5 @@
 *.gor
 
 *.out
+
+*.bin

Makefile

@@ -20,7 +20,7 @@ drace:
 	docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor go test ./... $(ARGS) -v -race -timeout 15s
 
 dtest:
-	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go test ./... $(ARGS) -v -timeout 15s
+	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go test ./... $(ARGS) -v -timeout 10s
 
 dcover:
 	docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor go test $(ARGS) -race -v -timeout 15s -coverprofile=coverage.out
@@ -37,7 +37,13 @@ dbench:
 
 # Used mainly for debugging, because docker container do not have access to parent machine ports
 drun:
-	docker run -v `pwd`:/gopath/src/gor -t -i gor go run $(SOURCE) --input-modifier="bash ./examples/echo_modifier.sh" --input-dummy=0 --output-http="http://localhost:9000"  --verbose
+	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go run $(SOURCE) --input-dummy=0 --output-http="http://localhost:9000"  --verbose
+
+drecord:
+	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go run $(SOURCE) --input-dummy=0 --output-file=requests.bin --verbose
+
+dreplay:
+	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go run $(SOURCE) --input-file=requests.bin --output-tcp=:9000 --verbose -h
 
 dbash:
 	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor /bin/bash

README.md

@@ -71,6 +71,12 @@ gor --input-tcp replay.local:28020 --output-http http://staging.com --output-htt
 ```
 The given example will follow up to 2 redirects per request.
 
+### HTTP timeouts
+By default http timeout for both request and response is 5 seconds. You can override it like this:
+```
+gor --input-tcp replay.local:28020 --output-http http://staging.com --output-http-timeout 30s
+```
+
 ### Rate limiting
 Rate limiting can be useful if you want forward only part of production traffic and not overload your staging environment. There is 2 strategies: dropping random requests or dropping fraction of requests based on Header or URL param value. 
 
@@ -303,14 +309,14 @@ https://github.com/buger/gor/releases
   -output-http-elasticsearch="": Send request and response stats to ElasticSearch:
   gor --input-raw :8080 --output-http staging.com --output-http-elasticsearch 'es_host:api_port/index_name'
   -output-http-header-filter=[]: WARNING: `--output-http-header-filter` DEPRECATED, use `--http-allow-header` instead  -output-http-redirects=0: Enable how often redirects should be followed.
-  -output-http-stats=false: Report http output queue stats to console every 5 seconds.
+  -output-http-stats=false: Report http output queue stats to console every 5 seconds. Remember to include also `--stats`
   -output-http-workers=0: Gor uses dynamic worker scaling by default.  Enter a number to run a set number of workers.
   -output-tcp=[]: Used for internal communication between Gor instances. Example:
   # Listen for requests on 80 port and forward them to other Gor instance on 28020 port
   gor --input-raw :80 --output-tcp replay.local:28020
-  -output-tcp-stats=false: Report TCP output queue stats to console every 5 seconds.
+  -output-tcp-stats=false: Report TCP output queue stats to console every 5 seconds. Remember to include also `--stats`
   -split-output=false: By default each output gets same traffic. If set to `true` it splits traffic equally among all outputs.
-  -stats=false: Turn on queue stats output
+  -stats=false: Turn on queue stats output. Use in combination with the other *-stats flags.
   -verbose=false: Turn on verbose/debug output
 ```
 

http_client.go

@@ -20,6 +20,9 @@ var defaultPorts = map[string]string{
 type HTTPClientConfig struct {
 	FollowRedirects int
 	Debug           bool
+	OriginalHost    bool
+	Timeout         time.Duration
+	ResponseBufferSize  int
 }
 
 type HTTPClient struct {
@@ -42,11 +45,19 @@ func NewHTTPClient(baseURL string, config *HTTPClientConfig) *HTTPClient {
 		u.Host += ":" + defaultPorts[u.Scheme]
 	}
 
+	if config.Timeout.Nanoseconds() == 0 {
+		config.Timeout = 5 * time.Second
+	}
+
+	if config.ResponseBufferSize == 0 {
+		config.ResponseBufferSize = 512*1024 // 500kb
+	}
+
 	client := new(HTTPClient)
 	client.baseURL = u.String()
 	client.host = u.Host
 	client.scheme = u.Scheme
-	client.respBuf = make([]byte, 4096*10)
+	client.respBuf = make([]byte, config.ResponseBufferSize)
 	client.config = config
 
 	return client
@@ -111,11 +122,13 @@ func (c *HTTPClient) Send(data []byte) (response []byte, err error) {
 		}
 	}
 
-	timeout := time.Now().Add(5 * time.Second)
+	timeout := time.Now().Add(c.config.Timeout)
 
 	c.conn.SetWriteDeadline(timeout)
 
-	data = proto.SetHost(data, []byte(c.baseURL), []byte(c.host))
+	if !c.config.OriginalHost {
+		data = proto.SetHost(data, []byte(c.baseURL), []byte(c.host))
+	}
 
 	if c.config.Debug {
 		Debug("[HTTPClient] Sending:", string(data))
@@ -129,6 +142,16 @@ func (c *HTTPClient) Send(data []byte) (response []byte, err error) {
 	c.conn.SetReadDeadline(timeout)
 	n, err := c.conn.Read(c.respBuf)
 
+	// If response large then our buffer, we need to read all response buffer
+	// Otherwise it will corrupt response of next request
+	// Parsing response body is non trivial thing, especially with keep-alive
+	// Simples case is to to close connection if response too large
+	//
+	// See https://github.com/buger/gor/issues/184
+	if n == len(c.respBuf) {
+		c.Disconnect()
+	}
+
 	if err != nil {
 		Debug("[HTTPClient] Response read error", err, c.conn)
 		return
@@ -136,6 +159,8 @@ func (c *HTTPClient) Send(data []byte) (response []byte, err error) {
 
 	payload := c.respBuf[:n]
 
+	Debug("[HTTPClient] Received:", n)
+
 	if c.config.Debug {
 		Debug("[HTTPClient] Received:", string(payload))
 	}

http_client_test.go

@@ -1,6 +1,8 @@
 package main
 
 import (
+	"bytes"
+	"crypto/rand"
 	"io/ioutil"
 	"net"
 	"net/http"
@@ -81,6 +83,36 @@ func TestHTTPClientSend(t *testing.T) {
 	wg.Wait()
 }
 
+// https://github.com/buger/gor/issues/184
+func TestHTTPClientResponseBuffer(t *testing.T) {
+	wg := new(sync.WaitGroup)
+
+	payload := []byte("GET / HTTP/1.1\r\n\r\n")
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+
+		size := 10 * 1024 // 10kb
+		rb := make([]byte, size)
+		rand.Read(rb)
+
+		w.Write(rb)
+
+		wg.Done()
+	}))
+
+	client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: false, ResponseBufferSize: 1024})
+
+	wg.Add(2)
+	client.Send(payload)
+	resp, _ := client.Send(payload)
+
+	if !bytes.Equal(resp[0:8], []byte("HTTP/1.1")) {
+		t.Error("Response buffer contains data from previous request", string(resp[0:5]))
+	}
+
+	wg.Wait()
+}
+
 func TestHTTPClientHTTPSSend(t *testing.T) {
 	wg := new(sync.WaitGroup)
 

input_tcp.go

@@ -67,7 +67,10 @@ func (i *TCPInput) handleConnection(conn net.Conn) {
 		encodedPayload := scanner.Bytes()
 		// Hex encoding always 2x number of bytes
 		decoded := make([]byte, len(encodedPayload)/2)
-		hex.Decode(decoded, encodedPayload)
+		_, err := hex.Decode(decoded, encodedPayload)
+		if err != nil {
+			log.Println("[TCPInput] failed to hex decode TCP payload:", err)
+		}
 		i.data <- decoded
 	}
 

output_http.go

@@ -18,6 +18,9 @@ type HTTPOutputConfig struct {
 
 	elasticSearch string
 
+	Timeout      time.Duration
+	OriginalHost bool
+
 	Debug bool
 }
 
@@ -94,6 +97,8 @@ 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,
 	})
 
 	deathCount := 0

output_http_test.go

@@ -76,6 +76,40 @@ func TestHTTPOutput(t *testing.T) {
 	Settings.modifierConfig = HTTPModifierConfig{}
 }
 
+func TestHTTPOutputKeepOriginalHost(t *testing.T) {
+	wg := new(sync.WaitGroup)
+	quit := make(chan int)
+
+	input := NewTestInput()
+
+	listener := startHTTP(func(req *http.Request) {
+		if req.Host != "custom-host.com" {
+			t.Error("Wrong header", req.Host)
+		}
+
+		wg.Done()
+	})
+
+	headers := HTTPHeaders{HTTPHeader{"Host", "custom-host.com"}}
+	Settings.modifierConfig = HTTPModifierConfig{headers: headers}
+
+	output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{Debug: false, OriginalHost: true})
+
+	Plugins.Inputs = []io.Reader{input}
+	Plugins.Outputs = []io.Writer{output}
+
+	go Start(quit)
+
+	wg.Add(1)
+	input.EmitGET()
+
+	wg.Wait()
+
+	close(quit)
+
+	Settings.modifierConfig = HTTPModifierConfig{}
+}
+
 func TestOutputHTTPSSL(t *testing.T) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)

output_tcp.go

@@ -60,7 +60,8 @@ func (o *TCPOutput) Write(data []byte) (n int, err error) {
 	// Hex encoding always 2x number of bytes
 	encoded := make([]byte, len(data)*2+1)
 	hex.Encode(encoded, data)
-	o.buf <- append(encoded, '\n')
+	encoded[len(encoded)-1] = '\n'
+	o.buf <- encoded
 
 	if Settings.outputTCPStats {
 		o.bufStats.Write(len(o.buf))

plugins.go

@@ -98,6 +98,15 @@ func InitPlugins() {
 		registerPlugin(NewHTTPInput, options)
 	}
 
+	// If we explicitly set Host header http output should not rewrite it
+	// Fix: https://github.com/buger/gor/issues/174
+	for _, header := range Settings.modifierConfig.headers {
+		if header.Name == "Host" {
+			Settings.outputHTTPConfig.OriginalHost = true
+			break
+		}
+	}
+
 	for _, options := range Settings.outputHTTP {
 		registerPlugin(NewHTTPOutput, options, &Settings.outputHTTPConfig)
 	}

settings.go

@@ -91,6 +91,7 @@ func init() {
 	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.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")
 
 	flag.BoolVar(&Settings.outputHTTPConfig.stats, "output-http-stats", false, "Report http output queue stats to console every 5 seconds.")
 
@@ -126,7 +127,7 @@ func init() {
 // Debug gets called only if --verbose flag specified
 func Debug(args ...interface{}) {
 	if Settings.verbose {
-		fmt.Print("[DEBUG] ")
+		fmt.Printf("[DEBUG][PID %d] ", os.Getpid())
 		log.Println(args...)
 	}
 }