diff --git a/emitter.go b/emitter.go index 266615c..8e2221d 100644 --- a/emitter.go +++ b/emitter.go @@ -46,6 +46,15 @@ func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) { e.close() } }() + go func() { + for { + select { + case <-e.quit: + middleware.Close() + return + } + } + }() } else { for _, in := range plugins.Inputs { e.Add(1) diff --git a/input_tcp.go b/input_tcp.go index 5486a75..7a0eaf1 100644 --- a/input_tcp.go +++ b/input_tcp.go @@ -51,7 +51,6 @@ func (i *TCPInput) Read(data []byte) (int, error) { return len(buf), nil } -// Close closes the data channel so that data func (i *TCPInput) Close() error { close(i.stop) return nil diff --git a/middleware.go b/middleware.go index 0fa0031..7611725 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 +}