diff --git a/input_raw_test.go b/input_raw_test.go index ba9b537..60bcf49 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -82,9 +82,7 @@ func TestInputRAW100Expect(t *testing.T) { }) replay_address := listener.Addr().String() - headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}} - methods := HTTPMethods{"GET", "PUT", "POST"} - http_output := NewHTTPOutput(replay_address, headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{}, 0) + http_output := NewHTTPOutput(replay_address, &HTTPOutputConfig{}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{test_output, http_output} @@ -141,9 +139,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) { }) replay_address := listener.Addr().String() - headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}} - methods := HTTPMethods{"GET", "PUT", "POST"} - http_output := NewHTTPOutput(replay_address, headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{}, 0) + http_output := NewHTTPOutput(replay_address, &HTTPOutputConfig{}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{test_output, http_output} diff --git a/output_http.go b/output_http.go index 4b1d08c..cd39eed 100644 --- a/output_http.go +++ b/output_http.go @@ -22,7 +22,7 @@ func (e *RedirectNotAllowed) Error() string { // customCheckRedirect disables redirects https://github.com/buger/gor/pull/15 func (o *HTTPOutput) customCheckRedirect(req *http.Request, via []*http.Request) error { - if len(via) >= o.redirectLimit { + if len(via) >= o.config.redirectLimit { return new(RedirectNotAllowed) } return nil @@ -56,6 +56,23 @@ func ParseRequest(data []byte) (request *http.Request, err error) { const InitialDynamicWorkers = 10 +type HTTPOutputConfig struct { + redirectLimit int + + urlRegexp HTTPUrlRegexp + urlRewrite UrlRewriteMap + headerFilters HTTPHeaderFilters + headerHashFilters HTTPHeaderHashFilters + + stats bool + workers int + + headers HTTPHeaders + methods HTTPMethods + + elasticSearch string +} + 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 @@ -66,24 +83,16 @@ type HTTPOutput struct { limit int queue chan []byte - redirectLimit int - needWorker chan int - urlRegexp HTTPUrlRegexp - headerFilters HTTPHeaderFilters - headerHashFilters HTTPHeaderHashFilters - outputHTTPUrlRewrite UrlRewriteMap - - headers HTTPHeaders - methods HTTPMethods - - elasticSearch *ESPlugin + config *HTTPOutputConfig queueStats *GorStat + + elasticSearch *ESPlugin } -func NewHTTPOutput(address string, headers HTTPHeaders, methods HTTPMethods, urlRegexp HTTPUrlRegexp, headerFilters HTTPHeaderFilters, headerHashFilters HTTPHeaderHashFilters, elasticSearchAddr string, outputHTTPUrlRewrite UrlRewriteMap, outputHTTPRedirects int) io.Writer { +func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer { o := new(HTTPOutput) @@ -92,33 +101,25 @@ func NewHTTPOutput(address string, headers HTTPHeaders, methods HTTPMethods, url } o.address = address - o.headers = headers - o.methods = methods + o.config = config - o.redirectLimit = Settings.outputHTTPRedirects - - o.urlRegexp = urlRegexp - o.headerFilters = headerFilters - o.headerHashFilters = headerHashFilters - o.outputHTTPUrlRewrite = outputHTTPUrlRewrite - - o.queue = make(chan []byte, 100) - if Settings.outputHTTPStats { + if o.config.stats { o.queueStats = NewGorStat("output_http") } + o.queue = make(chan []byte, 100) o.needWorker = make(chan int, 1) // Initial workers count - if Settings.outputHTTPWorkers == -1 { + if o.config.workers == 0 { o.needWorker <- InitialDynamicWorkers } else { - o.needWorker <- Settings.outputHTTPWorkers + o.needWorker <- o.config.workers } - if elasticSearchAddr != "" { + if o.config.elasticSearch != "" { o.elasticSearch = new(ESPlugin) - o.elasticSearch.Init(elasticSearchAddr) + o.elasticSearch.Init(o.config.elasticSearch) } go o.WorkerMaster() @@ -134,7 +135,7 @@ func (o *HTTPOutput) WorkerMaster() { } // Disable dynamic scaling if workers poll fixed size - if Settings.outputHTTPWorkers != -1 { + if o.config.workers != 0 { return } } @@ -161,7 +162,7 @@ func (o *HTTPOutput) Worker() { death_count = 0 case <-time.After(time.Millisecond * 100): // When dynamic scaling enabled workers die after 2s of inactivity - if Settings.outputHTTPWorkers == -1 { + if o.config.workers == 0 { death_count += 1 } else { continue @@ -186,11 +187,11 @@ func (o *HTTPOutput) Write(data []byte) (n int, err error) { o.queue <- buf - if Settings.outputHTTPStats { + if o.config.stats { o.queueStats.Write(len(o.queue)) } - if Settings.outputHTTPWorkers == -1 { + if o.config.workers == 0 { workersCount := atomic.LoadInt64(&o.activeWorkers) if len(o.queue) > int(workersCount) { @@ -209,16 +210,16 @@ func (o *HTTPOutput) sendRequest(client *http.Client, data []byte) { return } - if len(o.methods) > 0 && !o.methods.Contains(request.Method) { + if len(o.config.methods) > 0 && !o.config.methods.Contains(request.Method) { return } - if !(o.urlRegexp.Good(request) && o.headerFilters.Good(request) && o.headerHashFilters.Good(request)) { + if !(o.config.urlRegexp.Good(request) && o.config.headerFilters.Good(request) && o.config.headerHashFilters.Good(request)) { return } // Rewrite the path as necessary - request.URL.Path = o.outputHTTPUrlRewrite.Rewrite(request.URL.Path) + request.URL.Path = o.config.urlRewrite.Rewrite(request.URL.Path) // Change HOST of original request URL := o.address + request.URL.Path + "?" + request.URL.RawQuery @@ -226,7 +227,7 @@ func (o *HTTPOutput) sendRequest(client *http.Client, data []byte) { request.RequestURI = "" request.URL, _ = url.ParseRequestURI(URL) - for _, header := range o.headers { + for _, header := range o.config.headers { SetHeader(request, header.Name, header.Value) } diff --git a/output_http_test.go b/output_http_test.go index af13c1a..82c8877 100644 --- a/output_http_test.go +++ b/output_http_test.go @@ -75,7 +75,7 @@ func TestHTTPOutput(t *testing.T) { headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}} methods := HTTPMethods{"GET", "PUT", "POST"} - output := NewHTTPOutput(listener.Addr().String(), headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{}, 0) + output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{headers: headers, methods: methods}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} @@ -104,11 +104,7 @@ func TestOutputHTTPSSL(t *testing.T) { })) input := NewTestInput() - - headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}} - methods := HTTPMethods{"GET", "PUT", "POST"} - - http_output := NewHTTPOutput(server.URL, headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{}, 0) + http_output := NewHTTPOutput(server.URL, &HTTPOutputConfig{}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{http_output} @@ -128,17 +124,13 @@ func BenchmarkHTTPOutput(b *testing.B) { 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) { time.Sleep(50 * time.Millisecond) wg.Done() }) - output := NewHTTPOutput(listener.Addr().String(), headers, methods, HTTPUrlRegexp{}, HTTPHeaderFilters{}, HTTPHeaderHashFilters{}, "", UrlRewriteMap{}, 0) + input := NewTestInput() + output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} diff --git a/plugins.go b/plugins.go index 20234a0..65d7937 100644 --- a/plugins.go +++ b/plugins.go @@ -96,6 +96,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, Settings.outputHTTPRedirects) + registerPlugin(NewHTTPOutput, options, &HTTPOutputSettings) } } diff --git a/settings.go b/settings.go index c5a8452..6d917e9 100644 --- a/settings.go +++ b/settings.go @@ -44,6 +44,7 @@ type AppSettings struct { } var Settings AppSettings = AppSettings{} +var HTTPOutputSettings = HTTPOutputConfig{} func usage() { fmt.Printf("Gor is a simple http traffic replication tool written in Go. Its main goal is to replay traffic from production servers to staging and dev environments.\nProject page: https://github.com/buger/gor\nAuthor: leonsbox@gmail.com\nCurrent Version: %s\n\n", VERSION) @@ -74,21 +75,22 @@ 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.Var(&Settings.outputHTTPHeaders, "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.outputHTTPMethods, "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.outputHTTPUrlRegexp, "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.outputHTTPHeaderFilters, "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.outputHTTPHeaderHashFilters, "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.IntVar(&Settings.outputHTTPWorkers, "output-http-workers", -1, "Gor uses dynamic worker scaling by default. Enter a number to run a set number of workers.") - flag.BoolVar(&Settings.outputHTTPStats, "output-http-stats", false, "Report http output queue stats to console every 5 seconds.") + flag.Var(&HTTPOutputSettings.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(&HTTPOutputSettings.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(&HTTPOutputSettings.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(&HTTPOutputSettings.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(&HTTPOutputSettings.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.IntVar(&HTTPOutputSettings.workers, "output-http-workers", 0, "Gor uses dynamic worker scaling by default. Enter a number to run a set number of workers.") + flag.BoolVar(&HTTPOutputSettings.stats, "output-http-stats", false, "Report http output queue stats to console every 5 seconds.") - 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.") + flag.StringVar(&HTTPOutputSettings.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'") + flag.Var(&HTTPOutputSettings.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.IntVar(&HTTPOutputSettings.redirectLimit, "output-http-redirects", 0, "Enable how often redirects should be followed.") } func Debug(args ...interface{}) { if Settings.verbose { + log.Print("[DEBUG] ") log.Println(args...) } }