diff --git a/emitter.go b/emitter.go index cac54d8..8e2221d 100644 --- a/emitter.go +++ b/emitter.go @@ -8,13 +8,25 @@ import ( "time" ) -var wg sync.WaitGroup -var closeOnce sync.Once +type emitter struct { + sync.WaitGroup + quit chan int +} + +// NewEmitter creates and initializes new `emitter` object. +func NewEmitter(quit chan int) *emitter { + return &emitter{ + quit: quit, + } +} // Start initialize loop for sending data from inputs to outputs -func Start(plugins *InOutPlugins, stop chan int) { - if Settings.middleware != "" { - middleware := NewMiddleware(Settings.middleware) +func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) { + e.Add(1) + defer e.Done() + + if middlewareCmd != "" { + middleware := NewMiddleware(middlewareCmd) for _, in := range plugins.Inputs { middleware.ReadFrom(in) @@ -26,31 +38,43 @@ func Start(plugins *InOutPlugins, stop chan int) { middleware.ReadFrom(r) } } - wg.Add(1) + e.Add(1) go func() { - if err := CopyMulty(middleware, plugins.Outputs...); err != nil { + defer e.Done() + if err := CopyMulty(e.quit, middleware, plugins.Outputs...); err != nil { log.Println("Error during copy: ", err) - Close(stop) + e.close() + } + }() + go func() { + for { + select { + case <-e.quit: + middleware.Close() + return + } } }() } else { for _, in := range plugins.Inputs { - wg.Add(1) + e.Add(1) go func(in io.Reader) { - if err := CopyMulty(in, plugins.Outputs...); err != nil { + defer e.Done() + if err := CopyMulty(e.quit, in, plugins.Outputs...); err != nil { log.Println("Error during copy: ", err) - Close(stop) + e.close() } }(in) } for _, out := range plugins.Outputs { if r, ok := out.(io.Reader); ok { - wg.Add(1) + e.Add(1) go func(r io.Reader) { - if err := CopyMulty(r, plugins.Outputs...); err != nil { + defer e.Done() + if err := CopyMulty(e.quit, r, plugins.Outputs...); err != nil { log.Println("Error during copy: ", err) - Close(stop) + e.close() } }(r) } @@ -59,7 +83,7 @@ func Start(plugins *InOutPlugins, stop chan int) { for { select { - case <-stop: + case <-e.quit: finalize(plugins) return case <-time.After(100 * time.Millisecond): @@ -67,17 +91,22 @@ func Start(plugins *InOutPlugins, stop chan int) { } } +func (e *emitter) close() { + select { + case <-e.quit: + default: + close(e.quit) + } +} + // Close closes all the goroutine and waits for it to finish. -func Close(quit chan int) { - closeOnce.Do(func() { - close(quit) - }) - wg.Wait() +func (e *emitter) Close() { + e.close() + e.Wait() } // CopyMulty copies from 1 reader to multiple writers -func CopyMulty(src io.Reader, writers ...io.Writer) error { - defer wg.Done() +func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error { buf := make([]byte, Settings.copyBufferSize) wIndex := 0 modifier := NewHTTPModifier(&Settings.modifierConfig) @@ -89,7 +118,13 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error { var nr int nr, err := src.Read(buf) - if err == io.EOF { + select { + case <-stop: + return nil + default: + } + + if err == io.EOF || err == ErrorStopped { return nil } if err != nil { @@ -192,5 +227,4 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error { i++ } - } diff --git a/emitter_test.go b/emitter_test.go index 0a40fbd..868f624 100644 --- a/emitter_test.go +++ b/emitter_test.go @@ -21,8 +21,10 @@ func TestEmitter(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 1000; i++ { wg.Add(1) @@ -30,8 +32,7 @@ func TestEmitter(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func TestEmitterFiltered(t *testing.T) { @@ -49,10 +50,13 @@ func TestEmitterFiltered(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) + methods := HTTPMethods{[]byte("GET")} Settings.modifierConfig = HTTPModifierConfig{methods: methods} - go Start(plugins, quit) + emitter := &emitter{quit: quit} + go emitter.Start(plugins, "") wg.Add(2) @@ -77,8 +81,7 @@ func TestEmitterFiltered(t *testing.T) { input.EmitBytes(respb) wg.Wait() - - Close(quit) + emitter.Close() Settings.modifierConfig = HTTPModifierConfig{} } @@ -105,10 +108,12 @@ func TestEmitterRoundRobin(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output1, output2}, } + plugins.All = append(plugins.All, input, output1, output2) Settings.splitOutput = true - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 1000; i++ { wg.Add(1) @@ -116,8 +121,7 @@ func TestEmitterRoundRobin(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() if counter1 == 0 || counter2 == 0 { t.Errorf("Round robin should split traffic equally: %d vs %d", counter1, counter2) @@ -140,8 +144,10 @@ func BenchmarkEmitter(b *testing.B) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) b.ResetTimer() @@ -151,5 +157,5 @@ func BenchmarkEmitter(b *testing.B) { } wg.Wait() - close(quit) + emitter.Close() } diff --git a/gor.go b/gor.go index f81caf0..04c4bdd 100644 --- a/gor.go +++ b/gor.go @@ -86,6 +86,7 @@ func main() { }() } + emitter := NewEmitter(closeCh) c := make(chan os.Signal, 1) signal.Notify(c, os.Interrupt, syscall.SIGTERM) go func() { @@ -103,7 +104,7 @@ func main() { }) } - Start(plugins, closeCh) + emitter.Start(plugins, Settings.middleware) } func finalize(plugins *InOutPlugins) { diff --git a/input_file.go b/input_file.go index c7a20c7..5ea97c1 100644 --- a/input_file.go +++ b/input_file.go @@ -113,7 +113,7 @@ type FileInput struct { func NewFileInput(path string, loop bool) (i *FileInput) { i = new(FileInput) i.data = make(chan []byte, 1000) - i.exit = make(chan bool, 1) + i.exit = make(chan bool) i.path = path i.speedFactor = 1 i.loop = loop @@ -153,9 +153,13 @@ func (i *FileInput) init() (err error) { } func (i *FileInput) Read(data []byte) (int, error) { - buf := <-i.data + var buf []byte + select { + case <-i.exit: + return 0, ErrorStopped + case buf = <-i.data: + } copy(data, buf) - return len(buf), nil } @@ -214,7 +218,13 @@ func (i *FileInput) emit() { lastTime = reader.timestamp } - i.data <- reader.ReadPayload() + // Recheck if we have exited since last check. + select { + case <-i.exit: + return + default: + i.data <- reader.ReadPayload() + } } log.Printf("FileInput: end of file '%s'\n", i.path) @@ -231,8 +241,7 @@ func (i *FileInput) Close() error { defer i.mu.Unlock() i.mu.Lock() - i.exit <- true - + close(i.exit) for _, r := range i.readers { r.Close() } diff --git a/input_file_test.go b/input_file_test.go index e5fe3a0..89cc572 100644 --- a/input_file_test.go +++ b/input_file_test.go @@ -18,7 +18,6 @@ import ( var _ = log.Println func TestInputFileWithGET(t *testing.T) { - input := NewTestInput() rg := NewRequestGenerator([]io.Reader{input}, func() { input.EmitGET() }, 1) readPayloads := [][]byte{} @@ -305,7 +304,6 @@ func (expectedCaptureFile *CaptureFile) PayloadsEqual(other [][]byte) bool { } func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile { - f, err := ioutil.TempFile("", "testmainconf") if err != nil { panic(err) @@ -316,7 +314,6 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile { readPayloads := [][]byte{} output := NewTestOutput(func(data []byte) { readPayloads = append(readPayloads, Duplicate(data)) - requestGenerator.wg.Done() }) @@ -326,23 +323,25 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile { Inputs: requestGenerator.inputs, Outputs: []io.Writer{output, outputFile}, } + for _, input := range requestGenerator.inputs { + plugins.All = append(plugins.All, input) + } + plugins.All = append(plugins.All, output, outputFile) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) requestGenerator.emit() requestGenerator.wg.Wait() time.Sleep(100 * time.Millisecond) - outputFile.Close() - - close(quit) + emitter.Close() return NewExpectedCaptureFile(readPayloads, f) } func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback) (err error) { - quit := make(chan int) wg := new(sync.WaitGroup) @@ -356,9 +355,11 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) wg.Add(count) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) done := make(chan int, 1) go func() { @@ -372,8 +373,7 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback case <-time.After(2 * time.Second): err = errors.New("Timed out") } - close(quit) - + emitter.close() return } diff --git a/input_http.go b/input_http.go index f3a6419..6789330 100644 --- a/input_http.go +++ b/input_http.go @@ -13,6 +13,7 @@ type HTTPInput struct { data chan []byte address string listener net.Listener + stop chan bool // Channel used only to indicate goroutine should shutdown } // NewHTTPInput constructor for HTTPInput. Accepts address with port which he will listen on. @@ -20,6 +21,7 @@ func NewHTTPInput(address string) (i *HTTPInput) { i = new(HTTPInput) i.data = make(chan []byte, 10000) i.address = address + i.stop = make(chan bool) i.listen(address) @@ -27,7 +29,12 @@ func NewHTTPInput(address string) (i *HTTPInput) { } func (i *HTTPInput) Read(data []byte) (int, error) { - buf := <-i.data + var buf []byte + select { + case <-i.stop: + return 0, ErrorStopped + case buf = <-i.data: + } header := payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1) @@ -37,6 +44,11 @@ func (i *HTTPInput) Read(data []byte) (int, error) { return len(buf) + len(header), nil } +func (i *HTTPInput) Close() error { + close(i.stop) + return nil +} + func (i *HTTPInput) handler(w http.ResponseWriter, r *http.Request) { r.URL.Scheme = "http" r.URL.Host = i.listener.Addr().String() diff --git a/input_http_test.go b/input_http_test.go index 297cdb0..a4f8e14 100644 --- a/input_http_test.go +++ b/input_http_test.go @@ -25,8 +25,10 @@ func TestHTTPInput(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) address := strings.Replace(input.listener.Addr().String(), "[::]", "127.0.0.1", -1) @@ -36,8 +38,7 @@ func TestHTTPInput(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func TestInputHTTPLargePayload(t *testing.T) { @@ -61,8 +62,10 @@ func TestInputHTTPLargePayload(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) wg.Add(1) address := strings.Replace(input.listener.Addr().String(), "[::]", "127.0.0.1", -1) @@ -73,5 +76,5 @@ func TestInputHTTPLargePayload(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() } diff --git a/input_raw.go b/input_raw.go index 9fa54d0..1626407 100644 --- a/input_raw.go +++ b/input_raw.go @@ -14,7 +14,7 @@ type RAWInput struct { data chan *raw.TCPMessage address string expire time.Duration - quit chan bool + quit chan bool // Channel used only to indicate goroutine should shutdown engine int realIPHeader []byte trackResponse bool @@ -50,7 +50,13 @@ func NewRAWInput(address string, engine int, trackResponse bool, expire time.Dur } func (i *RAWInput) Read(data []byte) (int, error) { - msg := <-i.data + var msg *raw.TCPMessage + select { + case <-i.quit: + return 0, ErrorStopped + case msg = <-i.data: + } + buf := msg.Bytes() var header []byte diff --git a/input_raw_test.go b/input_raw_test.go index 5df2052..b0dff03 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -45,8 +45,6 @@ func TestRAWInputIPv4(t *testing.T) { var respCounter, reqCounter int64 input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "X-Real-IP", "", "", 0) - defer input.Close() - output := NewTestOutput(func(data []byte) { if data[0] == '1' { body := payloadBody(data) @@ -69,10 +67,12 @@ func TestRAWInputIPv4(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) client := NewHTTPClient("http://"+listener.Addr().String(), &HTTPClientConfig{}) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 100; i++ { // request + response @@ -82,8 +82,7 @@ func TestRAWInputIPv4(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func TestRAWInputNoKeepAlive(t *testing.T) { @@ -109,7 +108,6 @@ func TestRAWInputNoKeepAlive(t *testing.T) { originAddr := listener.Addr().String() input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0) - defer input.Close() output := NewTestOutput(func(data []byte) { wg.Done() @@ -119,10 +117,12 @@ func TestRAWInputNoKeepAlive(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) client := NewHTTPClient("http://"+listener.Addr().String(), &HTTPClientConfig{}) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 100; i++ { // request + response @@ -132,8 +132,7 @@ func TestRAWInputNoKeepAlive(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func TestRAWInputIPv6(t *testing.T) { @@ -157,7 +156,6 @@ func TestRAWInputIPv6(t *testing.T) { var respCounter, reqCounter int64 input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0) - defer input.Close() output := NewTestOutput(func(data []byte) { if data[0] == '1' { @@ -177,10 +175,12 @@ func TestRAWInputIPv6(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) client := NewHTTPClient("http://"+listener.Addr().String(), &HTTPClientConfig{}) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 100; i++ { // request + response @@ -190,7 +190,7 @@ func TestRAWInputIPv6(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() } func TestInputRAW100Expect(t *testing.T) { @@ -210,7 +210,6 @@ func TestInputRAW100Expect(t *testing.T) { originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "", "", "", 0) - defer input.Close() // We will use it to get content of raw HTTP request testOutput := NewTestOutput(func(data []byte) { @@ -244,8 +243,10 @@ func TestInputRAW100Expect(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{testOutput, httpOutput}, } + plugins.All = append(plugins.All, input, testOutput, httpOutput) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) // Origin + Response/Request Test Output + Request Http Output wg.Add(4) @@ -256,7 +257,7 @@ func TestInputRAW100Expect(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() } func TestInputRAWChunkedEncoding(t *testing.T) { @@ -275,7 +276,6 @@ func TestInputRAWChunkedEncoding(t *testing.T) { originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "", "", "", 0) - defer input.Close() replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { defer r.Body.Close() @@ -296,9 +296,10 @@ func TestInputRAWChunkedEncoding(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{httpOutput}, } + plugins.All = append(plugins.All, input, httpOutput) - go Start(plugins, quit) - + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) wg.Add(2) curl := exec.Command("curl", "http://"+originAddr, "--header", "Transfer-Encoding: chunked", "--header", "Expect:", "--data-binary", "@README.md") @@ -308,8 +309,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func TestInputRAWLargePayload(t *testing.T) { @@ -341,7 +341,6 @@ func TestInputRAWLargePayload(t *testing.T) { originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0) - defer input.Close() replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { body, _ := ioutil.ReadAll(req.Body) @@ -364,8 +363,10 @@ func TestInputRAWLargePayload(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{httpOutput}, } + plugins.All = append(plugins.All, input, httpOutput) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) wg.Add(2) curl := exec.Command("curl", "http://"+originAddr, "--header", "Transfer-Encoding: chunked", "--header", "Expect:", "--data-binary", "@/tmp/large") @@ -375,7 +376,7 @@ func TestInputRAWLargePayload(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() } func BenchmarkRAWInput(b *testing.B) { @@ -394,7 +395,6 @@ func BenchmarkRAWInput(b *testing.B) { upstreamAddr := strings.Replace(upstream.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "", "", 0) - defer input.Close() output := NewTestOutput(func(data []byte) { if data[0] == '1' { @@ -410,8 +410,10 @@ func BenchmarkRAWInput(b *testing.B) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output, httpOutput}, } + plugins.All = append(plugins.All, input, output, httpOutput) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) emitted := 0 fileContent, _ := ioutil.ReadFile("LICENSE.txt") @@ -442,5 +444,5 @@ func BenchmarkRAWInput(b *testing.B) { time.Sleep(400 * time.Millisecond) log.Println("Emitted ", emitted, ", Captured ", reqCounter, "requests and ", respCounter, " responses", "and replayed", replayCounter) - close(quit) + emitter.Close() } diff --git a/input_tcp.go b/input_tcp.go index 8230bad..7a0eaf1 100644 --- a/input_tcp.go +++ b/input_tcp.go @@ -17,6 +17,7 @@ type TCPInput struct { listener net.Listener address string config *TCPInputConfig + stop chan bool // Channel used only to indicate goroutine should shutdown } type TCPInputConfig struct { @@ -31,6 +32,7 @@ func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) { i.data = make(chan []byte, 1000) i.address = address i.config = config + i.stop = make(chan bool) i.listen(address) @@ -38,12 +40,22 @@ func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) { } func (i *TCPInput) Read(data []byte) (int, error) { - buf := <-i.data + var buf []byte + select { + case <-i.stop: + return 0, ErrorStopped + case buf = <-i.data: + } copy(data, buf) return len(buf), nil } +func (i *TCPInput) Close() error { + close(i.stop) + return nil +} + func (i *TCPInput) listen(address string) { if i.config.secure { cer, err := tls.LoadX509KeyPair(i.config.certificatePath, i.config.keyPath) diff --git a/input_tcp_test.go b/input_tcp_test.go index 43ca7b7..5fd745d 100644 --- a/input_tcp_test.go +++ b/input_tcp_test.go @@ -31,8 +31,10 @@ func TestTCPInput(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) tcpAddr, err := net.ResolveTCPAddr("tcp", input.listener.Addr().String()) @@ -55,8 +57,7 @@ func TestTCPInput(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func genCertificate(template *x509.Certificate) ([]byte, []byte) { @@ -113,8 +114,10 @@ func TestTCPInputSecure(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) conf := &tls.Config{ InsecureSkipVerify: true, @@ -135,6 +138,5 @@ func TestTCPInputSecure(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } diff --git a/limiter.go b/limiter.go index c28ff62..77ddfff 100644 --- a/limiter.go +++ b/limiter.go @@ -77,7 +77,6 @@ func (l *Limiter) Write(data []byte) (n int, err error) { } n, err = l.plugin.(io.Writer).Write(data) - return } @@ -98,3 +97,11 @@ func (l *Limiter) Read(data []byte) (n int, err error) { func (l *Limiter) String() string { return fmt.Sprintf("Limiting %s to: %d (isPercent: %v)", l.plugin, l.limit, l.isPercent) } + +// Close closes the resources. +func (l *Limiter) Close() error { + if fi, ok := l.plugin.(io.ReadCloser); ok { + fi.Close() + } + return nil +} diff --git a/limiter_test.go b/limiter_test.go index 9813fdb..1dafb9a 100644 --- a/limiter_test.go +++ b/limiter_test.go @@ -22,16 +22,17 @@ func TestOutputLimiter(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 100; i++ { input.EmitGET() } wg.Wait() - - close(quit) + emitter.Close() } func TestInputLimiter(t *testing.T) { @@ -48,16 +49,17 @@ func TestInputLimiter(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 100; i++ { input.(*Limiter).plugin.(*TestInput).EmitGET() } wg.Wait() - - close(quit) + emitter.Close() } // Should limit all requests @@ -74,16 +76,17 @@ func TestPercentLimiter1(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 100; i++ { input.EmitGET() } wg.Wait() - - close(quit) + emitter.Close() } // Should not limit at all @@ -101,14 +104,15 @@ func TestPercentLimiter2(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 100; i++ { input.EmitGET() } wg.Wait() - - close(quit) + emitter.Close() } diff --git a/middleware.go b/middleware.go index 0fa0031..daf6257 100644 --- a/middleware.go +++ b/middleware.go @@ -21,12 +21,15 @@ type Middleware struct { Stdin io.Writer Stdout io.Reader + + stop chan bool // Channel used only to indicate goroutine should shutdown } func NewMiddleware(command string) *Middleware { m := new(Middleware) m.command = command m.data = make(chan []byte, 1000) + m.stop = make(chan bool) commands := strings.Split(command, " ") cmd := exec.Command(commands[0], commands[1:]...) @@ -122,19 +125,33 @@ func (m *Middleware) read(from io.Reader) { Debug("[MIDDLEWARE-MASTER] Received:", string(buf)) } - m.data <- buf + select { + case <-m.stop: + return + case m.data <- buf: + } } return } func (m *Middleware) Read(data []byte) (int, error) { - buf := <-m.data - copy(data, buf) + var buf []byte + select { + case <-m.stop: + return 0, ErrorStopped + case buf = <-m.data: + } + copy(data, buf) return len(buf), nil } func (m *Middleware) String() string { return fmt.Sprintf("Modifying traffic using '%s' command", m.command) } + +func (m *Middleware) Close() error { + close(m.stop) + return nil +} diff --git a/middleware_test.go b/middleware_test.go index 02c94a0..3830eeb 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -119,7 +119,6 @@ func TestEchoMiddleware(t *testing.T) { // Catch traffic from one service fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "", "", "", 0) - defer input.Close() // And redirect to another output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: false}) @@ -128,9 +127,11 @@ func TestEchoMiddleware(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) // Start Gor - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) // Wait till middleware initialization time.Sleep(100 * time.Millisecond) @@ -148,7 +149,7 @@ func TestEchoMiddleware(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() time.Sleep(200 * time.Millisecond) Settings.middleware = "" @@ -183,7 +184,6 @@ func TestTokenMiddleware(t *testing.T) { fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1) // Catch traffic from one service input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "", "", "", 0) - defer input.Close() // And redirect to another output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true}) @@ -192,9 +192,11 @@ func TestTokenMiddleware(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) // Start Gor - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) // Wait for middleware to initialize // Give go compiller time to build programm @@ -219,7 +221,7 @@ func TestTokenMiddleware(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() time.Sleep(100 * time.Millisecond) Settings.middleware = "" } diff --git a/output_file_test.go b/output_file_test.go index 10f831e..2cfd903 100644 --- a/output_file_test.go +++ b/output_file_test.go @@ -24,8 +24,10 @@ func TestFileOutput(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 100; i++ { wg.Add(2) @@ -34,10 +36,7 @@ func TestFileOutput(t *testing.T) { } time.Sleep(100 * time.Millisecond) output.flush() - - close(quit) - - quit = make(chan int) + emitter.Close() var counter int64 input2 := NewFileInput("/tmp/test_requests.gor", false) @@ -50,11 +49,14 @@ func TestFileOutput(t *testing.T) { Inputs: []io.Reader{input2}, Outputs: []io.Writer{output2}, } + plugins2.All = append(plugins2.All, input2, output2) - go Start(plugins2, quit) + quit2 := make(chan int) + emitter2 := NewEmitter(quit2) + go emitter2.Start(plugins2, Settings.middleware) wg.Wait() - close(quit) + emitter2.Close() } func TestFileOutputWithNameCleaning(t *testing.T) { diff --git a/output_http.go b/output_http.go index b8f06df..538a913 100644 --- a/output_http.go +++ b/output_http.go @@ -64,6 +64,8 @@ type HTTPOutput struct { queueStats *GorStat elasticSearch *ESPlugin + + stop chan bool // Channel used only to indicate goroutine should shutdown } // NewHTTPOutput constructor for HTTPOutput @@ -73,6 +75,7 @@ func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer { o.address = address o.config = config + o.stop = make(chan bool) if o.config.stats { o.queueStats = NewGorStat("output_http", o.config.statsMs) @@ -124,6 +127,8 @@ func (o *HTTPOutput) startWorker() { for { select { + case <-o.stop: + return case data := <-o.queue: o.sendRequest(client, data) deathCount = 0 @@ -155,7 +160,11 @@ func (o *HTTPOutput) Write(data []byte) (n int, err error) { buf := make([]byte, len(data)) copy(buf, data) - o.queue <- buf + select { + case <-o.stop: + return 0, ErrorStopped + case o.queue <- buf: + } if o.config.stats { o.queueStats.Write(len(o.queue)) @@ -180,7 +189,12 @@ func (o *HTTPOutput) Write(data []byte) (n int, err error) { } func (o *HTTPOutput) Read(data []byte) (int, error) { - resp := <-o.responses + var resp response + select { + case <-o.stop: + return 0, ErrorStopped + case resp = <-o.responses: + } if Settings.debug { Debug("[OUTPUT-HTTP] Received response:", string(resp.payload)) @@ -231,3 +245,9 @@ func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) { func (o *HTTPOutput) String() string { return "HTTP output: " + o.address } + +// Close closes the data channel so that data +func (o *HTTPOutput) Close() error { + close(o.stop) + return nil +} diff --git a/output_http_test.go b/output_http_test.go index 930f9be..e2ecbff 100644 --- a/output_http_test.go +++ b/output_http_test.go @@ -52,8 +52,10 @@ func TestHTTPOutput(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{http_output, output}, } + plugins.All = append(plugins.All, input, output, http_output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 1; i++ { // 2 http-output, 2 - test output request, 2 - test output http response @@ -64,9 +66,7 @@ func TestHTTPOutput(t *testing.T) { } wg.Wait() - - close(quit) - + emitter.Close() Settings.modifierConfig = HTTPModifierConfig{} } @@ -94,16 +94,16 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) wg.Add(1) input.EmitGET() wg.Wait() - - close(quit) - + emitter.Close() Settings.modifierConfig = HTTPModifierConfig{} } @@ -123,8 +123,10 @@ func TestOutputHTTPSSL(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) wg.Add(2) @@ -132,7 +134,7 @@ func TestOutputHTTPSSL(t *testing.T) { input.EmitGET() wg.Wait() - close(quit) + emitter.Close() } func BenchmarkHTTPOutput(b *testing.B) { @@ -152,8 +154,10 @@ func BenchmarkHTTPOutput(b *testing.B) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < b.N; i++ { wg.Add(1) @@ -161,6 +165,5 @@ func BenchmarkHTTPOutput(b *testing.B) { } wg.Wait() - - close(quit) + emitter.Close() } diff --git a/output_tcp_test.go b/output_tcp_test.go index b9d5de0..cf3fd90 100644 --- a/output_tcp_test.go +++ b/output_tcp_test.go @@ -24,8 +24,10 @@ func TestTCPOutput(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 100; i++ { wg.Add(1) @@ -33,8 +35,7 @@ func TestTCPOutput(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func startTCP(cb func([]byte)) net.Listener { @@ -78,8 +79,10 @@ func BenchmarkTCPOutput(b *testing.B) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } + plugins.All = append(plugins.All, input, output) - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) b.ResetTimer() for i := 0; i < b.N; i++ { @@ -88,8 +91,7 @@ func BenchmarkTCPOutput(b *testing.B) { } wg.Wait() - - close(quit) + emitter.Close() } func TestStickyDisable(t *testing.T) { diff --git a/test_input.go b/test_input.go index fcc596c..0dda8f9 100644 --- a/test_input.go +++ b/test_input.go @@ -3,41 +3,51 @@ package main import ( "crypto/rand" "encoding/base64" - "fmt" + "errors" "time" ) +// ErrorStopped is the error returned when the go routines reading the input is stopped. +var ErrorStopped = errors.New("reading stopped") + // TestInput used for testing purpose, it allows emitting requests on demand type TestInput struct { data chan []byte skipHeader bool + stop chan bool // Channel used only to indicate goroutine should shutdown } // NewTestInput constructor for TestInput func NewTestInput() (i *TestInput) { i = new(TestInput) i.data = make(chan []byte, 100) - + i.stop = make(chan bool) return } func (i *TestInput) Read(data []byte) (int, error) { + var buf []byte select { - case buf := <-i.data: - 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 - case <-time.After(10* time.Second): - return 0, fmt.Errorf("timed out waiting for read") + case <-i.stop: + return 0, ErrorStopped + case buf = <-i.data: } + + 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) Close() error { + close(i.stop) + return nil } func (i *TestInput) EmitBytes(data []byte) {