diff --git a/capture/tcp_message_test.go b/capture/tcp_message_test.go index e5bf9aa..f1f642c 100644 --- a/capture/tcp_message_test.go +++ b/capture/tcp_message_test.go @@ -36,7 +36,7 @@ func buildMessage(p *TCPPacket) *TCPMessage { if p.SrcPort == 1 { isIncoming = true } - m := NewTCPMessage(p.Seq, p.Ack, isIncoming, ProtocolHTTP, "") + m := NewTCPMessage(p.Seq, p.Ack, isIncoming, ProtocolHTTP, p.timestamp) m.AddPacket(p) return m diff --git a/emitter.go b/emitter.go index 87c8a69..396b7b0 100644 --- a/emitter.go +++ b/emitter.go @@ -9,13 +9,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) @@ -27,31 +39,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) } @@ -60,7 +84,7 @@ func Start(plugins *InOutPlugins, stop chan int) { for { select { - case <-stop: + case <-e.quit: finalize(plugins) return case <-time.After(100 * time.Millisecond): @@ -68,17 +92,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) @@ -90,7 +119,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 { @@ -206,5 +241,4 @@ func CopyMulty(src io.Reader, writers ...io.Writer) error { i++ } - } diff --git a/emitter_test.go b/emitter_test.go index 70030b1..b34fb3c 100644 --- a/emitter_test.go +++ b/emitter_test.go @@ -29,8 +29,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) @@ -38,8 +40,7 @@ func TestEmitter(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func TestEmitterFiltered(t *testing.T) { @@ -57,10 +58,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) @@ -85,8 +89,7 @@ func TestEmitterFiltered(t *testing.T) { input.EmitBytes(respb) wg.Wait() - - Close(quit) + emitter.Close() Settings.modifierConfig = HTTPModifierConfig{} } @@ -116,7 +119,8 @@ func TestEmitterSplitRoundRobin(t *testing.T) { 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) @@ -125,7 +129,7 @@ func TestEmitterSplitRoundRobin(t *testing.T) { wg.Wait() - close(quit) + emitter.Close() if counter1 == 0 || counter2 == 0 || counter1 != counter2 { t.Errorf("Round robin should split traffic equally: %d vs %d", counter1, counter2) @@ -156,10 +160,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) @@ -167,8 +173,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) @@ -220,7 +225,8 @@ func TestEmitterSplitSession(t *testing.T) { Settings.splitOutput = true Settings.recognizeTCPSessions = true - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) for i := 0; i < 1000; i++ { // Keep session but randomize ACK @@ -237,14 +243,13 @@ func TestEmitterSplitSession(t *testing.T) { wg1.Wait() wg2.Wait() - close(quit) - if counter1 != 1000 || counter2 != 1000 { t.Errorf("Round robin should split traffic equally: %d vs %d", counter1, counter2) } Settings.splitOutput = false Settings.recognizeTCPSessions = false + emitter.Close() } func BenchmarkEmitter(b *testing.B) { @@ -261,8 +266,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() @@ -272,5 +279,5 @@ func BenchmarkEmitter(b *testing.B) { } wg.Wait() - close(quit) + emitter.Close() } diff --git a/gor.go b/gor.go index 4d3f985..d7caca7 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 c04a436..446b174 100644 --- a/input_file.go +++ b/input_file.go @@ -125,7 +125,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 @@ -187,9 +187,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 } @@ -248,7 +252,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) @@ -265,8 +275,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 c802c25..3f1c60a 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 @@ -64,7 +64,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 a3c1ef9..4a1114f 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -69,10 +69,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 +84,7 @@ func TestRAWInputIPv4(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func TestRAWInputNoKeepAlive(t *testing.T) { @@ -119,10 +120,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 +135,7 @@ func TestRAWInputNoKeepAlive(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func TestRAWInputIPv6(t *testing.T) { @@ -177,10 +179,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 +194,7 @@ func TestRAWInputIPv6(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() } func TestInputRAW100Expect(t *testing.T) { @@ -244,8 +248,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 +262,7 @@ func TestInputRAW100Expect(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() } func TestInputRAWChunkedEncoding(t *testing.T) { @@ -296,9 +302,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 +315,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) { } wg.Wait() - - close(quit) + emitter.Close() } func TestInputRAWLargePayload(t *testing.T) { @@ -364,8 +370,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 +383,7 @@ func TestInputRAWLargePayload(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() } func BenchmarkRAWInput(b *testing.B) { @@ -410,8 +418,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 +452,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 4d6d59a..07eb3c7 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 0d246cb..db94603 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -128,9 +128,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 +150,7 @@ func TestEchoMiddleware(t *testing.T) { } wg.Wait() - close(quit) + emitter.Close() time.Sleep(200 * time.Millisecond) Settings.middleware = "" @@ -192,9 +194,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 +223,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 9ace964..0ecb9d1 100644 --- a/output_http.go +++ b/output_http.go @@ -110,6 +110,8 @@ type HTTPOutput struct { queueStats *GorStat elasticSearch *ESPlugin + + stop chan bool // Channel used only to indicate goroutine should shutdown } // NewHTTPOutput constructor for HTTPOutput @@ -119,6 +121,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) @@ -207,6 +210,8 @@ func (o *HTTPOutput) startWorker() { for { select { + case <-o.stop: + return case data := <-o.queue: o.sendRequest(client, data) case <-time.After(2 * time.Second): @@ -234,7 +239,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)) @@ -259,7 +268,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)) @@ -310,3 +324,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 1125b91..9166708 100644 --- a/output_http_test.go +++ b/output_http_test.go @@ -53,8 +53,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 < 10; i++ { // 2 http-output, 2 - test output request, 2 - test output http response @@ -100,16 +102,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{} } @@ -129,8 +131,10 @@ func TestHTTPOutputSSL(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) @@ -138,7 +142,7 @@ func TestHTTPOutputSSL(t *testing.T) { input.EmitGET() wg.Wait() - close(quit) + emitter.Close() } func TestHTTPOutputSessions(t *testing.T) { @@ -160,7 +164,8 @@ func TestHTTPOutputSessions(t *testing.T) { Inputs: []io.Reader{input}, Outputs: []io.Writer{output}, } - go Start(plugins, quit) + emitter := NewEmitter(quit) + go emitter.Start(plugins, Settings.middleware) uuid1 := []byte("1234567890123456789a0000") uuid2 := []byte("1234567890123456789d0000") @@ -183,7 +188,7 @@ func TestHTTPOutputSessions(t *testing.T) { t.Error("Should have only 2 workers", output.(*HTTPOutput).activeWorkers) } - close(quit) + emitter.Close() Settings.recognizeTCPSessions = false } @@ -205,8 +210,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) @@ -214,6 +221,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 be262a5..3444a6b 100644 --- a/test_input.go +++ b/test_input.go @@ -3,21 +3,26 @@ package main import ( "crypto/rand" "encoding/base64" + "errors" "fmt" "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 } @@ -40,6 +45,11 @@ func (i *TestInput) Read(data []byte) (int, error) { } } +func (i *TestInput) Close() error { + close(i.stop) + return nil +} + func (i *TestInput) EmitBytes(data []byte) { i.data <- data }