From 2b17989fd91addc200b2947d2d47c5dad5cadec8 Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Sun, 5 Jul 2015 23:17:57 +0500 Subject: [PATCH] Move HTTP modifier from HTTP output plugin --- Makefile | 4 +- emitter.go | 16 ++++- http_modifier.go | 87 ++++++++++++++++++++++++++++ http_modifier_test.go | 38 ++++++++++++ output_http.go | 17 ------ output_http_test.go | 31 ++-------- proto/proto.go | 6 ++ settings.go | 13 +++-- settings_header_filters.go | 16 +---- settings_header_filters_test.go | 42 +++++--------- settings_header_hash_filters.go | 22 +------ settings_header_hash_filters_test.go | 21 +------ settings_methods.go | 9 +-- settings_methods_test.go | 6 +- settings_url_map.go | 15 +---- settings_url_map_test.go | 53 +---------------- 16 files changed, 196 insertions(+), 200 deletions(-) create mode 100644 http_modifier_test.go diff --git a/Makefile b/Makefile index 9ff6068..07d3f7c 100644 --- a/Makefile +++ b/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 \ No newline at end of file diff --git a/emitter.go b/emitter.go index 867fc75..5e7e78e 100644 --- a/emitter.go +++ b/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) } } diff --git a/http_modifier.go b/http_modifier.go index 73b5203..95fb819 100644 --- a/http_modifier.go +++ b/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 } \ No newline at end of file diff --git a/http_modifier_test.go b/http_modifier_test.go new file mode 100644 index 0000000..6ca055d --- /dev/null +++ b/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") + } +} diff --git a/output_http.go b/output_http.go index 14b630e..c1b42f9 100644 --- a/output_http.go +++ b/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() diff --git a/output_http_test.go b/output_http_test.go index 589cab4..7cf8c24 100644 --- a/output_http_test.go +++ b/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) { diff --git a/proto/proto.go b/proto/proto.go index 8a68949..f38645a 100644 --- a/proto/proto.go +++ b/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] +} \ No newline at end of file diff --git a/settings.go b/settings.go index 669d145..411ca62 100644 --- a/settings.go +++ b/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'") diff --git a/settings_header_filters.go b/settings_header_filters.go index ea2b034..7a017a6 100644 --- a/settings_header_filters.go +++ b/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 -} +} \ No newline at end of file diff --git a/settings_header_filters_test.go b/settings_header_filters_test.go index 2d0af19..499f910 100644 --- a/settings_header_filters_test.go +++ b/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:^:$") + } +} \ No newline at end of file diff --git a/settings_header_hash_filters.go b/settings_header_hash_filters.go index f76aa9e..0dc9801 100644 --- a/settings_header_hash_filters.go +++ b/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 -} +} \ No newline at end of file diff --git a/settings_header_hash_filters_test.go b/settings_header_hash_filters_test.go index f725761..171b69b 100644 --- a/settings_header_hash_filters_test.go +++ b/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") - } } diff --git a/settings_methods.go b/settings_methods.go index e83109f..cd5a977 100644 --- a/settings_methods.go +++ b/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 } } diff --git a/settings_methods_test.go b/settings_methods_test.go index 774e952..4c56ed2 100644 --- a/settings_methods_test.go +++ b/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") } } diff --git a/settings_url_map.go b/settings_url_map.go index 5453610..d07a313 100644 --- a/settings_url_map.go +++ b/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 -} +} \ No newline at end of file diff --git a/settings_url_map_test.go b/settings_url_map_test.go index ed5ed21..85f0393 100644 --- a/settings_url_map_test.go +++ b/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 \ No newline at end of file