Merge pull request #769 from buger/arijitAD/refactor_emitter

Refactor emitter.go and fix test accordingly.

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++
 	}
-
 }

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()
 }

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) {

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()
 	}

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
 
 }

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()

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()
 }

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

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()
 }

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)

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()
 }

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
+}

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()
 }

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
+}

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 = ""
 }

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) {

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
+}

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()
 }

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) {

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) {