Merge branch 'master' into input-modifier

.travis.yml

@@ -1,3 +1,3 @@
 language: go 
 go: 1.3.3
-script: sudo -E bash -c "source /etc/profile && gvm use go1.3.3 && export GOPATH=$HOME/gopath:$GOPATH && go get && GORACE='halt_on_error=1' go test -race -v"
+script: sudo -E bash -c "source /etc/profile && eval '$(gimme 1.4)' && export GOPATH=$HOME/gopath:$GOPATH && go get && GORACE='halt_on_error=1' go test -race -v"

README.md

@@ -34,7 +34,7 @@ It's recommended to use separate server for replaying traffic, but if you have e
 sudo gor --input-raw :80 --output-http "http://staging.com"
 ```
 
-### Guarante of replay and HTTP input
+### Guarantee of replay and HTTP input
 Due to how traffic interception works, there is chance of missing requests. If you want guarantee that requests will be replayed you can use http input, but it will require changes in your app as well. 
 
 ```
@@ -43,6 +43,13 @@ sudo gor --input-http :28019 --output-http "http://staging.com"
 
 Then in your application you should send copy (e.g. like reverse proxy) all incoming requests to Gor http input. 
 
+### Following redirects
+If you have a scenario where following redirects is usefull you can do it like with:
+
+```
+gor --input-tcp replay.local:28020 --output-http http://staging.com --output-http-redirects 10
+```
+The given example will follow up to 10 redirects per request.
 
 ## Advanced use
 
@@ -186,59 +193,49 @@ https://github.com/buger/gor/releases
 `gor -h` output:
 ```
   -cpuprofile="": write cpu profile to file
-  -memprofile="": write memory profile to this file
-  
   -input-dummy=[]: Used for testing outputs. Emits 'Get /' request every 1s
-
   -input-file=[]: Read requests from file: 
-    gor --input-file ./requests.gor --output-http staging.com
-
+	gor --input-file ./requests.gor --output-http staging.com
+  -input-http=[]: Read requests from HTTP, should be explicitly sent from your application:
+	# Listen for http on 9000
+	gor --input-http :9000 --output-http staging.com
   -input-raw=[]: Capture traffic from given port (use RAW sockets and require *sudo* access):
-    # Capture traffic from 8080 port
-    gor --input-raw :8080 --output-http staging.com
-
+	# Capture traffic from 8080 port
+	gor --input-raw :8080 --output-http staging.com
   -input-tcp=[]: Used for internal communication between Gor instances. Example: 
-    # Receive requests from other Gor instances on 28020 port, and redirect output to staging
-    gor --input-tcp :28020 --output-http staging.com
-
+	# Receive requests from other Gor instances on 28020 port, and redirect output to staging
+	gor --input-tcp :28020 --output-http staging.com
+  -memprofile="": write memory profile to this file
   -output-dummy=[]: Used for testing inputs. Just prints data coming from inputs.
-
   -output-file=[]: Write incoming requests to file: 
-    gor --input-raw :80 --output-file ./requests.gor
-
+	gor --input-raw :80 --output-file ./requests.gor
   -output-http=[]: Forwards incoming requests to given http address.
-    # Redirect all incoming requests to staging.com address 
-    gor --input-raw :80 --output-http http://staging.com
-
+	# Redirect all incoming requests to staging.com address 
+	gor --input-raw :80 --output-http http://staging.com
   -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'
-
+	gor --input-raw :8080 --output-http staging.com --output-http-elasticsearch 'es_host:api_port/index_name'
   -output-http-header=[]: Inject additional headers to http reqest:
-    gor --input-raw :8080 --output-http staging.com --output-http-header 'User-Agent: Gor'
-
+	gor --input-raw :8080 --output-http staging.com --output-http-header 'User-Agent: Gor'
   -output-http-header-filter=[]: A regexp to match a specific header against. Requests with non-matching headers will be dropped:
-    gor --input-raw :8080 --output-http staging.com --output-http-header-filter api-version:^v1
-
+	 gor --input-raw :8080 --output-http staging.com --output-http-header-filter api-version:^v1
   -output-http-header-hash-filter=[]: Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific header. The fraction must have a denominator that is a power of two:
-    gor --input-raw :8080 --output-http staging.com --output-http-header-hash-filter user-id:1/4
-
+	 gor --input-raw :8080 --output-http staging.com --output-http-header-hash-filter user-id:1/4
+  -output-http-method=[]: Whitelist of HTTP methods to replay. Anything else will be dropped:
+	gor --input-raw :8080 --output-http staging.com --output-http-method GET --output-http-method OPTIONS
+  -output-http-redirects=0: Enable how often redirects should be followed.
+  -output-http-rewrite-url=[]: Rewrite the requst url based on a mapping:
+	gor --input-raw :8080 --output-http staging.com --output-http-rewrite-url /xml_test/interface.php:/api/service.do
+  -output-http-stats=false: Report http output queue stats to console every 5 seconds.
   -output-http-url-regexp=: A regexp to match requests against. Anything else will be dropped:
-    gor --input-raw :8080 --output-http staging.com --output-http-url-regexp ^www.
-    
-  -output-http-workers=-1: Number of http output workers desired. Use default -1 for dynamic worker scaling.  Gor will add http workers if its work queue starts getting too full and kill them .
-    
-  -output-http-stats=false: If set to `true` it gives out queuing stats for the HTTP output every 5 seconds in the form latest,mean,max,count,count/second.
-    
+	 gor --input-raw :8080 --output-http staging.com --output-http-url-regexp ^www.
+  -output-http-workers=-1: 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: If set to `true` it gives out queuing stats for the TCP output every 5 seconds in the form latest,mean,max,count,count/second.
-    
+	# 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.
   -split-output=false: By default each output gets same traffic. If set to `true` it splits traffic equally among all outputs.
-
-  -output-http-rewrite-url=[]: Rewrites the url in the request based on a mapping
-    gor --input-raw :8080 --output-http staging.com --output-http-rewrite-url /xml_test/interface.php:/api/service.do
+  -stats=false: Turn on queue stats output
+  -verbose=false: Turn on verbose/debug output
 ```
 
 ## Building from source

input_tcp.go

@@ -62,6 +62,12 @@ func (i *TCPInput) handleConnection(conn net.Conn) {
 
 	for {
 		buf, err := reader.ReadBytes('¶')
+		if err == io.EOF {
+			return
+		} else if err != nil {
+			log.Println("Unexpected error in input tcp connection", err)
+			return
+		}
 		buf_len := len(buf)
 		if buf_len > 0 {
 			new_buf_len := len(buf) - 2
@@ -69,11 +75,6 @@ func (i *TCPInput) handleConnection(conn net.Conn) {
 				new_buf := make([]byte, new_buf_len)
 				copy(new_buf, buf[:new_buf_len])
 				i.data <- new_buf
-				if err != nil {
-					if err != io.EOF {
-						log.Printf("error: %s\n", err)
-					}
-				}
 			}
 		}
 	}

output_http.go

@@ -4,8 +4,10 @@ import (
 	"bufio"
 	"bytes"
 	"io"
+	"io/ioutil"
 	"log"
 	"net/http"
+	"net/http/httputil"
 	"net/url"
 	"strings"
 	"sync/atomic"
@@ -19,8 +21,8 @@ func (e *RedirectNotAllowed) Error() string {
 }
 
 // customCheckRedirect disables redirects https://github.com/buger/gor/pull/15
-func customCheckRedirect(req *http.Request, via []*http.Request) error {
-	if len(via) >= 0 {
+func (o *HTTPOutput) customCheckRedirect(req *http.Request, via []*http.Request) error {
+	if len(via) >= o.redirectLimit {
 		return new(RedirectNotAllowed)
 	}
 	return nil
@@ -28,23 +30,53 @@ 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 err != nil {
+		return
+	}
+
+	if request.Method == "POST" {
+		// 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())
+	}
+
 	return
 }
 
 const InitialDynamicWorkers = 10
 
 type HTTPOutput struct {
+	// Keep this as first element of struct because it guarantees 64bit
+	// alignment. atomic.* functions crash on 32bit machines if operand is not
+	// aligned at 64bit. See https://github.com/golang/go/issues/599
+	activeWorkers int64
+
 	address string
 	limit   int
 	queue   chan []byte
 
-	activeWorkers int64
-	needWorker    chan int
+	redirectLimit int
+
+	needWorker chan int
 
 	urlRegexp            HTTPUrlRegexp
 	headerFilters        HTTPHeaderFilters
@@ -59,7 +91,7 @@ type HTTPOutput struct {
 	queueStats *GorStat
 }
 
-func NewHTTPOutput(address string, headers HTTPHeaders, methods HTTPMethods, urlRegexp HTTPUrlRegexp, headerFilters HTTPHeaderFilters, headerHashFilters HTTPHeaderHashFilters, elasticSearchAddr string, outputHTTPUrlRewrite UrlRewriteMap) io.Writer {
+func NewHTTPOutput(address string, headers HTTPHeaders, methods HTTPMethods, urlRegexp HTTPUrlRegexp, headerFilters HTTPHeaderFilters, headerHashFilters HTTPHeaderHashFilters, elasticSearchAddr string, outputHTTPUrlRewrite UrlRewriteMap, outputHTTPRedirects int) io.Writer {
 
 	o := new(HTTPOutput)
 
@@ -71,6 +103,8 @@ func NewHTTPOutput(address string, headers HTTPHeaders, methods HTTPMethods, url
 	o.headers = headers
 	o.methods = methods
 
+	o.redirectLimit = Settings.outputHTTPRedirects
+
 	o.urlRegexp = urlRegexp
 	o.headerFilters = headerFilters
 	o.headerHashFilters = headerHashFilters
@@ -116,7 +150,7 @@ func (o *HTTPOutput) WorkerMaster() {
 
 func (o *HTTPOutput) Worker() {
 	client := &http.Client{
-		CheckRedirect: customCheckRedirect,
+		CheckRedirect: o.customCheckRedirect,
 	}
 
 	death_count := 0

output_http_test.go

@@ -2,8 +2,11 @@ package main
 
 import (
 	"io"
+	"io/ioutil"
 	"net"
 	"net/http"
+	"net/http/httputil"
+	_ "strings"
 	"sync"
 	"testing"
 	"time"
@@ -11,7 +14,7 @@ import (
 
 func startHTTP(cb func(*http.Request)) net.Listener {
 	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		go cb(r)
+		cb(r)
 	})
 
 	listener, _ := net.Listen("tcp", ":0")
@@ -60,10 +63,20 @@ func TestHTTPOutput(t *testing.T) {
 			t.Error("Wrong method")
 		}
 
+		if req.Method == "POST" {
+			defer req.Body.Close()
+			body, _ := ioutil.ReadAll(req.Body)
+
+			if string(body) != "a=1&b=2" {
+				buf, _ := httputil.DumpRequest(req, true)
+				t.Error("Wrong POST body:", string(buf))
+			}
+		}
+
 		wg.Done()
 	})
 
-	output := NewHTTPOutput(listener.Addr().String(), headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{})
+	output := NewHTTPOutput(listener.Addr().String(), headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{}, 0)
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
@@ -82,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{}, 0)
+
+	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)
@@ -96,7 +145,7 @@ func BenchmarkHTTPOutput(b *testing.B) {
 		wg.Done()
 	})
 
-	output := NewHTTPOutput(listener.Addr().String(), headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{})
+	output := NewHTTPOutput(listener.Addr().String(), headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{}, 0)
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}

output_tcp.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"log"
 	"net"
+	"time"
 )
 
 type TCPOutput struct {
@@ -32,11 +33,20 @@ func NewTCPOutput(address string) io.Writer {
 }
 
 func (o *TCPOutput) worker() {
-	conn, _ := o.connect(o.address)
+	conn, err := o.connect(o.address)
+	for ; err != nil; conn, err = o.connect(o.address) {
+		time.Sleep(2 * time.Second)
+	}
+
 	defer conn.Close()
 
 	for {
-		conn.Write(<-o.buf)
+		_, err := conn.Write(<-o.buf)
+		if err != nil {
+			log.Println("Worker failed on write, exitings and starting new worker")
+			go o.worker()
+			break
+		}
 	}
 }
 

plugins.go

@@ -101,6 +101,6 @@ func InitPlugins() {
 	}
 
 	for _, options := range Settings.outputHTTP {
-		registerPlugin(NewHTTPOutput, options, Settings.outputHTTPHeaders, Settings.outputHTTPMethods, Settings.outputHTTPUrlRegexp, Settings.outputHTTPHeaderFilters, Settings.outputHTTPHeaderHashFilters, Settings.outputHTTPElasticSearch, Settings.outputHTTPUrlRewrite)
+		registerPlugin(NewHTTPOutput, options, Settings.outputHTTPHeaders, Settings.outputHTTPMethods, Settings.outputHTTPUrlRegexp, Settings.outputHTTPHeaderFilters, Settings.outputHTTPHeaderHashFilters, Settings.outputHTTPElasticSearch, Settings.outputHTTPUrlRewrite, Settings.outputHTTPRedirects)
 	}
 }

raw_socket_listener/listener.go

@@ -61,12 +61,13 @@ func (t *Listener) listen() {
 
 func (t *Listener) readRAWSocket() {
     conn, e := net.ListenPacket("ip4:tcp", t.addr)
-    defer conn.Close()
 
     if e != nil {
         log.Fatal(e)
     }
 
+    defer conn.Close()
+
     buf := make([]byte, 4096*2)
 
     for {

settings.go

@@ -8,7 +8,7 @@ import (
 )
 
 const (
-	VERSION = "0.9.2"
+	VERSION = "0.9.4"
 )
 
 type AppSettings struct {
@@ -42,6 +42,7 @@ type AppSettings struct {
 	outputHTTPElasticSearch     string
 	outputHTTPWorkers           int
 	outputHTTPStats             bool
+	outputHTTPRedirects         int
 }
 
 var Settings AppSettings = AppSettings{}
@@ -87,6 +88,7 @@ func init() {
 
 	flag.StringVar(&Settings.outputHTTPElasticSearch, "output-http-elasticsearch", "", "Send request and response stats to ElasticSearch:\n\tgor --input-raw :8080 --output-http staging.com --output-http-elasticsearch 'es_host:api_port/index_name'")
 	flag.Var(&Settings.outputHTTPUrlRewrite, "output-http-rewrite-url", "Rewrite the requst url based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --output-http-rewrite-url /xml_test/interface.php:/api/service.do")
+	flag.IntVar(&Settings.outputHTTPRedirects, "output-http-redirects", 0, "Enable how often redirects should be followed.")
 }
 
 func Debug(args ...interface{}) {

settings_header_hash_filters.go

@@ -39,10 +39,8 @@ func (h *HTTPHeaderHashFilters) Set(value string) error {
 		panic("need positive numerators and denominators, with the former less than the latter.")
 	}
 
-	for test := den; test != 1; test /= 2 {
-		if test%2 == 1 {
-			return errors.New("must have a denominator which is a power of two.")
-		}
+	if den&(den-1) != 0 {
+		return errors.New("must have a denominator which is a power of two.")
 	}
 
 	var f headerHashFilter

settings_header_hash_filters_test.go

@@ -23,6 +23,11 @@ func TestHTTPHeaderHashFilters(t *testing.T) {
 		t.Error("Should error on HeaderIrrelevant:1/3")
 	}
 
+	err = filters.Set("Pow2Denom:1/31")
+	if err == nil {
+		t.Error("Should error on Pow2Denom:1/31")
+	}
+
 	req := http.Request{}
 	req.Header = make(map[string][]string)
 	req.Header.Add("Header1", "test3414")

settings_url_map.go

@@ -3,11 +3,12 @@ package main
 import (
 	"errors"
 	"fmt"
+	"regexp"
 	"strings"
 )
 
 type urlRewrite struct {
-	src    string
+	src    *regexp.Regexp
 	target string
 }
 
@@ -22,14 +23,18 @@ func (r *UrlRewriteMap) Set(value string) error {
 	if len(valArr) < 2 {
 		return errors.New("need both src and target, colon-delimited (ex. /a:/b).")
 	}
-	*r = append(*r, urlRewrite{src: valArr[0], target: valArr[1]})
+	regexp, err := regexp.Compile(valArr[0])
+	if err != nil {
+		return err
+	}
+	*r = append(*r, urlRewrite{src: regexp, target: valArr[1]})
 	return nil
 }
 
 func (r *UrlRewriteMap) Rewrite(path string) string {
 	for _, f := range *r {
-		if f.src == path {
-			return f.target
+		if f.src.MatchString(path) {
+			path = f.src.ReplaceAllString(path, f.target)
 		}
 	}
 	return path

settings_url_map_test.go

@@ -4,7 +4,7 @@ import (
 	"testing"
 )
 
-func TestUrlRewriteMap(t *testing.T) {
+func TestUrlRewriteMap_1(t *testing.T) {
 	var url string
 
 	rewrites := UrlRewriteMap{}
@@ -15,7 +15,6 @@ func TestUrlRewriteMap(t *testing.T) {
 	}
 
 	url = "/abc"
-
 	if rewrites.Rewrite(url) == url {
 		t.Error("Request url should have been rewritten, wasn't")
 	}
@@ -25,3 +24,29 @@ func TestUrlRewriteMap(t *testing.T) {
 		t.Error("Request url should not have been rewritten, was")
 	}
 }
+
+func TestUrlRewriteMap_2(t *testing.T) {
+	var url string
+
+	rewrites := UrlRewriteMap{}
+
+	err := rewrites.Set("/v1/user/([^\\/]+)/ping:/v2/user/$1/ping")
+	if err != nil {
+		t.Error("Should not error on /v1/user/([^\\/]+)/ping:/v2/user/$1/ping")
+	}
+
+	url = "/v1/user/joe/ping"
+	if rewrites.Rewrite(url) == url {
+		t.Error("Request url should have been rewritten, wasn't")
+	}
+
+	url = "/v1/user/joe/ping"
+	if rewrites.Rewrite(url) != "/v2/user/joe/ping" {
+		t.Error("Request url should have been rewritten, wasn't")
+	}
+
+	url = "/v1/user/ping"
+	if rewrites.Rewrite(url) != url {
+		t.Error("Request url should not have been rewritten, was")
+	}
+}

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() {