Move HTTP modifier from HTTP output plugin

Makefile

@@ -1,4 +1,4 @@
-SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go settings_header_filters.go settings_header_hash_filters.go settings_headers.go settings_methods.go settings_option.go settings_url_regexp.go test_input.go elasticsearch.go settings_url_map.go
+SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go settings_header_filters.go settings_header_hash_filters.go settings_headers.go settings_methods.go settings_option.go settings_url_regexp.go test_input.go elasticsearch.go settings_url_map.go http_modifier.go
 
 SOURCE_PATH = /gopath/src/github.com/buger/gor/
 
@@ -27,7 +27,7 @@ dbench:
 
 # Used mainly for debugging, because docker container do not have access to parent machine ports
 drun:
-	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go run $(SOURCE) --input-dummy=0 --input-http=:9000 --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
 
 dbash: 
 	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor /bin/bash

emitter.go

@@ -23,13 +23,25 @@ func Start(stop chan int) {
 func CopyMulty(src io.Reader, writers ...io.Writer) (err error) {
 	buf := make([]byte, 5*1024*1024)
 	wIndex := 0
+	modifier := NewHTTPModifier(&Settings.modifierConfig)
 
 	for {
 		nr, er := src.Read(buf)
 		if nr > 0 && len(buf) > nr {
+			payload := buf[0:nr]
+
+			if modifier != nil {
+				payload = modifier.Rewrite(payload)
+
+				// If modifier tells to skip request
+				if len(payload) == 0 {
+					continue
+				}
+			}
+
 			if Settings.splitOutput {
 				// Simple round robin
-				writers[wIndex].Write(buf[0:nr])
+				writers[wIndex].Write(payload)
 
 				wIndex++
 
@@ -38,7 +50,7 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) {
 				}
 			} else {
 				for _, dst := range writers {
-					dst.Write(buf[0:nr])
+					dst.Write(payload)
 				}
 			}
 

http_modifier.go

@@ -1,5 +1,11 @@
 package main
 
+import (
+    "github.com/buger/gor/proto"
+    "hash/fnv"
+)
+
+
 type HTTPModifierConfig struct {
     urlRegexp            HTTPUrlRegexp
     urlRewrite           UrlRewriteMap
@@ -8,4 +14,85 @@ type HTTPModifierConfig struct {
 
     headers HTTPHeaders
     methods HTTPMethods
+}
+
+type HTTPModifier struct {
+    config *HTTPModifierConfig
+}
+
+func NewHTTPModifier(config *HTTPModifierConfig) *HTTPModifier {
+    // Optimization to skip modifier completely if we do not need it
+    if config.urlRegexp.regexp == nil &&
+       len(config.urlRewrite) == 0 &&
+       len(config.headerFilters) == 0 &&
+       len(config.headerHashFilters) == 0 &&
+       len(config.headers) == 0 &&
+       len(config.methods) == 0 {
+        return nil
+    }
+
+    return &HTTPModifier{config: config}
+}
+
+func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) {
+    if len(m.config.methods) > 0 && !m.config.methods.Contains(proto.Method(payload)) {
+        return
+    }
+
+    if m.config.urlRegexp.regexp != nil {
+        host, _, _, _ := proto.Header(payload, []byte("Host"))
+        fullPath := append(host, proto.Path(payload)...)
+
+        if !m.config.urlRegexp.regexp.Match(fullPath) {
+            return
+        }
+    }
+
+    if len(m.config.headerFilters) > 0 {
+        for _, f := range m.config.headerFilters {
+            value, s, _, _ := proto.Header(payload, f.name)
+
+            if s != -1 && !f.regexp.Match(value) {
+                return
+            }
+        }
+    }
+
+    if len(m.config.headerHashFilters) > 0 {
+        for _, f := range m.config.headerHashFilters {
+            value, s, _, _ := proto.Header(payload, f.name)
+
+            if s == -1 {
+                return
+            }
+
+            hasher := fnv.New32a()
+            hasher.Write(value)
+            if hasher.Sum32() > f.maxHash {
+                return
+            }
+        }
+    }
+
+    if len(m.config.urlRewrite) > 0 {
+        path := proto.Path(payload)
+
+        for _, f := range m.config.urlRewrite {
+            if f.src.Match(path) {
+                path = f.src.ReplaceAll(path, f.target)
+                payload = proto.SetPath(payload, path)
+
+                break
+            }
+        }
+    }
+
+    if len(m.config.headers) > 0 {
+        for _, header := range m.config.headers {
+           payload = proto.SetHeader(payload, []byte(header.Name), []byte(header.Value))
+       }
+    }
+
+
+    return payload
 }

http_modifier_test.go

@@ -0,0 +1,38 @@
+package main
+
+import (
+    "testing"
+)
+
+func TestHTTPModifierWithoutConfig(t *testing.T) {
+    if NewHTTPModifier(&HTTPModifierConfig{}) != nil {
+        t.Error("If no config specified should not be initialized")
+    }
+}
+
+func TestHTTPModifierHeaderFilters(t *testing.T) {
+    filters := HTTPHeaderFilters{}
+    filters.Set("Host:^www.w3.org$")
+
+    modifier := NewHTTPModifier(&HTTPModifierConfig{
+        headerFilters: filters,
+    })
+
+    payload := []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
+
+    if len(modifier.Rewrite(payload)) == 0 {
+        t.Error("Request should pass filters")
+    }
+
+    filters = HTTPHeaderFilters{}
+    // Setting filter that not match our header
+    filters.Set("Host:^www.w4.org$")
+
+    modifier = NewHTTPModifier(&HTTPModifierConfig{
+        headerFilters: filters,
+    })
+
+    if len(modifier.Rewrite(payload)) != 0 {
+        t.Error("Request should not pass filters")
+    }
+}

output_http.go

@@ -62,8 +62,6 @@ type HTTPOutputConfig struct {
 	stats bool
 	workers int
 
-	modifier HTTPModifierConfig
-
 	elasticSearch string
 }
 
@@ -204,27 +202,12 @@ func (o *HTTPOutput) sendRequest(client *http.Client, data []byte) {
 		return
 	}
 
-	if len(o.config.modifier.methods) > 0 && !o.config.modifier.methods.Contains(request.Method) {
-		return
-	}
-
-	if !(o.config.modifier.urlRegexp.Good(request) && o.config.modifier.headerFilters.Good(request) && o.config.modifier.headerHashFilters.Good(request)) {
-		return
-	}
-
-	// Rewrite the path as necessary
-	request.URL.Path = o.config.modifier.urlRewrite.Rewrite(request.URL.Path)
-
 	// Change HOST of original request
 	URL := o.address + request.URL.Path + "?" + request.URL.RawQuery
 
 	request.RequestURI = ""
 	request.URL, _ = url.ParseRequestURI(URL)
 
-	for _, header := range o.config.modifier.headers {
-		SetHeader(request, header.Name, header.Value)
-	}
-
 	start := time.Now()
 	resp, err := client.Do(request)
 	stop := time.Now()

output_http_test.go

@@ -24,27 +24,6 @@ func startHTTP(cb func(*http.Request)) net.Listener {
 	return listener
 }
 
-func TestSetHeader(t *testing.T) {
-
-	req := &http.Request{
-		Header: make(map[string][]string),
-	}
-	req.Host = "test.com"
-
-	SetHeader(req, "Host", "test2.com")
-
-	if req.Host != "test2.com" {
-		t.Error("Expected test2.com - got ", req.Host)
-	}
-
-	SetHeader(req, "test_header", "test_value")
-
-	if req.Header.Get("test_header") != "test_value" {
-		t.Error("Wrong header value found")
-	}
-
-}
-
 func TestHTTPOutput(t *testing.T) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
@@ -73,10 +52,10 @@ func TestHTTPOutput(t *testing.T) {
 	})
 
 	headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}}
-	methods := HTTPMethods{"GET", "PUT", "POST"}
-	modifierConfig := HTTPModifierConfig{headers: headers, methods: methods}
+	methods := HTTPMethods{[]byte("GET"), []byte("PUT"), []byte("POST")}
+	Settings.modifierConfig = HTTPModifierConfig{headers: headers, methods: methods}
 
-	output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{modifier: modifierConfig})
+	output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{})
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
@@ -84,7 +63,7 @@ func TestHTTPOutput(t *testing.T) {
 	go Start(quit)
 
 	for i := 0; i < 100; i++ {
-		wg.Add(2)
+		wg.Add(2) // OPTIONS should be ignored
 		input.EmitPOST()
 		input.EmitOPTIONS()
 		input.EmitGET()
@@ -93,6 +72,8 @@ func TestHTTPOutput(t *testing.T) {
 	wg.Wait()
 
 	close(quit)
+
+	Settings.modifierConfig = HTTPModifierConfig{}
 }
 
 func TestOutputHTTPSSL(t *testing.T) {

proto/proto.go

@@ -79,3 +79,9 @@ func SetPath(payload, path []byte) []byte {
 
     return byteutils.Replace(payload, start, start+end, path)
 }
+
+func Method(payload []byte) []byte {
+    end := bytes.IndexByte(payload, ' ')
+
+    return payload[:end]
+}

settings.go

@@ -33,6 +33,7 @@ type AppSettings struct {
 	outputHTTP                  MultiOption
 
 	outputHTTPConfig HTTPOutputConfig
+	modifierConfig HTTPModifierConfig
 }
 
 var Settings AppSettings = AppSettings{}
@@ -68,12 +69,12 @@ 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.Var(&Settings.outputHTTPConfig.modifier.headers, "output-http-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --output-http-header 'User-Agent: Gor'")
-	flag.Var(&Settings.outputHTTPConfig.modifier.methods, "output-http-method", "Whitelist of HTTP methods to replay. Anything else will be dropped:\n\tgor --input-raw :8080 --output-http staging.com --output-http-method GET --output-http-method OPTIONS")
-	flag.Var(&Settings.outputHTTPConfig.modifier.urlRegexp, "output-http-url-regexp", "A regexp to match requests against. Anything else will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --output-http-url-regexp ^www.")
-	flag.Var(&Settings.outputHTTPConfig.modifier.urlRewrite, "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.Var(&Settings.outputHTTPConfig.modifier.headerFilters, "output-http-header-filter", "A regexp to match a specific header against. Requests with non-matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --output-http-header-filter api-version:^v1")
-	flag.Var(&Settings.outputHTTPConfig.modifier.headerHashFilters, "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:\n\t gor --input-raw :8080 --output-http staging.com --output-http-header-hash-filter user-id:1/4")
+	flag.Var(&Settings.modifierConfig.headers, "output-http-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --output-http-header 'User-Agent: Gor'")
+	flag.Var(&Settings.modifierConfig.methods, "output-http-method", "Whitelist of HTTP methods to replay. Anything else will be dropped:\n\tgor --input-raw :8080 --output-http staging.com --output-http-method GET --output-http-method OPTIONS")
+	flag.Var(&Settings.modifierConfig.urlRegexp, "output-http-url-regexp", "A regexp to match requests against. Anything else will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --output-http-url-regexp ^www.")
+	flag.Var(&Settings.modifierConfig.urlRewrite, "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.Var(&Settings.modifierConfig.headerFilters, "output-http-header-filter", "A regexp to match a specific header against. Requests with non-matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --output-http-header-filter api-version:^v1")
+	flag.Var(&Settings.modifierConfig.headerHashFilters, "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:\n\t gor --input-raw :8080 --output-http staging.com --output-http-header-hash-filter user-id:1/4")
 	flag.BoolVar(&Settings.outputHTTPConfig.stats, "output-http-stats", false, "Report http output queue stats to console every 5 seconds.")
 
 	flag.StringVar(&Settings.outputHTTPConfig.elasticSearch, "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'")

settings_header_filters.go

@@ -3,13 +3,12 @@ package main
 import (
 	"errors"
 	"fmt"
-	"net/http"
 	"regexp"
 	"strings"
 )
 
 type headerFilter struct {
-	name   string
+	name   []byte
 	regexp *regexp.Regexp
 }
 
@@ -29,16 +28,7 @@ func (h *HTTPHeaderFilters) Set(value string) error {
 		return err
 	}
 
-	*h = append(*h, headerFilter{name: valArr[0], regexp: r})
+	*h = append(*h, headerFilter{name: []byte(valArr[0]), regexp: r})
 
 	return nil
-}
-
-func (h *HTTPHeaderFilters) Good(req *http.Request) bool {
-	for _, f := range *h {
-		if !f.regexp.MatchString(req.Header.Get(f.name)) {
-			return false
-		}
-	}
-	return true
-}
+}

settings_header_filters_test.go

@@ -1,35 +1,25 @@
 package main
 
 import (
-	"net/http"
-	"testing"
+    "testing"
 )
 
 func TestHTTPHeaderFilters(t *testing.T) {
-	filters := HTTPHeaderFilters{}
+    filters := HTTPHeaderFilters{}
 
-	err := filters.Set("Header1:^$")
-	if err != nil {
-		t.Error("Should not error on Header1:^$")
-	}
+    err := filters.Set("Header1:^$")
+    if err != nil {
+        t.Error("Should not error on Header1:^$")
+    }
 
-	err = filters.Set("Header2:^:$")
-	if err != nil {
-		t.Error("Should not error on Header2:^:$")
-	}
+    err = filters.Set("Header2:^:$")
+    if err != nil {
+        t.Error("Should not error on Header2:^:$")
+    }
 
-	err = filters.Set("Header3-^$")
-	if err == nil {
-		t.Error("Should error on Header2:^:$")
-	}
-
-	req := http.Request{}
-	req.Header = make(map[string][]string)
-	req.Header.Add("Header1", "")
-	req.Header.Add("Header2", ":")
-	req.Header.Add("Header3", "Irrelevant")
-
-	if !filters.Good(&req) {
-		t.Error("Request should pass filters")
-	}
-}
+    // Missing colon
+    err = filters.Set("Header3-^$")
+    if err == nil {
+        t.Error("Should error on Header2:^:$")
+    }
+}

settings_header_hash_filters.go

@@ -3,14 +3,12 @@ package main
 import (
 	"errors"
 	"fmt"
-	"hash/fnv"
-	"net/http"
 	"strconv"
 	"strings"
 )
 
 type headerHashFilter struct {
-	name    string
+	name    []byte
 	maxHash uint32
 }
 
@@ -44,23 +42,9 @@ func (h *HTTPHeaderHashFilters) Set(value string) error {
 	}
 
 	var f headerHashFilter
-	f.name = valArr[0]
+	f.name = []byte(valArr[0])
 	f.maxHash = (uint32)(num * (((uint64)(2 << 31)) / den))
 	*h = append(*h, f)
 
 	return nil
-}
-
-func (h *HTTPHeaderHashFilters) Good(req *http.Request) bool {
-	for _, f := range *h {
-		if req.Header.Get(f.name) == "" {
-			return false
-		}
-		hasher := fnv.New32a()
-		hasher.Write([]byte(req.Header.Get(f.name)))
-		if hasher.Sum32() > f.maxHash {
-			return false
-		}
-	}
-	return true
-}
+}

settings_header_hash_filters_test.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"net/http"
 	"testing"
 )
 
@@ -18,31 +17,15 @@ func TestHTTPHeaderHashFilters(t *testing.T) {
 		t.Error("Should not error on Header2:^:$")
 	}
 
+    // Denominator must be power of 2
 	err = filters.Set("HeaderIrrelevant:1/3")
 	if err == nil {
 		t.Error("Should error on HeaderIrrelevant:1/3")
 	}
 
+    // Denominator must be power of 2
 	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")
-
-	if filters.Good(&req) {
-		t.Error("Request should not pass filters, Header2 does not exist")
-	}
-
-	req.Header.Add("Header2", "test2")
-	if filters.Good(&req) {
-		t.Error("Request should not pass filters, Header2 hash too high")
-	}
-
-	req.Header.Set("Header2", "test3414")
-	if !filters.Good(&req) {
-		t.Error("Request should pass filters")
-	}
 }

settings_methods.go

@@ -3,22 +3,23 @@ package main
 import (
 	"fmt"
 	"strings"
+	"bytes"
 )
 
-type HTTPMethods []string
+type HTTPMethods [][]byte
 
 func (h *HTTPMethods) String() string {
 	return fmt.Sprint(*h)
 }
 
 func (h *HTTPMethods) Set(value string) error {
-	*h = append(*h, strings.ToUpper(value))
+	*h = append(*h, []byte(strings.ToUpper(value)))
 	return nil
 }
 
-func (h *HTTPMethods) Contains(value string) bool {
+func (h *HTTPMethods) Contains(value []byte) bool {
 	for _, method := range *h {
-		if value == method {
+		if bytes.Equal(value, method) {
 			return true
 		}
 	}

settings_methods_test.go

@@ -10,15 +10,15 @@ func TestHTTPMethods(t *testing.T) {
 	methods.Set("lower")
 	methods.Set("UPPER")
 
-	if !methods.Contains("LOWER") {
+	if !methods.Contains([]byte("LOWER")) {
 		t.Error("Does not contain LOWER")
 	}
 
-	if !methods.Contains("UPPER") {
+	if !methods.Contains([]byte("UPPER")) {
 		t.Error("Does not contain UPPER")
 	}
 
-	if methods.Contains("ABSENT") {
+	if methods.Contains([]byte("ABSENT")) {
 		t.Error("Does contain ABSENT")
 	}
 }

settings_url_map.go

@@ -9,7 +9,7 @@ import (
 
 type urlRewrite struct {
 	src    *regexp.Regexp
-	target string
+	target []byte
 }
 
 type UrlRewriteMap []urlRewrite
@@ -27,15 +27,6 @@ func (r *UrlRewriteMap) Set(value string) error {
 	if err != nil {
 		return err
 	}
-	*r = append(*r, urlRewrite{src: regexp, target: valArr[1]})
+	*r = append(*r, urlRewrite{src: regexp, target: []byte(valArr[1]) })
 	return nil
-}
-
-func (r *UrlRewriteMap) Rewrite(path string) string {
-	for _, f := range *r {
-		if f.src.MatchString(path) {
-			path = f.src.ReplaceAllString(path, f.target)
-		}
-	}
-	return path
-}
+}

settings_url_map_test.go

@@ -1,52 +1 @@
-package main
-
-import (
-	"testing"
-)
-
-func TestUrlRewriteMap_1(t *testing.T) {
-	var url string
-
-	rewrites := UrlRewriteMap{}
-
-	err := rewrites.Set("/abc:/123")
-	if err != nil {
-		t.Error("Should not error on /abc:/123")
-	}
-
-	url = "/abc"
-	if rewrites.Rewrite(url) == url {
-		t.Error("Request url should have been rewritten, wasn't")
-	}
-
-	url = "/wibble"
-	if rewrites.Rewrite(url) != url {
-		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")
-	}
-}
+package main