diff --git a/elasticsearch.go b/elasticsearch.go index 8bfaa5d..af57db1 100644 --- a/elasticsearch.go +++ b/elasticsearch.go @@ -124,6 +124,7 @@ func (p *ESPlugin) RttDurationToMs(d time.Duration) int64 { return int64(fl) } +// ResponseAnalyze send req and resp to ES func (p *ESPlugin) ResponseAnalyze(req, resp []byte, start, stop time.Time) { if len(resp) == 0 { // nil http response - skipped elasticsearch export for this request @@ -131,7 +132,6 @@ func (p *ESPlugin) ResponseAnalyze(req, resp []byte, start, stop time.Time) { } t := time.Now() rtt := p.RttDurationToMs(stop.Sub(start)) - req = payloadBody(req) esResp := ESRequestResponse{ ReqURL: string(proto.Path(req)), diff --git a/emitter.go b/emitter.go index c58ff2d..644b784 100644 --- a/emitter.go +++ b/emitter.go @@ -1,30 +1,30 @@ package main import ( - "bytes" + "fmt" "hash/fnv" "io" "log" "sync" "time" + + "github.com/buger/goreplay/byteutils" ) -type emitter struct { +// Emitter represents an abject to manage plugins communication +type Emitter struct { sync.Mutex sync.WaitGroup - quit chan int plugins *InOutPlugins } -// NewEmitter creates and initializes new `emitter` object. -func NewEmitter(quit chan int) *emitter { - return &emitter{ - quit: quit, - } +// NewEmitter creates and initializes new Emitter object. +func NewEmitter() *Emitter { + return &Emitter{} } // Start initialize loop for sending data from inputs to outputs -func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) { +func (e *Emitter) Start(plugins *InOutPlugins, middlewareCmd string) { defer e.Wait() if Settings.CopyBufferSize < 1 { Settings.CopyBufferSize = 5 << 20 @@ -38,147 +38,92 @@ func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) { middleware.ReadFrom(in) } - // We are going only to read responses, so using same ReadFrom method - for _, out := range plugins.Outputs { - if r, ok := out.(io.Reader); ok { - middleware.ReadFrom(r) - } - } + e.plugins.Inputs = append(e.plugins.Inputs, middleware) + e.plugins.All = append(e.plugins.All, middleware) e.Add(1) go func() { defer e.Done() - if err := CopyMulty(e.quit, middleware, plugins.Outputs...); err != nil { - Debug(2, "Error during copy: ", err) - e.Close() - } - }() - go func() { - for { - select { - case <-e.quit: - middleware.Close() - return - } + if err := CopyMulty(middleware, plugins.Outputs...); err != nil { + Debug(2, fmt.Sprintf("[EMITTER] error during copy: %q", err)) } }() } else { for _, in := range plugins.Inputs { e.Add(1) - go func(in io.Reader) { + go func(in PluginReader) { defer e.Done() - if err := CopyMulty(e.quit, in, plugins.Outputs...); err != nil { - Debug(2, "Error during copy: ", err) - e.Close() + if err := CopyMulty(in, plugins.Outputs...); err != nil { + Debug(2, fmt.Sprintf("[EMITTER] error during copy: %q", err)) } }(in) } - - for _, out := range plugins.Outputs { - if r, ok := out.(io.Reader); ok { - e.Add(1) - go func(r io.Reader) { - defer e.Done() - if err := CopyMulty(e.quit, r, plugins.Outputs...); err != nil { - Debug(2, "Error during copy: ", err) - e.Close() - } - }(r) - } - } - } -} - -func (e *emitter) close() { - select { - case <-e.quit: - default: - close(e.quit) } } // Close closes all the goroutine and waits for it to finish. -func (e *emitter) Close() { - e.close() - for _, p := range e.plugins.Inputs { +func (e *Emitter) Close() { + for _, p := range e.plugins.All { if cp, ok := p.(io.Closer); ok { cp.Close() } } - for _, p := range e.plugins.Outputs { - if cp, ok := p.(io.Closer); ok { - cp.Close() - } - } - e.plugins = nil // avoid further accidental usage + e.plugins.All = nil // avoid further accidental close } // CopyMulty copies from 1 reader to multiple writers -func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error { - buf := make([]byte, Settings.CopyBufferSize) +func CopyMulty(src PluginReader, writers ...PluginWriter) error { wIndex := 0 modifier := NewHTTPModifier(&Settings.ModifierConfig) - filteredRequests := make(map[string]time.Time) - filteredRequestsLastCleanTime := time.Now() + filteredRequests := make(map[string]int64) + filteredRequestsLastCleanTime := time.Now().UnixNano() + filteredCount := 0 - i := 0 for { - var nr int - nr, err := src.Read(buf) - - select { - case <-stop: - return nil - default: - } + msg, err := src.PluginRead() if err != nil { + if err == ErrorStopped || err == io.EOF { + return nil + } return err } - - _maxN := nr - if nr > 500 { - _maxN = 500 - } - if nr > 0 { - payload := buf[:nr] - meta := payloadMeta(payload) + if msg != nil && len(msg.Data) > 0 { + if len(msg.Data) > int(Settings.CopyBufferSize) { + msg.Data = msg.Data[:Settings.CopyBufferSize] + } + meta := payloadMeta(msg.Meta) if len(meta) < 3 { - Debug(2, "[EMITTER] Found malformed record", string(payload[0:_maxN]), nr, "from:", src) + Debug(2, fmt.Sprintf("[EMITTER] Found malformed record %q from %q", msg.Meta, src)) continue } - requestID := string(meta[1]) - - Debug(3, "[EMITTER] input:", string(payload[0:_maxN]), nr, "from:", src) - + requestID := byteutils.SliceToString(meta[1]) + // start a subroutine only when necessary + if Settings.Verbose >= 3 { + Debug(3, "[EMITTER] input: ", byteutils.SliceToString(msg.Meta[:len(msg.Meta)-1]), " from: ", src) + } if modifier != nil { - if isRequestPayload(payload) { - headSize := bytes.IndexByte(payload, '\n') + 1 - body := payload[headSize:] - originalBodyLen := len(body) - body = modifier.Rewrite(body) - + Debug(3, "[EMITTER] modifier:", requestID, "from:", src) + if isRequestPayload(msg.Meta) { + msg.Data = modifier.Rewrite(msg.Data) // If modifier tells to skip request - if len(body) == 0 { - filteredRequests[requestID] = time.Now() + if len(msg.Data) == 0 { + filteredRequests[requestID] = time.Now().UnixNano() + filteredCount++ continue } - - if originalBodyLen != len(body) { - payload = append(payload[:headSize], body...) - } - - Debug(3, "[EMITTER] Rewritten input:", len(payload), "First %d bytes:", _maxN, string(payload[0:_maxN])) + Debug(3, "[EMITTER] Rewritten input:", requestID, "from:", src) } else { if _, ok := filteredRequests[requestID]; ok { delete(filteredRequests, requestID) + filteredCount-- continue } } } if Settings.PrettifyHTTP { - payload = prettifyHTTP(payload) - if len(payload) == 0 { + msg.Data = prettifyHTTP(msg.Data) + if len(msg.Data) == 0 { continue } } @@ -189,15 +134,15 @@ func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error { log.Fatal("Detailed TCP sessions work only with PRO license") } hasher := fnv.New32a() - // First 20 bytes contain tcp session - id := payloadID(payload) - hasher.Write(id) + hasher.Write(meta[1]) wIndex = int(hasher.Sum32()) % len(writers) - writers[wIndex].Write(payload) + if _, err := writers[wIndex].PluginWrite(msg); err != nil { + return err + } } else { // Simple round robin - if _, err := writers[wIndex].Write(payload); err != nil { + if _, err := writers[wIndex].PluginWrite(msg); err != nil { return err } @@ -209,7 +154,7 @@ func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error { } } else { for _, dst := range writers { - if _, err := dst.Write(payload); err != nil { + if _, err := dst.PluginWrite(msg); err != nil { return err } } @@ -217,19 +162,18 @@ func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error { } // Run GC on each 1000 request - if i%1000 == 0 { + if filteredCount > 0 && filteredCount%1000 == 0 { // Clean up filtered requests for which we didn't get a response to filter - now := time.Now() - if now.Sub(filteredRequestsLastCleanTime) > 60*time.Second { + now := time.Now().UnixNano() + if now-filteredRequestsLastCleanTime > int64(60*time.Second) { for k, v := range filteredRequests { - if now.Sub(v) > 60*time.Second { + if now-v > int64(60*time.Second) { delete(filteredRequests, k) + filteredCount-- } } - filteredRequestsLastCleanTime = time.Now() + filteredRequestsLastCleanTime = time.Now().UnixNano() } } - - i++ } } diff --git a/emitter_test.go b/emitter_test.go index 0b4f640..805ce45 100644 --- a/emitter_test.go +++ b/emitter_test.go @@ -2,7 +2,6 @@ package main import ( "fmt" - "io" "os" "sync" "sync/atomic" @@ -18,20 +17,19 @@ func TestMain(m *testing.M) { func TestEmitter(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() - output := NewTestOutput(func(data []byte) { + output := NewTestOutput(func(*Message) { wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 1000; i++ { @@ -45,32 +43,31 @@ func TestEmitter(t *testing.T) { func TestEmitterFiltered(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() input.skipHeader = true - output := NewTestOutput(func(data []byte) { + output := NewTestOutput(func(*Message) { wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) methods := HTTPMethods{[]byte("GET")} Settings.ModifierConfig = HTTPModifierConfig{Methods: methods} - emitter := &emitter{quit: quit} + emitter := &Emitter{} go emitter.Start(plugins, "") wg.Add(2) id := uuid() reqh := payloadHeader(RequestPayload, id, time.Now().UnixNano(), -1) - reqb := append(reqh, []byte("GET / HTTP/1.1\r\nHost: www.w3.org\r\nUser-Agent: Go 1.1 package http\r\nAccept-Encoding: gzip\r\n\r\n")...) + reqb := append(reqh, []byte("POST / HTTP/1.1\r\nHost: www.w3.org\r\nUser-Agent: Go 1.1 package http\r\nAccept-Encoding: gzip\r\n\r\n")...) resh := payloadHeader(ResponsePayload, id, time.Now().UnixNano()+1, 1) respb := append(resh, []byte("HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")...) @@ -80,7 +77,7 @@ func TestEmitterFiltered(t *testing.T) { id = uuid() reqh = payloadHeader(RequestPayload, id, time.Now().UnixNano(), -1) - reqb = append(reqh, []byte("POST / HTTP/1.1\r\nHost: www.w3.org\r\nUser-Agent: Go 1.1 package http\r\nAccept-Encoding: gzip\r\n\r\n")...) + reqb = append(reqh, []byte("GET / HTTP/1.1\r\nHost: www.w3.org\r\nUser-Agent: Go 1.1 package http\r\nAccept-Encoding: gzip\r\n\r\n")...) resh = payloadHeader(ResponsePayload, id, time.Now().UnixNano()+1, 1) respb = append(resh, []byte("HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")...) @@ -96,30 +93,29 @@ func TestEmitterFiltered(t *testing.T) { func TestEmitterSplitRoundRobin(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() var counter1, counter2 int32 - output1 := NewTestOutput(func(data []byte) { + output1 := NewTestOutput(func(*Message) { atomic.AddInt32(&counter1, 1) wg.Done() }) - output2 := NewTestOutput(func(data []byte) { + output2 := NewTestOutput(func(*Message) { atomic.AddInt32(&counter2, 1) wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output1, output2}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output1, output2}, } Settings.SplitOutput = true - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 1000; i++ { @@ -140,31 +136,30 @@ func TestEmitterSplitRoundRobin(t *testing.T) { func TestEmitterRoundRobin(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() var counter1, counter2 int32 - output1 := NewTestOutput(func(data []byte) { - atomic.AddInt32(&counter1, 1) + output1 := NewTestOutput(func(*Message) { + counter1++ wg.Done() }) - output2 := NewTestOutput(func(data []byte) { - atomic.AddInt32(&counter2, 1) + output2 := NewTestOutput(func(*Message) { + counter2++ wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output1, output2}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output1, output2}, } plugins.All = append(plugins.All, input, output1, output2) Settings.SplitOutput = true - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 1000; i++ { @@ -186,36 +181,34 @@ func TestEmitterSplitSession(t *testing.T) { wg := new(sync.WaitGroup) wg.Add(200) - quit := make(chan int) - input := NewTestInput() input.skipHeader = true var counter1, counter2 int32 - output1 := NewTestOutput(func(data []byte) { - if payloadID(data)[0] == 'a' { - atomic.AddInt32(&counter1, 1) + output1 := NewTestOutput(func(msg *Message) { + if payloadID(msg.Meta)[0] == 'a' { + counter1++ } wg.Done() }) - output2 := NewTestOutput(func(data []byte) { - if payloadID(data)[0] == 'b' { - atomic.AddInt32(&counter2, 1) + output2 := NewTestOutput(func(msg *Message) { + if payloadID(msg.Meta)[0] == 'b' { + counter2++ } wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output1, output2}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output1, output2}, } Settings.SplitOutput = true Settings.RecognizeTCPSessions = true - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 200; i++ { @@ -242,21 +235,20 @@ func TestEmitterSplitSession(t *testing.T) { func BenchmarkEmitter(b *testing.B) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() - output := NewTestOutput(func(data []byte) { + output := NewTestOutput(func(*Message) { wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) b.ResetTimer() diff --git a/gor.go b/gor.go index 6272b06..7827bc0 100644 --- a/gor.go +++ b/gor.go @@ -74,7 +74,7 @@ func main() { } closeCh := make(chan int) - emitter := NewEmitter(closeCh) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) if Settings.ExitAfter > 0 { log.Printf("Running gor for a duration of %s\n", Settings.ExitAfter) diff --git a/http_modifier_settings.go b/http_modifier_settings.go index e4263cf..7edb50b 100644 --- a/http_modifier_settings.go +++ b/http_modifier_settings.go @@ -19,10 +19,9 @@ type HTTPModifierConfig struct { HeaderBasicAuthFilters HTTPHeaderBasicAuthFilters `json:"http-basic-auth-filter"` HeaderHashFilters HTTPHashFilters `json:"http-header-limiter"` ParamHashFilters HTTPHashFilters `json:"http-param-limiter"` - - Params HTTPParams `json:"http-set-param"` - Headers HTTPHeaders `json:"http-set-header"` - Methods HTTPMethods `json:"http-allow-method"` + Params HTTPParams `json:"http-set-param"` + Headers HTTPHeaders `json:"http-set-header"` + Methods HTTPMethods `json:"http-allow-method"` } // @@ -40,10 +39,11 @@ func (h *HTTPHeaderFilters) String() string { return fmt.Sprint(*h) } +// Set method to implement flags.Value func (h *HTTPHeaderFilters) Set(value string) error { valArr := strings.SplitN(value, ":", 2) if len(valArr) < 2 { - return errors.New("need both header and value, colon-delimited (ex. user_id:^169$).") + return errors.New("need both header and value, colon-delimited (ex. user_id:^169$)") } val := strings.TrimSpace(valArr[1]) r, err := regexp.Compile(val) @@ -63,13 +63,14 @@ type basicAuthFilter struct { regexp *regexp.Regexp } -// HTTPHeaderFilters holds list of headers and their regexps +// HTTPHeaderBasicAuthFilters holds list of regxp to match basic Auth header values type HTTPHeaderBasicAuthFilters []basicAuthFilter func (h *HTTPHeaderBasicAuthFilters) String() string { return fmt.Sprint(*h) } +// Set method to implement flags.Value func (h *HTTPHeaderBasicAuthFilters) Set(value string) error { r, err := regexp.Compile(value) if err != nil { @@ -89,12 +90,14 @@ type hashFilter struct { percent uint32 } +// HTTPHashFilters represents a slice of header hash filters type HTTPHashFilters []hashFilter func (h *HTTPHashFilters) String() string { return fmt.Sprint(*h) } +// Set method to implement flags.Value func (h *HTTPHashFilters) Set(value string) error { valArr := strings.SplitN(value, ":", 2) if len(valArr) < 2 { @@ -129,23 +132,26 @@ func (h *HTTPHashFilters) Set(value string) error { // // Handling of --http-set-header option // -type HTTPHeaders []HTTPHeader -type HTTPHeader struct { +type httpHeader struct { Name string Value string } +// HTTPHeaders is a slice of headers that must appended +type HTTPHeaders []httpHeader + func (h *HTTPHeaders) String() string { return fmt.Sprint(*h) } +// Set method to implement flags.Value func (h *HTTPHeaders) Set(value string) error { v := strings.SplitN(value, ":", 2) if len(v) != 2 { return errors.New("Expected `Key: Value`") } - header := HTTPHeader{ + header := httpHeader{ strings.TrimSpace(v[0]), strings.TrimSpace(v[1]), } @@ -157,23 +163,26 @@ func (h *HTTPHeaders) Set(value string) error { // // Handling of --http-set-param option // -type HTTPParams []HTTPParam -type HTTPParam struct { +type httpParam struct { Name []byte Value []byte } +// HTTPParams filters for --http-set-param +type HTTPParams []httpParam + func (h *HTTPParams) String() string { return fmt.Sprint(*h) } +// Set method to implement flags.Value func (h *HTTPParams) Set(value string) error { v := strings.SplitN(value, "=", 2) if len(v) != 2 { return errors.New("Expected `Key=Value`") } - param := HTTPParam{ + param := httpParam{ []byte(strings.TrimSpace(v[0])), []byte(strings.TrimSpace(v[1])), } @@ -185,12 +194,15 @@ func (h *HTTPParams) Set(value string) error { // // Handling of --http-allow-method option // + +// HTTPMethods holds values for method allowed type HTTPMethods [][]byte func (h *HTTPMethods) String() string { return fmt.Sprint(*h) } +// Set method to implement flags.Value func (h *HTTPMethods) Set(value string) error { *h = append(*h, []byte(value)) return nil @@ -204,12 +216,14 @@ type urlRewrite struct { target []byte } +// URLRewriteMap holds regexp and data to modify URL type URLRewriteMap []urlRewrite func (r *URLRewriteMap) String() string { return fmt.Sprint(*r) } +// Set method to implement flags.Value func (r *URLRewriteMap) Set(value string) error { valArr := strings.SplitN(value, ":", 2) if len(valArr) < 2 { @@ -232,12 +246,14 @@ type headerRewrite struct { target []byte } +// HeaderRewriteMap holds regexp and data to rewrite headers type HeaderRewriteMap []headerRewrite func (r *HeaderRewriteMap) String() string { return fmt.Sprint(*r) } +// Set method to implement flags.Value func (r *HeaderRewriteMap) Set(value string) error { headerArr := strings.SplitN(value, ":", 2) if len(headerArr) < 2 { @@ -266,12 +282,14 @@ type urlRegexp struct { regexp *regexp.Regexp } +// HTTPURLRegexp a slice of regexp to match URLs type HTTPURLRegexp []urlRegexp func (r *HTTPURLRegexp) String() string { return fmt.Sprint(*r) } +// Set method to implement flags.Value func (r *HTTPURLRegexp) Set(value string) error { regexp, err := regexp.Compile(value) diff --git a/http_prettifier.go b/http_prettifier.go index 7a9169f..0aefd37 100644 --- a/http_prettifier.go +++ b/http_prettifier.go @@ -12,25 +12,22 @@ import ( ) func prettifyHTTP(p []byte) []byte { - headSize := bytes.IndexByte(p, '\n') + 1 - head := p[:headSize] - body := p[headSize:] - tEnc := bytes.Equal(proto.Header(body, []byte("Transfer-Encoding")), []byte("chunked")) - cEnc := bytes.Equal(proto.Header(body, []byte("Content-Encoding")), []byte("gzip")) + tEnc := bytes.Equal(proto.Header(p, []byte("Transfer-Encoding")), []byte("chunked")) + cEnc := bytes.Equal(proto.Header(p, []byte("Content-Encoding")), []byte("gzip")) if !(tEnc || cEnc) { return p } - headersPos := proto.MIMEHeadersEndPos(body) + headersPos := proto.MIMEHeadersEndPos(p) - if headersPos < 5 || headersPos > len(body) { + if headersPos < 5 || headersPos > len(p) { return p } - headers := body[:headersPos] - content := body[headersPos:] + headers := p[:headersPos] + content := p[headersPos:] if tEnc { buf := bytes.NewReader(content) @@ -64,7 +61,7 @@ func prettifyHTTP(p []byte) []byte { headers = proto.SetHeader(headers, []byte("Content-Length"), []byte(newLen)) } - newPayload := append(append(head, headers...), content...) + newPayload := append(headers, content...) return newPayload } diff --git a/http_prettifier_test.go b/http_prettifier_test.go index c48c97f..7763a70 100644 --- a/http_prettifier_test.go +++ b/http_prettifier_test.go @@ -5,6 +5,8 @@ import ( "compress/gzip" "strconv" "testing" + + "github.com/buger/goreplay/proto" ) func TestHTTPPrettifierGzip(t *testing.T) { @@ -15,22 +17,21 @@ func TestHTTPPrettifierGzip(t *testing.T) { size := strconv.Itoa(len(b.Bytes())) - payload := []byte("2 1 1\nHTTP/1.1 200 OK\r\nContent-Length: " + size + "\r\nContent-Encoding: gzip\r\n\r\n") + payload := []byte("HTTP/1.1 200 OK\r\nContent-Length: " + size + "\r\nContent-Encoding: gzip\r\n\r\n") payload = append(payload, b.Bytes()...) newPayload := prettifyHTTP(payload) - if string(newPayload) != "2 1 1\nHTTP/1.1 200 OK\r\nContent-Length: 4\r\n\r\ntest" { - t.Error("Payload not match:", string(newPayload)) + if string(newPayload) != "HTTP/1.1 200 OK\r\nContent-Length: 4\r\n\r\ntest" { + t.Errorf("Payload not match %q", string(newPayload)) } } func TestHTTPPrettifierChunked(t *testing.T) { payload := []byte("POST / HTTP/1.1\r\nHost: www.w3.org\r\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") - newPayload := prettifyHTTP(payload) - - if string(newPayload) != "POST / HTTP/1.1\r\nHost: www.w3.org\r\nContent-Length: 23\r\n\r\nWikipedia in\r\n\r\nchunks." { - t.Error("Payload not match:", string(newPayload)) + payload = prettifyHTTP(payload) + if string(proto.Header(payload, []byte("Content-Length"))) != "23" { + t.Errorf("payload should have content length of 23") } } diff --git a/input_dummy.go b/input_dummy.go index beb1f9c..5360d31 100644 --- a/input_dummy.go +++ b/input_dummy.go @@ -7,24 +7,30 @@ import ( // DummyInput used for debugging. It generate 1 "GET /"" request per second. type DummyInput struct { data chan []byte + quit chan struct{} } // NewDummyInput constructor for DummyInput func NewDummyInput(options string) (di *DummyInput) { di = new(DummyInput) di.data = make(chan []byte) + di.quit = make(chan struct{}) go di.emit() return } -func (i *DummyInput) Read(data []byte) (int, error) { - buf := <-i.data - - copy(data, buf) - - return len(buf), nil +// PluginRead reads message from this plugin +func (i *DummyInput) PluginRead() (*Message, error) { + var msg Message + select { + case <-i.quit: + return nil, ErrorStopped + case buf := <-i.data: + msg.Meta, msg.Data = payloadMetaWithBody(buf) + return &msg, nil + } } func (i *DummyInput) emit() { @@ -46,3 +52,9 @@ func (i *DummyInput) emit() { func (i *DummyInput) String() string { return "Dummy Input" } + +// Close closes this plugins +func (i *DummyInput) Close() error { + close(i.quit) + return nil +} diff --git a/input_file.go b/input_file.go index 528578b..36b72f1 100644 --- a/input_file.go +++ b/input_file.go @@ -5,8 +5,8 @@ import ( "bytes" "compress/gzip" "errors" + "fmt" "io" - "log" "os" "path/filepath" "strconv" @@ -64,6 +64,8 @@ func (f *fileInputReader) ReadPayload() []byte { return f.data } + +// Close closes this plugin func (f *fileInputReader) Close() error { if atomic.LoadInt32(&f.closed) == 0 { atomic.StoreInt32(&f.closed, 1) @@ -73,7 +75,7 @@ func (f *fileInputReader) Close() error { return nil } -func NewFileInputReader(path string) *fileInputReader { +func newFileInputReader(path string) *fileInputReader { var file io.ReadCloser var err error @@ -84,7 +86,7 @@ func NewFileInputReader(path string) *fileInputReader { } if err != nil { - log.Println(err) + Debug(0, fmt.Sprintf("[INPUT-FILE] err: %q", err)) return nil } @@ -92,7 +94,7 @@ func NewFileInputReader(path string) *fileInputReader { if strings.HasSuffix(path, ".gz") { gzReader, err := gzip.NewReader(file) if err != nil { - log.Println(err) + Debug(0, fmt.Sprintf("[INPUT-FILE] err: %q", err)) return nil } r.reader = bufio.NewReader(gzReader) @@ -153,7 +155,7 @@ func (i *FileInput) init() (err error) { resp, err := svc.ListObjects(params) if err != nil { - log.Println("Error while retreiving list of files from S3", i.path, err) + Debug(0, "[INPUT-FILE] Error while retreiving list of files from S3", i.path, err) return err } @@ -162,34 +164,35 @@ func (i *FileInput) init() (err error) { } } else { if matches, err = filepath.Glob(i.path); err != nil { - log.Println("Wrong file pattern", i.path, err) + Debug(0, "[INPUT-FILE] Wrong file pattern", i.path, err) return } } if len(matches) == 0 { - log.Println("No files match pattern: ", i.path) + Debug(0, "[INPUT-FILE] No files match pattern: ", i.path) return errors.New("No matching files") } i.readers = make([]*fileInputReader, len(matches)) for idx, p := range matches { - i.readers[idx] = NewFileInputReader(p) + i.readers[idx] = newFileInputReader(p) } return nil } -func (i *FileInput) Read(data []byte) (int, error) { - var buf []byte +// PluginRead reads message from this plugin +func (i *FileInput) PluginRead() (*Message, error) { + var msg Message select { case <-i.exit: - return 0, ErrorStopped - case buf = <-i.data: + return nil, ErrorStopped + case buf := <-i.data: + msg.Meta, msg.Data = payloadMetaWithBody(buf) + return &msg, nil } - n := copy(data, buf) - return n, nil } func (i *FileInput) String() string { @@ -256,7 +259,7 @@ func (i *FileInput) emit() { } } - log.Printf("FileInput: end of file '%s'\n", i.path) + Debug(0, fmt.Sprintf("[INPUT-FILE] FileInput: end of file '%s'\n", i.path)) } diff --git a/input_file_test.go b/input_file_test.go index 09f569e..7f92eff 100644 --- a/input_file_test.go +++ b/input_file_test.go @@ -4,9 +4,7 @@ import ( "bytes" "errors" "fmt" - "io" "io/ioutil" - "log" "math/rand" "os" "sync" @@ -14,20 +12,18 @@ import ( "time" ) -var _ = log.Println - func TestInputFileWithGET(t *testing.T) { input := NewTestInput() - rg := NewRequestGenerator([]io.Reader{input}, func() { input.EmitGET() }, 1) - readPayloads := [][]byte{} + rg := NewRequestGenerator([]PluginReader{input}, func() { input.EmitGET() }, 1) + readPayloads := []*Message{} // Given a capture file with a GET request expectedCaptureFile := CreateCaptureFile(rg) defer expectedCaptureFile.TearDown() // When the request is read from the capture file - err := ReadFromCaptureFile(expectedCaptureFile.file, 1, func(data []byte) { - readPayloads = append(readPayloads, Duplicate(data)) + err := ReadFromCaptureFile(expectedCaptureFile.file, 1, func(msg *Message) { + readPayloads = append(readPayloads, msg) }) // The read request should match the original request @@ -41,18 +37,17 @@ func TestInputFileWithGET(t *testing.T) { } func TestInputFileWithPayloadLargerThan64Kb(t *testing.T) { - input := NewTestInput() - rg := NewRequestGenerator([]io.Reader{input}, func() { input.EmitSizedPOST(64 * 1024) }, 1) - readPayloads := [][]byte{} + rg := NewRequestGenerator([]PluginReader{input}, func() { input.EmitSizedPOST(64 * 1024) }, 1) + readPayloads := []*Message{} // Given a capture file with a request over 64Kb expectedCaptureFile := CreateCaptureFile(rg) defer expectedCaptureFile.TearDown() // When the request is read from the capture file - err := ReadFromCaptureFile(expectedCaptureFile.file, 1, func(data []byte) { - readPayloads = append(readPayloads, Duplicate(data)) + err := ReadFromCaptureFile(expectedCaptureFile.file, 1, func(msg *Message) { + readPayloads = append(readPayloads, msg) }) // The read request should match the original request @@ -69,19 +64,19 @@ func TestInputFileWithPayloadLargerThan64Kb(t *testing.T) { func TestInputFileWithGETAndPOST(t *testing.T) { input := NewTestInput() - rg := NewRequestGenerator([]io.Reader{input}, func() { + rg := NewRequestGenerator([]PluginReader{input}, func() { input.EmitGET() input.EmitPOST() }, 2) - readPayloads := [][]byte{} + readPayloads := []*Message{} // Given a capture file with a GET request expectedCaptureFile := CreateCaptureFile(rg) defer expectedCaptureFile.TearDown() // When the requests are read from the capture file - err := ReadFromCaptureFile(expectedCaptureFile.file, 2, func(data []byte) { - readPayloads = append(readPayloads, Duplicate(data)) + err := ReadFromCaptureFile(expectedCaptureFile.file, 2, func(msg *Message) { + readPayloads = append(readPayloads, msg) }) // The read requests should match the original request @@ -113,12 +108,11 @@ func TestInputFileMultipleFilesWithRequestsOnly(t *testing.T) { file2.Close() input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false) - buf := make([]byte, 1000) for i := '1'; i <= '4'; i++ { - n, _ := input.Read(buf) - if buf[4] != byte(i) { - t.Error("Should emit requests in right order", string(buf[:n])) + msg, _ := input.PluginRead() + if msg.Meta[4] != byte(i) { + t.Error("Should emit requests in right order", string(msg.Meta)) } } @@ -140,11 +134,10 @@ func TestInputFileRequestsWithLatency(t *testing.T) { file.Write([]byte(payloadSeparator)) input := NewFileInput(fmt.Sprintf("/tmp/%d", rnd), false) - buf := make([]byte, 1000) start := time.Now().UnixNano() for i := 0; i < 3; i++ { - input.Read(buf) + input.PluginRead() } end := time.Now().UnixNano() @@ -181,17 +174,16 @@ func TestInputFileMultipleFilesWithRequestsAndResponses(t *testing.T) { file2.Close() input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false) - buf := make([]byte, 1000) for i := '1'; i <= '4'; i++ { - n, _ := input.Read(buf) - if buf[0] != '1' && buf[4] != byte(i) { - t.Error("Shound emit requests in right order", string(buf[:n])) + msg, _ := input.PluginRead() + if msg.Meta[0] != '1' && msg.Meta[4] != byte(i) { + t.Error("Shound emit requests in right order", string(msg.Meta)) } - n, _ = input.Read(buf) - if buf[0] != '2' && buf[4] != byte(i) { - t.Error("Shound emit responses in right order", string(buf[:n])) + msg, _ = input.PluginRead() + if msg.Meta[0] != '2' && msg.Meta[4] != byte(i) { + t.Error("Shound emit responses in right order", string(msg.Meta)) } } @@ -210,11 +202,10 @@ func TestInputFileLoop(t *testing.T) { file.Close() input := NewFileInput(fmt.Sprintf("/tmp/%d", rnd), true) - buf := make([]byte, 1000) // Even if we have just 2 requests in file, it should indifinitly loop for i := 0; i < 1000; i++ { - input.Read(buf) + input.PluginRead() } input.Close() @@ -226,22 +217,21 @@ func TestInputFileCompressed(t *testing.T) { output := NewFileOutput(fmt.Sprintf("/tmp/%d_0.gz", rnd), &FileOutputConfig{FlushInterval: time.Minute, Append: true}) for i := 0; i < 1000; i++ { - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) } name1 := output.file.Name() output.Close() output2 := NewFileOutput(fmt.Sprintf("/tmp/%d_1.gz", rnd), &FileOutputConfig{FlushInterval: time.Minute, Append: true}) for i := 0; i < 1000; i++ { - output2.Write([]byte("1 1 1\r\ntest")) + output2.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) } name2 := output2.file.Name() output2.Close() input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false) - buf := make([]byte, 1000) for i := 0; i < 2000; i++ { - input.Read(buf) + input.PluginRead() } os.Remove(name1) @@ -249,14 +239,14 @@ func TestInputFileCompressed(t *testing.T) { } type CaptureFile struct { - data [][]byte + msgs []*Message file *os.File } -func NewExpectedCaptureFile(data [][]byte, file *os.File) *CaptureFile { +func NewExpectedCaptureFile(msgs []*Message, file *os.File) *CaptureFile { ecf := new(CaptureFile) ecf.file = file - ecf.data = data + ecf.msgs = msgs return ecf } @@ -267,12 +257,12 @@ func (expectedCaptureFile *CaptureFile) TearDown() { } type RequestGenerator struct { - inputs []io.Reader + inputs []PluginReader emit func() wg *sync.WaitGroup } -func NewRequestGenerator(inputs []io.Reader, emit func(), count int) (rg *RequestGenerator) { +func NewRequestGenerator(inputs []PluginReader, emit func(), count int) (rg *RequestGenerator) { rg = new(RequestGenerator) rg.inputs = inputs rg.emit = emit @@ -281,14 +271,17 @@ func NewRequestGenerator(inputs []io.Reader, emit func(), count int) (rg *Reques return } -func (expectedCaptureFile *CaptureFile) PayloadsEqual(other [][]byte) bool { +func (expectedCaptureFile *CaptureFile) PayloadsEqual(other []*Message) bool { - if len(expectedCaptureFile.data) != len(other) { + if len(expectedCaptureFile.msgs) != len(other) { return false } for i, payload := range other { - if !bytes.Equal(expectedCaptureFile.data[i], payload) { + if !bytes.Equal(expectedCaptureFile.msgs[i].Meta, payload.Meta) { + return false + } + if !bytes.Equal(expectedCaptureFile.msgs[i].Data, payload.Data) { return false } } @@ -303,26 +296,24 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile { panic(err) } - quit := make(chan int) - - readPayloads := [][]byte{} - output := NewTestOutput(func(data []byte) { - readPayloads = append(readPayloads, Duplicate(data)) + readPayloads := []*Message{} + output := NewTestOutput(func(msg *Message) { + readPayloads = append(readPayloads, msg) requestGenerator.wg.Done() }) - outputFile := NewFileOutput(f.Name(), &FileOutputConfig{FlushInterval: time.Minute, Append: true}) + outputFile := NewFileOutput(f.Name(), &FileOutputConfig{FlushInterval: time.Second, Append: true}) plugins := &InOutPlugins{ Inputs: requestGenerator.inputs, - Outputs: []io.Writer{output, outputFile}, + Outputs: []PluginWriter{output, outputFile}, } for _, input := range requestGenerator.inputs { plugins.All = append(plugins.All, input) } plugins.All = append(plugins.All, output, outputFile) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) requestGenerator.emit() @@ -336,23 +327,22 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile { } func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback) (err error) { - quit := make(chan int) wg := new(sync.WaitGroup) input := NewFileInput(captureFile.Name(), false) - output := NewTestOutput(func(data []byte) { - callback(data) + output := NewTestOutput(func(msg *Message) { + callback(msg) wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) wg.Add(count) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) done := make(chan int, 1) @@ -371,10 +361,3 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback return } - -func Duplicate(data []byte) (duplicate []byte) { - duplicate = make([]byte, len(data)) - copy(duplicate, data) - - return -} diff --git a/input_http.go b/input_http.go index 91ff9cb..46583b4 100644 --- a/input_http.go +++ b/input_http.go @@ -27,25 +27,17 @@ func NewHTTPInput(address string) (i *HTTPInput) { return } -func (i *HTTPInput) Read(data []byte) (int, error) { - var buf []byte +// PluginRead reads message from this plugin +func (i *HTTPInput) PluginRead() (*Message, error) { + var msg Message select { case <-i.stop: - return 0, ErrorStopped - case buf = <-i.data: + return nil, ErrorStopped + case buf := <-i.data: + msg.Data = buf + msg.Meta = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1) + return &msg, nil } - header := payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1) - - n := copy(data, header) - if len(data) > len(header) { - n += copy(data[len(header):], buf) - } - dis := len(header) + len(buf) - n - if dis > 0 { - Debug(2, "[INPUT-HTTP] discarded", dis, "increase copy buffer size") - } - - return n, nil } // Close closes this plugin diff --git a/input_http_test.go b/input_http_test.go index 8ed3a38..afce748 100644 --- a/input_http_test.go +++ b/input_http_test.go @@ -2,33 +2,29 @@ package main import ( "bytes" - "io" "net/http" "strings" "sync" "testing" "time" - - "github.com/buger/goreplay/proto" ) func TestHTTPInput(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewHTTPInput("127.0.0.1:0") time.Sleep(time.Millisecond) - output := NewTestOutput(func(data []byte) { + output := NewTestOutput(func(*Message) { wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) address := strings.Replace(input.address, "[::]", "127.0.0.1", -1) @@ -44,27 +40,25 @@ func TestHTTPInput(t *testing.T) { func TestInputHTTPLargePayload(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int, 1) const n = 10 << 20 // 10MB var large [n]byte large[n-1] = '0' input := NewHTTPInput("127.0.0.1:0") - time.Sleep(time.Millisecond) - output := NewTestOutput(func(data []byte) { - _len := len(proto.Body(payloadBody(data))) + output := NewTestOutput(func(msg *Message) { + _len := len(msg.Data) if _len >= n { // considering http body CRLF t.Errorf("expected body to be >= %d", n) } wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() defer emitter.Close() go emitter.Start(plugins, Settings.Middleware) diff --git a/input_kafka.go b/input_kafka.go index 10f426f..9f660bd 100644 --- a/input_kafka.go +++ b/input_kafka.go @@ -15,6 +15,7 @@ type KafkaInput struct { config *InputKafkaConfig consumers []sarama.PartitionConsumer messages chan *sarama.ConsumerMessage + quit chan struct{} } // NewKafkaInput creates instance of kafka consumer client with TLS config @@ -43,6 +44,7 @@ func NewKafkaInput(address string, config *InputKafkaConfig, tlsConfig *KafkaTLS config: config, consumers: make([]sarama.PartitionConsumer, len(partitions)), messages: make(chan *sarama.ConsumerMessage, 256), + quit: make(chan struct{}), } for index, partition := range partitions { @@ -74,29 +76,45 @@ func (i *KafkaInput) ErrorHandler(consumer sarama.PartitionConsumer) { } } -func (i *KafkaInput) Read(data []byte) (int, error) { - message := <-i.messages - - if !i.config.UseJSON { - copy(data, message.Value) - return len(message.Value), nil +// PluginRead a reads message from this plugin +func (i *KafkaInput) PluginRead() (*Message, error) { + var message *sarama.ConsumerMessage + var msg Message + select { + case <-i.quit: + return nil, ErrorStopped + case message = <-i.messages: } - var kafkaMessage KafkaMessage - json.Unmarshal(message.Value, &kafkaMessage) + msg.Data = message.Value + if i.config.UseJSON { - buf, err := kafkaMessage.Dump() - if err != nil { - Debug(1, "Failed to decode access log entry:", err) - return 0, err + var kafkaMessage KafkaMessage + json.Unmarshal(message.Value, &kafkaMessage) + + var err error + msg.Data, err = kafkaMessage.Dump() + if err != nil { + Debug(1, "[INPUT-KAFKA] failed to decode access log entry:", err) + return nil, err + } } - n := copy(data, buf) + // does it have meta + if isOriginPayload(msg.Data) { + msg.Meta, msg.Data = payloadMetaWithBody(msg.Data) + } - return n, nil + return &msg, nil } func (i *KafkaInput) String() string { return "Kafka Input: " + i.config.Host + "/" + i.config.Topic } + +// Close closes this plugin +func (i *KafkaInput) Close() error { + close(i.quit) + return nil +} diff --git a/input_kafka_test.go b/input_kafka_test.go index 544da6b..3ae6974 100644 --- a/input_kafka_test.go +++ b/input_kafka_test.go @@ -20,17 +20,16 @@ func TestInputKafkaRAW(t *testing.T) { consumer: consumer, Topic: "test", UseJSON: false, - },nil) + }, nil) - buf := make([]byte, 1024) - n, err := input.Read(buf) + msg, err := input.PluginRead() if err != nil { t.Fatal(err) } - if string(buf[:n]) != "1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n" { - t.Error("Message not properly decoded: ", string(buf[:n]), n) + if string(append(msg.Meta, msg.Data...)) != "1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n" { + t.Error("Message not properly decoded") } } @@ -47,16 +46,15 @@ func TestInputKafkaJSON(t *testing.T) { consumer: consumer, Topic: "test", UseJSON: true, - },nil) + }, nil) - buf := make([]byte, 1024) - n, err := input.Read(buf) + msg, err := input.PluginRead() if err != nil { t.Fatal(err) } - if string(buf[:n]) != "1 2 3\nGET / HTTP/1.1\r\nHeader: 1\r\n\r\n" { - t.Error("Message not properly decoded: ", string(buf[:n]), n) + if string(append(msg.Meta, msg.Data...)) != "1 2 3\nGET / HTTP/1.1\r\nHeader: 1\r\n\r\n" { + t.Error("Message not properly decoded") } } diff --git a/input_raw.go b/input_raw.go index ab3e7e4..cc39a0c 100644 --- a/input_raw.go +++ b/input_raw.go @@ -102,44 +102,39 @@ func NewRAWInput(address string, config RAWInputConfig) (i *RAWInput) { return } -func (i *RAWInput) Read(data []byte) (n int, err error) { - var msg *tcp.Message - var buf []byte +// PluginRead reads meassage from this plugin +func (i *RAWInput) PluginRead() (*Message, error) { + var msgTCP *tcp.Message + var msg Message select { case <-i.quit: - return 0, ErrorStopped - case msg = <-i.message: - buf = msg.Data() + return nil, ErrorStopped + case msgTCP = <-i.message: + msg.Data = msgTCP.Data() } - var header []byte - var msgType byte = ResponsePayload - if msg.IsIncoming { + if msgTCP.IsIncoming { msgType = RequestPayload if i.RealIPHeader != "" { - buf = proto.SetHeader(buf, []byte(i.RealIPHeader), []byte(msg.SrcAddr)) + msg.Data = proto.SetHeader(msg.Data, []byte(i.RealIPHeader), []byte(msgTCP.SrcAddr)) } } - header = payloadHeader(msgType, msg.UUID(), msg.Start.UnixNano(), msg.End.UnixNano()-msg.Start.UnixNano()) + msg.Meta = payloadHeader(msgType, msgTCP.UUID(), msgTCP.Start.UnixNano(), msgTCP.End.UnixNano()-msgTCP.Start.UnixNano()) - n = copy(data, header) - if len(data) > len(header) { - n += copy(data[len(header):], buf) - } // to be removed.... - if msg.Truncated || len(header)+len(buf)-n > 0 { - go Debug(2, "[INPUT-RAW] message truncated, increase copy-buffer-size") + if msgTCP.Truncated { + Debug(2, "[INPUT-RAW] message truncated, increase copy-buffer-size") } // to be removed... - if msg.TimedOut { - go Debug(2, "[INPUT-RAW] message timeout reached, increase input-raw-expire") + if msgTCP.TimedOut { + Debug(2, "[INPUT-RAW] message timeout reached, increase input-raw-expire") } if i.Stats { - stat := msg.Stats + stat := msgTCP.Stats go i.addStats(stat) } - msg = nil - return n, nil + msgTCP = nil + return &msg, nil } func (i *RAWInput) listen(address string) { diff --git a/input_raw_test.go b/input_raw_test.go index a41bf2a..9071396 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -2,7 +2,6 @@ package main import ( "bytes" - "io" "io/ioutil" "net" "net/http" @@ -23,9 +22,8 @@ const testRawExpire = time.Millisecond * 200 func TestRAWInputIPv4(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) - listener, err := net.Listen("tcp", ":0") + listener, err := net.Listen("tcp4", "127.0.0.1:0") if err != nil { t.Error(err) return @@ -49,32 +47,31 @@ func TestRAWInputIPv4(t *testing.T) { TrackResponse: true, RealIPHeader: "X-Real-IP", } - input := NewRAWInput(":"+port, conf) + input := NewRAWInput(listener.Addr().String(), conf) - output := NewTestOutput(func(data []byte) { - if data[0] == '1' { - body := payloadBody(data) - if len(proto.Header(body, []byte("X-Real-IP"))) == 0 { - t.Error("Should have X-Real-IP header", string(body)) + output := NewTestOutput(func(msg *Message) { + if msg.Meta[0] == '1' { + if len(proto.Header(msg.Data, []byte("X-Real-IP"))) == 0 { + t.Error("Should have X-Real-IP header") } - atomic.AddInt64(&reqCounter, 1) + reqCounter++ } else { - atomic.AddInt64(&respCounter, 1) + respCounter++ } wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) addr := "http://127.0.0.1:" + port - emitter := NewEmitter(quit) + emitter := NewEmitter() defer emitter.Close() go emitter.Start(plugins, Settings.Middleware) - for i := 0; i < 10; i++ { + for i := 0; i < 1; i++ { wg.Add(2) _, err = http.Get(addr) if err != nil { @@ -91,7 +88,6 @@ func TestRAWInputIPv4(t *testing.T) { func TestRAWInputNoKeepAlive(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { @@ -117,8 +113,8 @@ func TestRAWInputNoKeepAlive(t *testing.T) { } input := NewRAWInput(":"+port, conf) var respCounter, reqCounter int64 - output := NewTestOutput(func(data []byte) { - if data[0] == '1' { + output := NewTestOutput(func(msg *Message) { + if msg.Meta[0] == '1' { atomic.AddInt64(&reqCounter, 1) } else { atomic.AddInt64(&respCounter, 1) @@ -127,14 +123,14 @@ func TestRAWInputNoKeepAlive(t *testing.T) { }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) addr := "http://127.0.0.1:" + port - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 10; i++ { @@ -157,7 +153,6 @@ func TestRAWInputNoKeepAlive(t *testing.T) { func TestRAWInputIPv6(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) listener, err := net.Listen("tcp", "[::1]:0") if err != nil { @@ -183,8 +178,8 @@ func TestRAWInputIPv6(t *testing.T) { } input := NewRAWInput(originAddr, conf) - output := NewTestOutput(func(data []byte) { - if data[0] == '1' { + output := NewTestOutput(func(msg *Message) { + if msg.Meta[0] == '1' { atomic.AddInt64(&reqCounter, 1) } else { atomic.AddInt64(&respCounter, 1) @@ -193,11 +188,11 @@ func TestRAWInputIPv6(t *testing.T) { }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } - emitter := NewEmitter(quit) + emitter := NewEmitter() addr := "http://" + originAddr go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 10; i++ { @@ -220,7 +215,6 @@ func TestRAWInputIPv6(t *testing.T) { func TestInputRAWChunkedEncoding(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) fileContent, _ := ioutil.ReadFile("README.md") @@ -257,12 +251,12 @@ func TestInputRAWChunkedEncoding(t *testing.T) { httpOutput := NewHTTPOutput(replay.URL, &HTTPOutputConfig{}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{httpOutput}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{httpOutput}, } plugins.All = append(plugins.All, input, httpOutput) - emitter := NewEmitter(quit) + emitter := NewEmitter() defer emitter.Close() go emitter.Start(plugins, Settings.Middleware) wg.Add(2) @@ -278,17 +272,15 @@ func TestInputRAWChunkedEncoding(t *testing.T) { } func BenchmarkRAWInputWithReplay(b *testing.B) { - var respCounter, reqCounter, replayCounter uint64 + var respCounter, reqCounter, replayCounter uint32 wg := &sync.WaitGroup{} - wg.Add(b.N * 3) // reqCounter + replayCounter + respCounter - quit := make(chan int) - listener, err := net.Listen("tcp", "127.0.0.1:0") + listener, err := net.Listen("tcp4", "127.0.0.1:0") if err != nil { b.Error(err) return } - listener0, err := net.Listen("tcp", "127.0.0.1:0") + listener0, err := net.Listen("tcp4", "127.0.0.1:0") if err != nil { b.Error(err) return @@ -305,6 +297,8 @@ func BenchmarkRAWInputWithReplay(b *testing.B) { replay := http.Server{ Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddUint32(&replayCounter, 1) + w.Write(nil) wg.Done() }), } @@ -320,34 +314,37 @@ func BenchmarkRAWInputWithReplay(b *testing.B) { } input := NewRAWInput(originAddr, conf) - testOutput := NewTestOutput(func(data []byte) { - if data[0] == '1' { - atomic.AddUint64(&reqCounter, 1) + testOutput := NewTestOutput(func(msg *Message) { + if msg.Meta[0] == '1' { + reqCounter++ } else { - atomic.AddUint64(&respCounter, 1) + respCounter++ } wg.Done() }) - httpOutput := NewHTTPOutput(replayAddr, &HTTPOutputConfig{}) + httpOutput := NewHTTPOutput("http://"+replayAddr, &HTTPOutputConfig{}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{testOutput, httpOutput}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{testOutput, httpOutput}, } - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) - now := time.Now() addr := "http://" + originAddr + b.ResetTimer() for i := 0; i < b.N; i++ { - _, err = http.Get(addr) + wg.Add(3) // reqCounter + replayCounter + respCounter + resp, err := http.Get(addr) if err != nil { - b.Log(err) wg.Add(-3) } + resp.Body.Close() } wg.Wait() - b.Logf("%d/%d Requests, %d/%d Responses, %d/%d Replayed in %s\n", reqCounter, b.N, respCounter, b.N, replayCounter, b.N, time.Since(now)) + b.ReportMetric(float64(reqCounter), "requests") + b.ReportMetric(float64(respCounter), "responses") + b.ReportMetric(float64(replayCounter), "replayed") emitter.Close() } diff --git a/input_tcp.go b/input_tcp.go index e7ea273..257bc12 100644 --- a/input_tcp.go +++ b/input_tcp.go @@ -8,18 +8,18 @@ import ( "io" "log" "net" - "os" ) // TCPInput used for internal communication type TCPInput struct { - data chan []byte + data chan *Message listener net.Listener address string config *TCPInputConfig stop chan bool // Channel used only to indicate goroutine should shutdown } +// TCPInputConfig represents configuration of a TCP input plugin type TCPInputConfig struct { Secure bool `json:"input-tcp-secure"` CertificatePath string `json:"input-tcp-certificate"` @@ -29,7 +29,7 @@ type TCPInputConfig struct { // NewTCPInput constructor for TCPInput, accepts address with port func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) { i = new(TCPInput) - i.data = make(chan []byte, 1000) + i.data = make(chan *Message, 1000) i.address = address i.config = config i.stop = make(chan bool) @@ -39,20 +39,21 @@ func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) { return } -func (i *TCPInput) Read(data []byte) (int, error) { - var buf []byte +// PluginRead returns data and details read from plugin +func (i *TCPInput) PluginRead() (msg *Message, err error) { select { case <-i.stop: - return 0, ErrorStopped - case buf = <-i.data: + return nil, ErrorStopped + case msg = <-i.data: + return msg, nil } - copy(data, buf) - return len(buf), nil } +// Close closes the plugin func (i *TCPInput) Close() error { close(i.stop) + i.listener.Close() return nil } @@ -60,64 +61,67 @@ func (i *TCPInput) listen(address string) { if i.config.Secure { cer, err := tls.LoadX509KeyPair(i.config.CertificatePath, i.config.KeyPath) if err != nil { - log.Fatal("Error while loading --input-file certificate:", err) + log.Fatalln("error while loading --input-tcp TLS certificate:", err) } config := &tls.Config{Certificates: []tls.Certificate{cer}} listener, err := tls.Listen("tcp", address, config) if err != nil { - log.Fatal("Can't start --input-tcp with secure connection:", err) + log.Fatalln("[INPUT-TCP] failed to start INPUT-TCP listener:", err) } i.listener = listener } else { listener, err := net.Listen("tcp", address) if err != nil { - log.Fatal("Can't start:", err) + log.Fatalln("failed to start INPUT-TCP listener:", err) } - i.listener = listener } - go func() { for { conn, err := i.listener.Accept() - - if err != nil { - log.Println("Error while Accept()", err) + if err == nil { + go i.handleConnection(conn) continue } - - go i.handleConnection(conn) + if isTemporaryNetworkError(err) { + continue + } + if operr, ok := err.(*net.OpError); ok && operr.Err.Error() != "use of closed network connection" { + Debug(0, fmt.Sprintf("[INPUT-TCP] listener closed, err: %q", err)) + } + break } }() } +var payloadSeparatorAsBytes = []byte(payloadSeparator) + func (i *TCPInput) handleConnection(conn net.Conn) { defer conn.Close() - payloadSeparatorAsBytes := []byte(payloadSeparator) reader := bufio.NewReader(conn) var buffer bytes.Buffer for { line, err := reader.ReadBytes('\n') - if err != nil { + if isTemporaryNetworkError(err) { + continue + } if err != io.EOF { - fmt.Fprintln(os.Stderr, "Unexpected error in input tcp connection:", err) + Debug(0, fmt.Sprintf("[INPUT-TCP] connection error: %q", err)) } break - } if bytes.Equal(payloadSeparatorAsBytes[1:], line) { - asBytes := buffer.Bytes() + // unread the '\n' before monkeys + buffer.UnreadByte() + var msg Message + msg.Meta, msg.Data = payloadMetaWithBody(buffer.Bytes()) + i.data <- &msg buffer.Reset() - - newBuf := make([]byte, len(asBytes)-1) - copy(newBuf, asBytes) - - i.data <- newBuf } else { buffer.Write(line) } @@ -127,3 +131,13 @@ func (i *TCPInput) handleConnection(conn net.Conn) { func (i *TCPInput) String() string { return "TCP input: " + i.address } + +func isTemporaryNetworkError(err error) bool { + if nerr, ok := err.(net.Error); ok && nerr.Temporary() { + return true + } + if operr, ok := err.(*net.OpError); ok && operr.Temporary() { + return true + } + return false +} diff --git a/input_tcp_test.go b/input_tcp_test.go index f470086..7714ce3 100644 --- a/input_tcp_test.go +++ b/input_tcp_test.go @@ -7,7 +7,6 @@ import ( "crypto/tls" "crypto/x509" "encoding/pem" - "io" "io/ioutil" "log" "math/big" @@ -20,20 +19,19 @@ import ( func TestTCPInput(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTCPInput("127.0.0.1:0", &TCPInputConfig{}) - output := NewTestOutput(func(data []byte) { + output := NewTestOutput(func(*Message) { wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) tcpAddr, err := net.ResolveTCPAddr("tcp", input.listener.Addr().String()) @@ -43,7 +41,6 @@ func TestTCPInput(t *testing.T) { } conn, err := net.DialTCP("tcp", nil, tcpAddr) - if err != nil { log.Fatal(err) } @@ -52,10 +49,14 @@ func TestTCPInput(t *testing.T) { for i := 0; i < 100; i++ { wg.Add(1) - conn.Write(msg) - conn.Write([]byte(payloadSeparator)) + if _, err = conn.Write(msg); err == nil { + _, err = conn.Write(payloadSeparatorAsBytes) + } + if err != nil { + t.Error(err) + return + } } - wg.Wait() emitter.Close() } @@ -99,24 +100,23 @@ func TestTCPInputSecure(t *testing.T) { }() wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTCPInput("127.0.0.1:0", &TCPInputConfig{ Secure: true, CertificatePath: serverCertPemFile.Name(), KeyPath: serverPrivPemFile.Name(), }) - output := NewTestOutput(func(data []byte) { + output := NewTestOutput(func(*Message) { wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) conf := &tls.Config{ diff --git a/limiter.go b/limiter.go index 77ddfff..2d89f8b 100644 --- a/limiter.go +++ b/limiter.go @@ -20,8 +20,8 @@ type Limiter struct { } func parseLimitOptions(options string) (limit int, isPercent bool) { - if strings.Contains(options, "%") { - limit, _ = strconv.Atoi(strings.Split(options, "%")[0]) + if n := strings.Index(options, "%"); n > 0 { + limit, _ = strconv.Atoi(options[:n]) isPercent = true } else { limit, _ = strconv.Atoi(options) @@ -33,7 +33,7 @@ func parseLimitOptions(options string) (limit int, isPercent bool) { // NewLimiter constructor for Limiter, accepts plugin and options // `options` allow to sprcify relatve or absolute limiting -func NewLimiter(plugin interface{}, options string) io.ReadWriter { +func NewLimiter(plugin interface{}, options string) PluginReadWriter { l := new(Limiter) l.limit, l.isPercent = parseLimitOptions(options) l.plugin = plugin @@ -71,24 +71,29 @@ func (l *Limiter) isLimited() bool { return false } -func (l *Limiter) Write(data []byte) (n int, err error) { +// PluginWrite writes message to this plugin +func (l *Limiter) PluginWrite(msg *Message) (n int, err error) { if l.isLimited() { return 0, nil } - - n, err = l.plugin.(io.Writer).Write(data) - return + if w, ok := l.plugin.(PluginWriter); ok { + return w.PluginWrite(msg) + } + // avoid further writing + return 0, io.ErrClosedPipe } -func (l *Limiter) Read(data []byte) (n int, err error) { - if r, ok := l.plugin.(io.Reader); ok { - n, err = r.Read(data) +// PluginRead reads message from this plugin +func (l *Limiter) PluginRead() (msg *Message, err error) { + if r, ok := l.plugin.(PluginReader); ok { + msg, err = r.PluginRead() } else { - return 0, nil + // avoid further reading + return nil, io.ErrClosedPipe } if l.isLimited() { - return 0, nil + return nil, nil } return @@ -100,7 +105,7 @@ func (l *Limiter) String() string { // Close closes the resources. func (l *Limiter) Close() error { - if fi, ok := l.plugin.(io.ReadCloser); ok { + if fi, ok := l.plugin.(io.Closer); ok { fi.Close() } return nil diff --git a/limiter_test.go b/limiter_test.go index 39e6b1c..e689002 100644 --- a/limiter_test.go +++ b/limiter_test.go @@ -3,28 +3,26 @@ package main import ( - "io" "sync" "testing" ) func TestOutputLimiter(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() - output := NewLimiter(NewTestOutput(func(data []byte) { + output := NewLimiter(NewTestOutput(func(*Message) { wg.Done() }), "10") wg.Add(10) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 100; i++ { @@ -37,21 +35,20 @@ func TestOutputLimiter(t *testing.T) { func TestInputLimiter(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewLimiter(NewTestInput(), "10") - output := NewTestOutput(func(data []byte) { + output := NewTestOutput(func(*Message) { wg.Done() }) wg.Add(10) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 100; i++ { @@ -65,20 +62,19 @@ func TestInputLimiter(t *testing.T) { // Should limit all requests func TestPercentLimiter1(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() - output := NewLimiter(NewTestOutput(func(data []byte) { + output := NewLimiter(NewTestOutput(func(*Message) { wg.Done() }), "0%") plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 100; i++ { @@ -91,21 +87,20 @@ func TestPercentLimiter1(t *testing.T) { // Should not limit at all func TestPercentLimiter2(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() - output := NewLimiter(NewTestOutput(func(data []byte) { + output := NewLimiter(NewTestOutput(func(*Message) { wg.Done() }), "100%") wg.Add(100) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 100; i++ { diff --git a/middleware.go b/middleware.go index edcdf43..726632a 100644 --- a/middleware.go +++ b/middleware.go @@ -2,6 +2,7 @@ package main import ( "bufio" + "context" "encoding/hex" "fmt" "io" @@ -9,32 +10,29 @@ import ( "os" "os/exec" "strings" - "sync" ) // Middleware represents a middleware object type Middleware struct { - command string - - data chan []byte - - mu sync.Mutex - - Stdin io.Writer - Stdout io.Reader - - stop chan bool // Channel used only to indicate goroutine should shutdown + command string + data chan *Message + Stdin io.Writer + Stdout io.Reader + commandCancel context.CancelFunc + stop chan bool // Channel used only to indicate goroutine should shutdown } // NewMiddleware returns new middleware func NewMiddleware(command string) *Middleware { m := new(Middleware) m.command = command - m.data = make(chan []byte, 1000) + m.data = make(chan *Message, 1000) m.stop = make(chan bool) commands := strings.Split(command, " ") - cmd := exec.Command(commands[0], commands[1:]...) + ctx, cancl := context.WithCancel(context.Background()) + m.commandCancel = cancl + cmd := exec.CommandContext(ctx, commands[0], commands[1:]...) m.Stdout, _ = cmd.StdoutPipe() m.Stdin, _ = cmd.StdinPipe() @@ -61,45 +59,32 @@ func NewMiddleware(command string) *Middleware { } // ReadFrom start a worker to read from this plugin -func (m *Middleware) ReadFrom(plugin io.Reader) { +func (m *Middleware) ReadFrom(plugin PluginReader) { Debug(2, "[MIDDLEWARE-MASTER] Starting reading from", plugin) go m.copy(m.Stdin, plugin) } -func (m *Middleware) copy(to io.Writer, from io.Reader) { - buf := make([]byte, 5*1024*1024) - dst := make([]byte, len(buf)*4) +func (m *Middleware) copy(to io.Writer, from PluginReader) { + var buf, dst []byte for { - nr, _ := from.Read(buf) - if nr == 0 || nr > len(buf) { + msg, err := from.PluginRead() + if err != nil { + return + } + if msg == nil || len(msg.Data) == 0 { continue } - payload := buf[0:nr] - if Settings.PrettifyHTTP { - payload = prettifyHTTP(payload) - nr = len(payload) - - if nr*2 > len(dst) { - continue - } + buf = prettifyHTTP(msg.Data) } + dst = make([]byte, len(buf)*2+1) + hex.Encode(dst, buf) + dst[len(buf)*2] = '\n' - if Settings.PrettifyHTTP { - payload = prettifyHTTP(payload) - nr = len(payload) - } + to.Write(dst) - hex.Encode(dst, payload) - dst[nr*2] = '\n' - - m.mu.Lock() - to.Write(dst[0 : nr*2+1]) - m.mu.Unlock() - - Debug(3, "[MIDDLEWARE-MASTER] Sending:", string(buf[0:nr]), "From:", from) } } @@ -117,41 +102,41 @@ func (m *Middleware) read(from io.Reader) { } } - buf := make([]byte, len(line)/2) + buf := make([]byte, len(line)/2-1) if _, err := hex.Decode(buf, line[:len(line)-1]); err != nil { - fmt.Fprintln(os.Stderr, "Failed to decode input payload", err, len(line), string(line[:len(line)-1])) + Debug(0, fmt.Sprintf("[MIDDLEWARE] failed to decode err: %q", err)) + continue } - - Debug(3, "[MIDDLEWARE-MASTER] Received:", string(buf)) - + var msg Message + msg.Meta, msg.Data = payloadMetaWithBody(buf) select { case <-m.stop: return - case m.data <- buf: + case m.data <- &msg: } } return } -func (m *Middleware) Read(data []byte) (int, error) { - var buf []byte +// PluginRead reads message from this plugin +func (m *Middleware) PluginRead() (msg *Message, err error) { select { case <-m.stop: - return 0, ErrorStopped - case buf = <-m.data: + return nil, ErrorStopped + case msg = <-m.data: } - n := copy(data, buf) - return n, nil + return } func (m *Middleware) String() string { - return fmt.Sprintf("Modifying traffic using '%s' command", m.command) + return fmt.Sprintf("Modifying traffic using %q command", m.command) } // Close closes this plugin func (m *Middleware) Close() error { + m.commandCancel() close(m.stop) return nil } diff --git a/output_binary.go b/output_binary.go index aba86fb..32b3e3a 100644 --- a/output_binary.go +++ b/output_binary.go @@ -1,7 +1,6 @@ package main import ( - "io" "sync/atomic" "time" @@ -25,30 +24,27 @@ type BinaryOutput struct { // 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 - queue chan []byte - - responses chan response - - needWorker chan int - - config *BinaryOutputConfig - - queueStats *GorStat + address string + queue chan *Message + responses chan response + needWorker chan int + quit chan struct{} + config *BinaryOutputConfig + queueStats *GorStat } // NewBinaryOutput constructor for BinaryOutput // Initialize workers -func NewBinaryOutput(address string, config *BinaryOutputConfig) io.Writer { +func NewBinaryOutput(address string, config *BinaryOutputConfig) PluginReadWriter { o := new(BinaryOutput) o.address = address o.config = config - o.queue = make(chan []byte, 1000) + o.queue = make(chan *Message, 1000) o.responses = make(chan response, 1000) o.needWorker = make(chan int, 1) + o.quit = make(chan struct{}) // Initial workers count if o.config.Workers == 0 { @@ -89,8 +85,8 @@ func (o *BinaryOutput) startWorker() { for { select { - case data := <-o.queue: - o.sendRequest(client, data) + case msg := <-o.queue: + o.sendRequest(client, msg) deathCount = 0 case <-time.After(time.Millisecond * 100): // When dynamic scaling enabled workers die after 2s of inactivity @@ -113,15 +109,13 @@ func (o *BinaryOutput) startWorker() { } } -func (o *BinaryOutput) Write(data []byte) (n int, err error) { - if !isRequestPayload(data) { - return len(data), nil +// PluginWrite writes a message tothis plugin +func (o *BinaryOutput) PluginWrite(msg *Message) (n int, err error) { + if !isRequestPayload(msg.Meta) { + return len(msg.Data), nil } - buf := make([]byte, len(data)) - copy(buf, data) - - o.queue <- buf + o.queue <- msg if o.config.Workers == 0 { workersCount := atomic.LoadInt64(&o.activeWorkers) @@ -131,37 +125,33 @@ func (o *BinaryOutput) Write(data []byte) (n int, err error) { } } - return len(data), nil + return len(msg.Data) + len(msg.Meta), nil } -func (o *BinaryOutput) Read(data []byte) (int, error) { - resp := <-o.responses +// PluginRead reads a message from this plugin +func (o *BinaryOutput) PluginRead() (*Message, error) { + var resp response + var msg Message + select { + case <-o.quit: + return nil, ErrorStopped + case resp = <-o.responses: + } + msg.Data = resp.payload + msg.Meta = payloadHeader(ReplayedResponsePayload, resp.uuid, resp.startedAt, resp.roundTripTime) - Debug(2, "[OUTPUT-TCP] Received response:", string(resp.payload)) - - header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.startedAt, resp.roundTripTime) - copy(data[0:len(header)], header) - copy(data[len(header):], resp.payload) - - return len(resp.payload) + len(header), nil + return &msg, nil } -func (o *BinaryOutput) sendRequest(client *TCPClient, request []byte) { - meta := payloadMeta(request) - if len(meta) < 2 { +func (o *BinaryOutput) sendRequest(client *TCPClient, msg *Message) { + if !isRequestPayload(msg.Meta) { return } - if !isRequestPayload(request) { - return - } - - uuid := meta[1] - - body := payloadBody(request) + uuid := payloadID(msg.Meta) start := time.Now() - resp, err := client.Send(body) + resp, err := client.Send(msg.Data) stop := time.Now() if err != nil { @@ -176,3 +166,9 @@ func (o *BinaryOutput) sendRequest(client *TCPClient, request []byte) { func (o *BinaryOutput) String() string { return "Binary output: " + o.address } + +// Close closes this plugin for reading +func (o *BinaryOutput) Close() error { + close(o.quit) + return nil +} diff --git a/output_dummy.go b/output_dummy.go index eca0f53..e28653f 100644 --- a/output_dummy.go +++ b/output_dummy.go @@ -15,9 +15,15 @@ func NewDummyOutput() (di *DummyOutput) { return } -func (i *DummyOutput) Write(data []byte) (int, error) { - n, err := os.Stdout.Write(data) - os.Stdout.Write([]byte{'\n'}) +// PluginWrite writes message to this plugin +func (i *DummyOutput) PluginWrite(msg *Message) (int, error) { + var n, nn int + var err error + n, err = os.Stdout.Write(msg.Meta) + nn, err = os.Stdout.Write(msg.Data) + n += nn + nn, err = os.Stdout.Write(payloadSeparatorAsBytes) + n += nn return n, err } diff --git a/output_file.go b/output_file.go index 8b4dab5..c329ccc 100644 --- a/output_file.go +++ b/output_file.go @@ -195,10 +195,11 @@ func (o *FileOutput) updateName() { o.Unlock() } -func (o *FileOutput) Write(data []byte) (n int, err error) { +// PluginWrite writes message to this plugin +func (o *FileOutput) PluginWrite(msg *Message) (n int, err error) { if o.requestPerFile { o.Lock() - meta := payloadMeta(data) + meta := payloadMeta(msg.Meta) o.currentID = meta[1] o.payloadType = meta[0] o.Unlock() @@ -227,10 +228,12 @@ func (o *FileOutput) Write(data []byte) (n int, err error) { o.QueueLength = 0 } - n, _ = o.writer.Write(data) - nSeparator, _ := o.writer.Write([]byte(payloadSeparator)) - - n += nSeparator + var nn int + n, err = o.writer.Write(msg.Meta) + nn, err = o.writer.Write(msg.Data) + n += nn + nn, err = o.writer.Write(payloadSeparatorAsBytes) + n += nn o.totalFileSize += size.Size(n) o.QueueLength++ @@ -239,14 +242,14 @@ func (o *FileOutput) Write(data []byte) (n int, err error) { return n, errors.New("File output reached size limit") } - return n, nil + return n, err } func (o *FileOutput) flush() { // Don't exit on panic defer func() { if r := recover(); r != nil { - log.Println("PANIC while file flush: ", r, o, string(debug.Stack())) + Debug(0, "[OUTPUT-FILE] PANIC while file flush: ", r, o, string(debug.Stack())) } }() @@ -263,7 +266,7 @@ func (o *FileOutput) flush() { if stat, err := o.file.Stat(); err == nil { o.chunkSize = int(stat.Size()) } else { - log.Println("Error accessing file sats", err) + Debug(0, "[OUTPUT-HTTP] error accessing file size", err) } } } diff --git a/output_file_test.go b/output_file_test.go index 8dcccc0..d68c183 100644 --- a/output_file_test.go +++ b/output_file_test.go @@ -2,7 +2,6 @@ package main import ( "fmt" - "io" "math/rand" "os" "reflect" @@ -17,18 +16,17 @@ import ( func TestFileOutput(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() output := NewFileOutput("/tmp/test_requests.gor", &FileOutputConfig{FlushInterval: time.Minute, Append: true}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 100; i++ { @@ -42,19 +40,18 @@ func TestFileOutput(t *testing.T) { var counter int64 input2 := NewFileInput("/tmp/test_requests.gor", false) - output2 := NewTestOutput(func(data []byte) { + output2 := NewTestOutput(func(*Message) { atomic.AddInt64(&counter, 1) wg.Done() }) plugins2 := &InOutPlugins{ - Inputs: []io.Reader{input2}, - Outputs: []io.Writer{output2}, + Inputs: []PluginReader{input2}, + Outputs: []PluginWriter{output2}, } plugins2.All = append(plugins2.All, input2, output2) - quit2 := make(chan int) - emitter2 := NewEmitter(quit2) + emitter2 := NewEmitter() go emitter2.Start(plugins2, Settings.Middleware) wg.Wait() @@ -91,16 +88,16 @@ func TestFileOutputMultipleFiles(t *testing.T) { t.Error("Should not initialize file if no writes") } - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name1 := output.file.Name() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name2 := output.file.Name() time.Sleep(time.Second) output.updateName() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name3 := output.file.Name() if name2 != name1 { @@ -122,16 +119,16 @@ func TestFileOutputFilePerRequest(t *testing.T) { t.Error("Should not initialize file if no writes") } - output.Write([]byte("1 1 1\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name1 := output.file.Name() - output.Write([]byte("1 2 1\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 2 1\r\n"), Data: []byte("test")}) name2 := output.file.Name() time.Sleep(time.Second) output.updateName() - output.Write([]byte("1 3 1\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 3 1\r\n"), Data: []byte("test")}) name3 := output.file.Name() if name3 == name2 || name2 == name1 || name3 == name1 { @@ -151,7 +148,7 @@ func TestFileOutputCompression(t *testing.T) { } for i := 0; i < 1000; i++ { - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) } name := output.file.Name() @@ -210,15 +207,15 @@ func TestFileOutputAppendQueueLimitOverflow(t *testing.T) { output := NewFileOutput(name, &FileOutputConfig{Append: false, FlushInterval: time.Minute, QueueLimit: 2}) - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name1 := output.file.Name() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name2 := output.file.Name() output.updateName() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name3 := output.file.Name() if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0", rnd) { @@ -239,15 +236,15 @@ func TestFileOutputAppendQueueLimitNoOverflow(t *testing.T) { output := NewFileOutput(name, &FileOutputConfig{Append: false, FlushInterval: time.Minute, QueueLimit: 3}) - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name1 := output.file.Name() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name2 := output.file.Name() output.updateName() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name3 := output.file.Name() if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0", rnd) { @@ -268,15 +265,15 @@ func TestFileOutputAppendQueueLimitGzips(t *testing.T) { output := NewFileOutput(name, &FileOutputConfig{Append: false, FlushInterval: time.Minute, QueueLimit: 2}) - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name1 := output.file.Name() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name2 := output.file.Name() output.updateName() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name3 := output.file.Name() if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0.gz", rnd) { @@ -311,16 +308,16 @@ func TestFileOutputAppendSizeLimitOverflow(t *testing.T) { output := NewFileOutput(name, &FileOutputConfig{Append: false, FlushInterval: time.Minute, SizeLimit: size.Size(2 * messageSize)}) - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name1 := output.file.Name() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name2 := output.file.Name() output.flush() output.updateName() - output.Write([]byte("1 1 1\r\ntest")) + output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")}) name3 := output.file.Name() if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0", rnd) { diff --git a/output_http.go b/output_http.go index 7540fed..2dde5fe 100644 --- a/output_http.go +++ b/output_http.go @@ -5,7 +5,6 @@ import ( "bytes" "crypto/tls" "fmt" - "io" "log" "math" "net/http" @@ -59,14 +58,14 @@ type HTTPOutput struct { elasticSearch *ESPlugin client *HTTPClient stopWorker chan struct{} - queue chan []byte + queue chan *Message responses chan response stop chan bool // Channel used only to indicate goroutine should shutdown } // NewHTTPOutput constructor for HTTPOutput // Initialize workers -func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer { +func NewHTTPOutput(address string, config *HTTPOutputConfig) PluginReadWriter { o := new(HTTPOutput) var err error config.url, err = url.Parse(address) @@ -110,7 +109,7 @@ func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer { o.queueStats = NewGorStat("output_http", o.config.StatsMs) } - o.queue = make(chan []byte, o.config.QueueLen) + o.queue = make(chan *Message, o.config.QueueLen) o.responses = make(chan response, o.config.QueueLen) // it should not be buffered to avoid races o.stopWorker = make(chan struct{}) @@ -160,23 +159,22 @@ func (o *HTTPOutput) startWorker() { select { case <-o.stopWorker: return - case data := <-o.queue: - o.sendRequest(o.client, data) + case msg := <-o.queue: + o.sendRequest(o.client, msg) } } } -func (o *HTTPOutput) Write(data []byte) (n int, err error) { - if !isRequestPayload(data) { - return len(data), nil +// PluginWrite writes message to this plugin +func (o *HTTPOutput) PluginWrite(msg *Message) (n int, err error) { + if !isRequestPayload(msg.Meta) { + return len(msg.Data), nil } - buf := make([]byte, len(data)) - copy(buf, data) select { case <-o.stop: return 0, ErrorStopped - case o.queue <- buf: + case o.queue <- msg: } if o.config.Stats { @@ -189,51 +187,48 @@ func (o *HTTPOutput) Write(data []byte) (n int, err error) { atomic.AddInt32(&o.activeWorkers, 1) } } - return len(data), nil + return len(msg.Data) + len(msg.Meta), nil } -func (o *HTTPOutput) Read(data []byte) (int, error) { +// PluginRead reads message from this plugin +func (o *HTTPOutput) PluginRead() (*Message, error) { var resp response + var msg Message select { case <-o.stop: - return 0, ErrorStopped + return nil, ErrorStopped case resp = <-o.responses: + msg.Data = resp.payload } - header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.roundTripTime, resp.startedAt) - n := copy(data, header) - if len(data) > len(header) { - n += copy(data[len(header):], resp.payload) - } - dis := len(header) + len(data) - n - if dis > 0 { - Debug(2, fmt.Sprintf("[OUTPUT-HTTP] %dB discarded increase copy buffer size", dis)) - } + msg.Meta = payloadHeader(ReplayedResponsePayload, resp.uuid, resp.roundTripTime, resp.startedAt) - return n, nil + return &msg, nil } -func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) { - if !isRequestPayload(request) { +func (o *HTTPOutput) sendRequest(client *HTTPClient, msg *Message) { + if !isRequestPayload(msg.Meta) { return } - uuid := payloadID(request) - body := payloadBody(request) + uuid := payloadID(msg.Meta) start := time.Now() - resp, err := client.Send(body) + resp, err := client.Send(msg.Data) stop := time.Now() if err != nil { Debug(1, fmt.Sprintf("[HTTP-OUTPUT] error when sending: %q", err)) return } + if resp == nil { + return + } if o.config.TrackResponses { o.responses <- response{resp, uuid, start.UnixNano(), stop.UnixNano() - start.UnixNano()} } if o.elasticSearch != nil { - o.elasticSearch.ResponseAnalyze(request, resp, start, stop) + o.elasticSearch.ResponseAnalyze(msg.Data, resp, start, stop) } } diff --git a/output_http_test.go b/output_http_test.go index 279e4ed..91ec5b7 100644 --- a/output_http_test.go +++ b/output_http_test.go @@ -1,7 +1,6 @@ package main import ( - "io" "io/ioutil" "net/http" "net/http/httptest" @@ -12,7 +11,6 @@ import ( func TestHTTPOutput(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() @@ -38,27 +36,27 @@ func TestHTTPOutput(t *testing.T) { })) defer server.Close() - headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}} + headers := HTTPHeaders{httpHeader{"User-Agent", "Gor"}} methods := HTTPMethods{[]byte("GET"), []byte("PUT"), []byte("POST")} Settings.ModifierConfig = HTTPModifierConfig{Headers: headers, Methods: methods} - httpOutput := NewHTTPOutput(server.URL, &HTTPOutputConfig{TrackResponses: true}) - output := NewTestOutput(func(data []byte) { + httpOutput := NewHTTPOutput(server.URL, &HTTPOutputConfig{TrackResponses: false}) + output := NewTestOutput(func(*Message) { wg.Done() }) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{httpOutput, output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{httpOutput, output}, } plugins.All = append(plugins.All, input, output, httpOutput) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 10; i++ { - // 2 http-output, 2 - test output request, 2 - test output http response - wg.Add(6) // OPTIONS should be ignored + // 2 http-output, 2 - test output request + wg.Add(4) // OPTIONS should be ignored input.EmitPOST() input.EmitOPTIONS() input.EmitGET() @@ -72,7 +70,6 @@ func TestHTTPOutput(t *testing.T) { func TestHTTPOutputKeepOriginalHost(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() @@ -85,18 +82,18 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) { })) defer server.Close() - headers := HTTPHeaders{HTTPHeader{"Host", "custom-host.com"}} + headers := HTTPHeaders{httpHeader{"Host", "custom-host.com"}} Settings.ModifierConfig = HTTPModifierConfig{Headers: headers} output := NewHTTPOutput(server.URL, &HTTPOutputConfig{OriginalHost: true, SkipVerify: true}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) wg.Add(1) @@ -109,7 +106,6 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) { func TestHTTPOutputSSL(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) // Origing and Replay server initialization server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -120,12 +116,12 @@ func TestHTTPOutputSSL(t *testing.T) { output := NewHTTPOutput(server.URL, &HTTPOutputConfig{SkipVerify: true}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) wg.Add(2) @@ -139,7 +135,6 @@ func TestHTTPOutputSSL(t *testing.T) { func TestHTTPOutputSessions(t *testing.T) { wg := new(sync.WaitGroup) - quit := make(chan int) input := NewTestInput() input.skipHeader = true @@ -150,25 +145,27 @@ func TestHTTPOutputSessions(t *testing.T) { defer server.Close() Settings.RecognizeTCPSessions = true + Settings.SplitOutput = true output := NewHTTPOutput(server.URL, &HTTPOutputConfig{}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } - emitter := NewEmitter(quit) + plugins.All = append(plugins.All, input, output) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) uuid1 := []byte("1234567890123456789a0000") uuid2 := []byte("1234567890123456789d0000") - for i := 0; i < 100; i++ { + for i := 0; i < 10; i++ { wg.Add(1) // OPTIONS should be ignored copy(uuid1[20:], randByte(4)) input.EmitBytes([]byte("1 " + string(uuid1) + " 1\n" + "GET / HTTP/1.1\r\n\r\n")) } - for i := 0; i < 100; i++ { + for i := 0; i < 10; i++ { wg.Add(1) // OPTIONS should be ignored copy(uuid2[20:], randByte(4)) input.EmitBytes([]byte("1 " + string(uuid2) + " 1\n" + "GET / HTTP/1.1\r\n\r\n")) @@ -179,11 +176,11 @@ func TestHTTPOutputSessions(t *testing.T) { emitter.Close() Settings.RecognizeTCPSessions = false + Settings.SplitOutput = false } func BenchmarkHTTPOutput(b *testing.B) { wg := new(sync.WaitGroup) - quit := make(chan int) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { wg.Done() @@ -194,12 +191,12 @@ func BenchmarkHTTPOutput(b *testing.B) { output := NewHTTPOutput(server.URL, &HTTPOutputConfig{WorkersMax: 1}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < b.N; i++ { @@ -213,7 +210,6 @@ func BenchmarkHTTPOutput(b *testing.B) { func BenchmarkHTTPOutputTLS(b *testing.B) { wg := new(sync.WaitGroup) - quit := make(chan int) server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { wg.Done() @@ -224,12 +220,12 @@ func BenchmarkHTTPOutputTLS(b *testing.B) { output := NewHTTPOutput(server.URL, &HTTPOutputConfig{SkipVerify: true, WorkersMax: 1}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < b.N; i++ { diff --git a/output_kafka.go b/output_kafka.go index 61e7f30..40a2bf8 100644 --- a/output_kafka.go +++ b/output_kafka.go @@ -2,7 +2,6 @@ package main import ( "encoding/json" - "io" "log" "strings" "time" @@ -24,7 +23,7 @@ type KafkaOutput struct { const KafkaOutputFrequency = 500 // NewKafkaOutput creates instance of kafka producer client with TLS config -func NewKafkaOutput (address string, config *OutputKafkaConfig, tlsConfig *KafkaTLSConfig) io.Writer { +func NewKafkaOutput(address string, config *OutputKafkaConfig, tlsConfig *KafkaTLSConfig) PluginWriter { c := NewKafkaConfig(tlsConfig) var producer sarama.AsyncProducer @@ -63,20 +62,21 @@ func (o *KafkaOutput) ErrorHandler() { } } -func (o *KafkaOutput) Write(data []byte) (n int, err error) { +// PluginWrite writes a message to this plugin +func (o *KafkaOutput) PluginWrite(msg *Message) (n int, err error) { var message sarama.StringEncoder if !o.config.UseJSON { - message = sarama.StringEncoder(byteutils.SliceToString(data)) + message = sarama.StringEncoder(byteutils.SliceToString(msg.Meta) + byteutils.SliceToString(msg.Data)) } else { - mimeHeader := proto.ParseHeaders(data) + mimeHeader := proto.ParseHeaders(msg.Data) var header map[string]string for k, v := range mimeHeader { header[k] = strings.Join(v, ", ") } - meta := payloadMeta(data) - req := payloadBody(data) + meta := payloadMeta(msg.Meta) + req := msg.Data kafkaMessage := KafkaMessage{ ReqURL: byteutils.SliceToString(proto.Path(req)), diff --git a/output_kafka_test.go b/output_kafka_test.go index 67fb464..b7227c6 100644 --- a/output_kafka_test.go +++ b/output_kafka_test.go @@ -17,16 +17,16 @@ func TestOutputKafkaRAW(t *testing.T) { producer: producer, Topic: "test", UseJSON: false, - },nil) + }, nil) - output.Write([]byte("1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n")) + output.PluginWrite(&Message{Meta: []byte("1 2 3\n"), Data: []byte("GET / HTTP1.1\r\nHeader: 1\r\n\r\n")}) resp := <-producer.Successes() data, _ := resp.Value.Encode() if string(data) != "1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n" { - t.Error("Message not properly encoded: ", string(data)) + t.Errorf("Message not properly encoded: %q", data) } } @@ -42,7 +42,7 @@ func TestOutputKafkaJSON(t *testing.T) { UseJSON: true, }, nil) - output.Write([]byte("1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n")) + output.PluginWrite(&Message{Meta: []byte("1 2 3\n"), Data: []byte("GET / HTTP1.1\r\nHeader: 1\r\n\r\n")}) resp := <-producer.Successes() diff --git a/output_null.go b/output_null.go index ee5fdaa..0957ed9 100644 --- a/output_null.go +++ b/output_null.go @@ -9,8 +9,9 @@ func NewNullOutput() (o *NullOutput) { return new(NullOutput) } -func (o *NullOutput) Write(data []byte) (int, error) { - return len(data), nil +// PluginWrite writes message to this plugin +func (o *NullOutput) PluginWrite(msg *Message) (int, error) { + return len(msg.Data) + len(msg.Meta), nil } func (o *NullOutput) String() string { diff --git a/output_s3.go b/output_s3.go index 3f2d2d8..b135d4d 100644 --- a/output_s3.go +++ b/output_s3.go @@ -67,8 +67,9 @@ func (o *S3Output) connect() { } } -func (o *S3Output) Write(data []byte) (n int, err error) { - return o.buffer.Write(data) +// PluginWrite writes message to this plugin +func (o *S3Output) PluginWrite(msg *Message) (n int, err error) { + return o.buffer.PluginWrite(msg) } func (o *S3Output) String() string { diff --git a/output_tcp.go b/output_tcp.go index fa0d5fc..82b99d5 100644 --- a/output_tcp.go +++ b/output_tcp.go @@ -4,7 +4,6 @@ import ( "crypto/tls" "fmt" "hash/fnv" - "io" "net" "time" ) @@ -15,7 +14,7 @@ import ( type TCPOutput struct { address string limit int - buf []chan []byte + buf []chan *Message bufStats *GorStat config *TCPOutputConfig workerIndex uint32 @@ -31,7 +30,7 @@ type TCPOutputConfig struct { // NewTCPOutput constructor for TCPOutput // Initialize X workers which hold keep-alive connection -func NewTCPOutput(address string, config *TCPOutputConfig) io.Writer { +func NewTCPOutput(address string, config *TCPOutputConfig) PluginWriter { o := new(TCPOutput) o.address = address @@ -42,9 +41,9 @@ func NewTCPOutput(address string, config *TCPOutputConfig) io.Writer { } // create X buffers and send the buffer index to the worker - o.buf = make([]chan []byte, o.config.Workers) + o.buf = make([]chan *Message, o.config.Workers) for i := 0; i < o.config.Workers; i++ { - o.buf[i] = make(chan []byte, 100) + o.buf[i] = make(chan *Message, 100) go o.worker(i) } @@ -73,14 +72,16 @@ func (o *TCPOutput) worker(bufferIndex int) { defer conn.Close() for { - data := <-o.buf[bufferIndex] - if _, err = conn.Write(data); err == nil { - _, err = conn.Write([]byte(payloadSeparator)) + msg := <-o.buf[bufferIndex] + if _, err = conn.Write(msg.Meta); err == nil { + if _, err = conn.Write(msg.Data); err == nil { + _, err = conn.Write(payloadSeparatorAsBytes) + } } if err != nil { Debug(2, "INFO: TCP output connection closed, reconnecting") - o.buf[bufferIndex] <- data + o.buf[bufferIndex] <- msg go o.worker(bufferIndex) break } @@ -99,23 +100,20 @@ func (o *TCPOutput) getBufferIndex(data []byte) int { } -func (o *TCPOutput) Write(data []byte) (n int, err error) { - if !isOriginPayload(data) { - return len(data), nil +// PluginWrite writes message to this plugin +func (o *TCPOutput) PluginWrite(msg *Message) (n int, err error) { + if !isOriginPayload(msg.Meta) { + return len(msg.Data), nil } - // We have to copy, because sending data in multiple threads - newBuf := make([]byte, len(data)) - copy(newBuf, data) - - bufferIndex := o.getBufferIndex(data) - o.buf[bufferIndex] <- newBuf + bufferIndex := o.getBufferIndex(msg.Data) + o.buf[bufferIndex] <- msg if Settings.OutputTCPStats { o.bufStats.Write(len(o.buf[bufferIndex])) } - return len(data), nil + return len(msg.Data) + len(msg.Meta), nil } func (o *TCPOutput) connect(address string) (conn net.Conn, err error) { diff --git a/output_tcp_test.go b/output_tcp_test.go index 6ac77d9..6026e8d 100644 --- a/output_tcp_test.go +++ b/output_tcp_test.go @@ -2,7 +2,6 @@ package main import ( "bufio" - "io" "log" "net" "sync" @@ -11,9 +10,7 @@ import ( ) func TestTCPOutput(t *testing.T) { - Settings.Verbose = 2 wg := new(sync.WaitGroup) - quit := make(chan int) listener := startTCP(func(data []byte) { wg.Done() @@ -22,12 +19,11 @@ func TestTCPOutput(t *testing.T) { output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{Workers: 10}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } - plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() go emitter.Start(plugins, Settings.Middleware) for i := 0; i < 10; i++ { @@ -68,7 +64,6 @@ func startTCP(cb func([]byte)) net.Listener { func BenchmarkTCPOutput(b *testing.B) { wg := new(sync.WaitGroup) - quit := make(chan int) listener := startTCP(func(data []byte) { wg.Done() @@ -82,12 +77,11 @@ func BenchmarkTCPOutput(b *testing.B) { output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{Workers: 10}) plugins := &InOutPlugins{ - Inputs: []io.Reader{input}, - Outputs: []io.Writer{output}, + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, } - plugins.All = append(plugins.All, input, output) - emitter := NewEmitter(quit) + emitter := NewEmitter() // avoid counting above initialization b.ResetTimer() go emitter.Start(plugins, Settings.Middleware) @@ -109,7 +103,7 @@ func TestStickyDisable(t *testing.T) { func TestBufferDistribution(t *testing.T) { numberOfWorkers := 10 - numberOfMessages := 1000000 + numberOfMessages := 10000 percentDistributionErrorRange := 20 buffer := make([]int, numberOfWorkers) diff --git a/plugins.go b/plugins.go index 151d189..7ea6ab3 100644 --- a/plugins.go +++ b/plugins.go @@ -1,15 +1,36 @@ package main import ( - "io" "reflect" "strings" ) +// Message represents data accross plugins +type Message struct { + Meta []byte // metadata + Data []byte // actual data +} + +// PluginReader is an interface for input plugins +type PluginReader interface { + PluginRead() (msg *Message, err error) +} + +// PluginWriter is an interface for output plugins +type PluginWriter interface { + PluginWrite(msg *Message) (n int, err error) +} + +// PluginReadWriter is an interface for plugins that support reading and writing +type PluginReadWriter interface { + PluginReader + PluginWriter +} + // InOutPlugins struct for holding references to plugins type InOutPlugins struct { - Inputs []io.Reader - Outputs []io.Writer + Inputs []PluginReader + Outputs []PluginWriter All []interface{} } @@ -48,27 +69,21 @@ func (plugins *InOutPlugins) registerPlugin(constructor interface{}, options ... // Calling our constructor with list of given options plugin := vc.Call(vo)[0].Interface() - pluginWrapper := plugin if limit != "" { - pluginWrapper = NewLimiter(plugin, limit) - } else { - pluginWrapper = plugin + plugin = NewLimiter(plugin, limit) } - _, isR := plugin.(io.Reader) - _, isW := plugin.(io.Writer) - // Some of the output can be Readers as well because return responses - if isR && !isW { - plugins.Inputs = append(plugins.Inputs, pluginWrapper.(io.Reader)) + if r, ok := plugin.(PluginReader); ok { + plugins.Inputs = append(plugins.Inputs, r) } - if isW { - plugins.Outputs = append(plugins.Outputs, pluginWrapper.(io.Writer)) + if w, ok := plugin.(PluginWriter); ok { + plugins.Outputs = append(plugins.Outputs, w) } - plugins.All = append(plugins.All, plugin) + } // NewPlugins specify and initialize all available plugins diff --git a/plugins_test.go b/plugins_test.go index 06dde01..3df13b2 100644 --- a/plugins_test.go +++ b/plugins_test.go @@ -12,8 +12,8 @@ func TestPluginsRegistration(t *testing.T) { plugins := NewPlugins() - if len(plugins.Inputs) != 2 { - t.Errorf("Should be 2 inputs %d", len(plugins.Inputs)) + if len(plugins.Inputs) != 3 { + t.Errorf("Should be 3 inputs got %d", len(plugins.Inputs)) } if _, ok := plugins.Inputs[0].(*DummyInput); !ok { diff --git a/protocol.go b/protocol.go index 5a73152..e1d4a00 100644 --- a/protocol.go +++ b/protocol.go @@ -66,6 +66,16 @@ func payloadMeta(payload []byte) [][]byte { return bytes.Split(payload[:headerSize], []byte{' '}) } +func payloadMetaWithBody(payload []byte) (meta, body []byte) { + if i := bytes.IndexByte(payload, '\n'); i > 0 && len(payload) > i+1 { + meta = payload[:i+1] + body = payload[i+1:] + return + } + // we assume the message did not have meta data + return nil, payload +} + func payloadID(payload []byte) (id []byte) { meta := payloadMeta(payload) diff --git a/settings.go b/settings.go index 2a52e03..928de3b 100644 --- a/settings.go +++ b/settings.go @@ -93,14 +93,10 @@ func init() { } flag.BoolVar(&Settings.SplitOutput, "split-output", false, "By default each output gets same traffic. If set to `true` it splits traffic equally among all outputs.") - flag.BoolVar(&Settings.RecognizeTCPSessions, "recognize-tcp-sessions", false, "[PRO] If turned on http output will create separate worker for each TCP session. Splitting output will session based as well.") flag.Var(&Settings.InputDummy, "input-dummy", "Used for testing outputs. Emits 'Get /' request every 1s") - flag.Var(&Settings.OutputDummy, "output-dummy", "DEPRECATED: use --output-stdout instead") - flag.BoolVar(&Settings.OutputStdout, "output-stdout", false, "Used for testing inputs. Just prints to console data coming from inputs.") - flag.BoolVar(&Settings.OutputNull, "output-null", false, "Used for testing inputs. Drops all requests.") flag.Var(&Settings.InputTCP, "input-tcp", "Used for internal communication between Gor instances. Example: \n\t# Receive requests from other Gor instances on 28020 port, and redirect output to staging\n\tgor --input-tcp :28020 --output-http staging.com") @@ -192,32 +188,16 @@ func init() { flag.StringVar(&Settings.KafkaTLSConfig.ClientKey, "kafka-tls-client-key", "", "Client Key for Kafka TLS Config (mandatory with to kafka-tls-client-cert and kafka-tls-client-key)") flag.Var(&Settings.ModifierConfig.Headers, "http-set-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --http-set-header 'User-Agent: Gor'") - flag.Var(&Settings.ModifierConfig.Headers, "output-http-header", "WARNING: `--output-http-header` DEPRECATED, use `--http-set-header` instead") - flag.Var(&Settings.ModifierConfig.HeaderRewrite, "http-rewrite-header", "Rewrite the request header based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --http-rewrite-header Host: (.*).example.com,$1.beta.example.com") - flag.Var(&Settings.ModifierConfig.Params, "http-set-param", "Set request url param, if param already exists it will be overwritten:\n\tgor --input-raw :8080 --output-http staging.com --http-set-param api_key=1") - flag.Var(&Settings.ModifierConfig.Methods, "http-allow-method", "Whitelist of HTTP methods to replay. Anything else will be dropped:\n\tgor --input-raw :8080 --output-http staging.com --http-allow-method GET --http-allow-method OPTIONS") - flag.Var(&Settings.ModifierConfig.URLRegexp, "http-allow-url", "A regexp to match requests against. Filter get matched against full url with domain. Anything else will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-allow-url ^www.") - flag.Var(&Settings.ModifierConfig.URLNegativeRegexp, "http-disallow-url", "A regexp to match requests against. Filter get matched against full url with domain. Anything else will be forwarded:\n\t gor --input-raw :8080 --output-http staging.com --http-disallow-url ^www.") - flag.Var(&Settings.ModifierConfig.URLRewrite, "http-rewrite-url", "Rewrite the request url based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --http-rewrite-url /v1/user/([^\\/]+)/ping:/v2/user/$1/ping") - flag.Var(&Settings.ModifierConfig.URLRewrite, "output-http-rewrite-url", "WARNING: `--output-http-rewrite-url` DEPRECATED, use `--http-rewrite-url` instead") - flag.Var(&Settings.ModifierConfig.HeaderFilters, "http-allow-header", "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 --http-allow-header api-version:^v1") - flag.Var(&Settings.ModifierConfig.HeaderFilters, "output-http-header-filter", "WARNING: `--output-http-header-filter` DEPRECATED, use `--http-allow-header` instead") - flag.Var(&Settings.ModifierConfig.HeaderNegativeFilters, "http-disallow-header", "A regexp to match a specific header against. Requests with matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-disallow-header \"User-Agent: Replayed by Gor\"") - flag.Var(&Settings.ModifierConfig.HeaderBasicAuthFilters, "http-basic-auth-filter", "A regexp to match the decoded basic auth string against. Requests with non-matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-basic-auth-filter \"^customer[0-9].*\"") - flag.Var(&Settings.ModifierConfig.HeaderHashFilters, "http-header-limiter", "Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific header:\n\t gor --input-raw :8080 --output-http staging.com --http-header-limiter user-id:25%") - - flag.Var(&Settings.ModifierConfig.HeaderHashFilters, "output-http-header-hash-filter", "WARNING: `output-http-header-hash-filter` DEPRECATED, use `--http-header-hash-limiter` instead") - flag.Var(&Settings.ModifierConfig.ParamHashFilters, "http-param-limiter", "Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific GET param:\n\t gor --input-raw :8080 --output-http staging.com --http-param-limiter user_id:25%") // default values, using for tests diff --git a/test_input.go b/test_input.go index 4a461cf..954ccef 100644 --- a/test_input.go +++ b/test_input.go @@ -25,22 +25,21 @@ func NewTestInput() (i *TestInput) { return } -func (i *TestInput) Read(data []byte) (int, error) { +// PluginRead reads message from this plugin +func (i *TestInput) PluginRead() (*Message, error) { + var msg Message select { case buf := <-i.data: - var header []byte - + msg.Data = buf if !i.skipHeader { - header = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1) - copy(data[0:len(header)], header) - copy(data[len(header):], buf) + msg.Meta = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1) } else { - copy(data, buf) + msg.Meta, msg.Data = payloadMetaWithBody(msg.Data) } - return len(buf) + len(header), nil + return &msg, nil case <-i.stop: - return 0, ErrorStopped + return nil, ErrorStopped } } diff --git a/test_output.go b/test_output.go index e76fbb6..0452d42 100644 --- a/test_output.go +++ b/test_output.go @@ -1,6 +1,6 @@ package main -type writeCallback func(data []byte) +type writeCallback func(*Message) // TestOutput used in testing to intercept any output into callback type TestOutput struct { @@ -8,19 +8,20 @@ type TestOutput struct { } // NewTestOutput constructor for TestOutput, accepts callback which get called on each incoming Write -func NewTestOutput(cb writeCallback) (i *TestOutput) { - i = new(TestOutput) +func NewTestOutput(cb writeCallback) PluginWriter { + i := new(TestOutput) i.cb = cb - return + return i } -func (i *TestOutput) Write(data []byte) (int, error) { - i.cb(data) +// PluginWrite write message to this plugin +func (i *TestOutput) PluginWrite(msg *Message) (int, error) { + i.cb(msg) - return len(data), nil + return len(msg.Data) + len(msg.Meta), nil } func (i *TestOutput) String() string { - return "Test Input" + return "Test Output" }