diff --git a/Makefile b/Makefile index fb1e399..614ff1e 100644 --- a/Makefile +++ b/Makefile @@ -1,8 +1,8 @@ -SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_null.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go test_input.go elasticsearch.go http_modifier.go http_modifier_settings.go http_client.go middleware.go protocol.go output_file_settings.go +SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_null.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go test_input.go elasticsearch.go http_modifier.go http_modifier_settings.go http_client.go middleware.go protocol.go output_file_settings.go output_kafka.go SOURCE_PATH = /go/src/github.com/buger/gor/ PORT = 8000 FADDR = :8000 -RUN = docker run -v `pwd`:$(SOURCE_PATH) -p 0.0.0.0:$(PORT):$(PORT) -t -i gor +RUN = docker run -v `pwd`:$(SOURCE_PATH) -p 0.0.0.0:$(PORT):$(PORT) -t -i gor:go BENCHMARK = BenchmarkRAWInput TEST = TestRawListenerBench VERSION = DEV-$(shell date +%s) diff --git a/emitter.go b/emitter.go index f085a2e..e22595a 100644 --- a/emitter.go +++ b/emitter.go @@ -44,12 +44,16 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) { buf := make([]byte, 5*1024*1024) wIndex := 0 modifier := NewHTTPModifier(&Settings.modifierConfig) + filteredRequests := make(map[string]time.Time) + filteredRequestsLastCleanTime := time.Now() for { nr, er := src.Read(buf) if nr > 0 && len(buf) > nr { payload := buf[:nr] + meta := payloadMeta(payload) + requestID := string(meta[1]) _maxN := nr if nr > 500 { @@ -60,23 +64,31 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) { Debug("[EMITTER] input:", string(payload[0:_maxN]), nr, "from:", src) } - if modifier != nil && isRequestPayload(payload) { - headSize := bytes.IndexByte(payload, '\n') + 1 - body := payload[headSize:] - originalBodyLen := len(body) - body = modifier.Rewrite(body) + if modifier != nil { + if isRequestPayload(payload) { + headSize := bytes.IndexByte(payload, '\n') + 1 + body := payload[headSize:] + originalBodyLen := len(body) + body = modifier.Rewrite(body) - // If modifier tells to skip request - if len(body) == 0 { - continue - } + // If modifier tells to skip request + if len(body) == 0 { + filteredRequests[requestID] = time.Now() + continue + } - if originalBodyLen != len(body) { - payload = append(payload[:headSize], body...) - } + if originalBodyLen != len(body) { + payload = append(payload[:headSize], body...) + } - if Settings.debug { - Debug("[EMITTER] Rewrittern input:", len(payload), "First 500 bytes:", string(payload[0:_maxN])) + if Settings.debug { + Debug("[EMITTER] Rewritten input:", len(payload), "First 500 bytes:", string(payload[0:_maxN])) + } + } else { + if _, ok := filteredRequests[requestID]; ok { + delete(filteredRequests, requestID); + continue + } } } @@ -104,5 +116,17 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) { break } } + + // Clean up filtered requests for which we didn't get a response to filter + now := time.Now() + if now.Sub(filteredRequestsLastCleanTime) > 60 * time.Second { + for k, v := range filteredRequests { + if now.Sub(v) > 60 * time.Second { + delete(filteredRequests, k) + } + } + filteredRequestsLastCleanTime = time.Now() + } + return err } diff --git a/emitter_test.go b/emitter_test.go index aec1e30..500cecd 100644 --- a/emitter_test.go +++ b/emitter_test.go @@ -5,6 +5,7 @@ import ( "sync" "sync/atomic" "testing" + "time" ) func TestEmitter(t *testing.T) { @@ -31,6 +32,51 @@ func TestEmitter(t *testing.T) { close(quit) } +func TestEmitterFiltered(t *testing.T) { + wg := new(sync.WaitGroup) + quit := make(chan int) + + input := NewTestInput() + input.skipHeader = true + + output := NewTestOutput(func(data []byte) { + wg.Done() + }) + + Plugins.Inputs = []io.Reader{input} + Plugins.Outputs = []io.Writer{output} + methods := HTTPMethods{[]byte("GET")} + Settings.modifierConfig = HTTPModifierConfig{methods: methods} + + go Start(quit) + + 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")...) + + 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")...) + + input.EmitBytes(reqb) + input.EmitBytes(respb) + + 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")...) + + 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")...) + + input.EmitBytes(reqb) + input.EmitBytes(respb) + + wg.Wait() + + close(quit) +} + func TestEmitterRoundRobin(t *testing.T) { wg := new(sync.WaitGroup) quit := make(chan int) diff --git a/test_input.go b/test_input.go index 0f990e7..a69debf 100644 --- a/test_input.go +++ b/test_input.go @@ -9,6 +9,7 @@ import ( // TestInput used for testing purpose, it allows emitting requests on demand type TestInput struct { data chan []byte + skipHeader bool } // NewTestInput constructor for TestInput @@ -22,13 +23,23 @@ func NewTestInput() (i *TestInput) { func (i *TestInput) Read(data []byte) (int, error) { buf := <-i.data - header := payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1) - copy(data[0:len(header)], header) - copy(data[len(header):], buf) + var header []byte + + if !i.skipHeader { + header = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1) + copy(data[0:len(header)], header) + copy(data[len(header):], buf) + } else { + copy(data, buf) + } return len(buf) + len(header), nil } +func (i *TestInput) EmitBytes(data []byte) { + i.data <- data +} + // EmitGET emits GET request without headers func (i *TestInput) EmitGET() { i.data <- []byte("GET / HTTP/1.1\r\n\r\n")