diff --git a/emitter.go b/emitter.go index 00ba070..8ad053c 100644 --- a/emitter.go +++ b/emitter.go @@ -32,8 +32,19 @@ func Start(stop chan int) { for { select { case <-stop: + for _, in := range Plugins.Inputs { + if c, ok := in.(io.Closer); ok { + c.Close() + } + } + + for _, out := range Plugins.Outputs { + if c, ok := out.(io.Closer); ok { + c.Close() + } + } return - case <-time.After(1 * time.Second): + case <-time.After(time.Second): } } } diff --git a/input_raw.go b/input_raw.go index 8c096c6..7ce6a30 100644 --- a/input_raw.go +++ b/input_raw.go @@ -13,6 +13,8 @@ type RAWInput struct { data chan *raw.TCPMessage address string expire time.Duration + quit chan bool + listener *raw.Listener } // NewRAWInput constructor for RAWInput. Accepts address with port as argument. @@ -21,6 +23,7 @@ func NewRAWInput(address string, expire time.Duration) (i *RAWInput) { i.data = make(chan *raw.TCPMessage) i.address = address i.expire = expire + i.quit = make(chan bool) go i.listen(address) @@ -56,11 +59,17 @@ func (i *RAWInput) listen(address string) { log.Fatal("input-raw: error while parsing address", err) } - listener := raw.NewListener(host, port, i.expire, true) + i.listener = raw.NewListener(host, port, i.expire, true) for { + select { + case <-i.quit: + return + default: + } + // Receiving TCPMessage object - m := listener.Receive() + m := i.listener.Receive() i.data <- m } @@ -69,3 +78,8 @@ func (i *RAWInput) listen(address string) { func (i *RAWInput) String() string { return "RAW Socket input: " + i.address } + +func (i *RAWInput) Close() { + i.listener.Close() + close(i.quit) +} diff --git a/input_raw_test.go b/input_raw_test.go index 608c06e..4c7c1ad 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -28,8 +28,13 @@ func TestRAWInput(t *testing.T) { var respCounter, reqCounter int64 + defer func(){ + log.Println(reqCounter, respCounter) + }() input := NewRAWInput(originAddr, testRawExpire) + defer input.Close() + output := NewTestOutput(func(data []byte) { if data[0] == '1' { atomic.AddInt64(&reqCounter, 1) @@ -37,8 +42,6 @@ func TestRAWInput(t *testing.T) { atomic.AddInt64(&respCounter, 1) } - log.Println(reqCounter, respCounter) - wg.Done() }) @@ -58,7 +61,6 @@ func TestRAWInput(t *testing.T) { } wg.Wait() - close(quit) } @@ -80,6 +82,7 @@ func TestInputRAW100Expect(t *testing.T) { originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(originAddr, testRawExpire) + defer input.Close() // We will use it to get content of raw HTTP request testOutput := NewTestOutput(func(data []byte) { @@ -142,6 +145,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) { originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(originAddr, testRawExpire) + defer input.Close() replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { defer r.Body.Close() @@ -200,6 +204,7 @@ func TestInputRAWLargePayload(t *testing.T) { originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) input := NewRAWInput(originAddr, testRawExpire) + defer input.Close() replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { req.Body = http.MaxBytesReader(w, req.Body, 1*1024*1024) diff --git a/middleware_test.go b/middleware_test.go index 8458654..c94154d 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -117,6 +117,7 @@ func TestEchoMiddleware(t *testing.T) { // Catch traffic from one service input := NewRAWInput(from.Listener.Addr().String(), testRawExpire) + defer input.Close() // And redirect to another output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: false}) @@ -174,6 +175,7 @@ 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, testRawExpire) + defer input.Close() // And redirect to another output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true}) diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index fc3de75..65b1899 100644 --- a/raw_socket_listener/listener.go +++ b/raw_socket_listener/listener.go @@ -20,6 +20,7 @@ import ( "runtime/debug" "strconv" "time" + "strings" ) // Listener handle traffic capture @@ -50,6 +51,9 @@ type Listener struct { messageExpire time.Duration captureResponse bool + + conn net.PacketConn + quit chan bool } type request struct { @@ -58,30 +62,31 @@ type request struct { } // NewListener creates and initializes new Listener object -func NewListener(addr string, port string, expire time.Duration, captureResponse bool) (rawListener *Listener) { - rawListener = &Listener{captureResponse: captureResponse} +func NewListener(addr string, port string, expire time.Duration, captureResponse bool) (l *Listener) { + l = &Listener{captureResponse: captureResponse} - rawListener.packetsChan = make(chan *TCPPacket, 10000) - rawListener.messagesChan = make(chan *TCPMessage, 10000) - rawListener.messageDelChan = make(chan *TCPMessage, 10000) + l.packetsChan = make(chan *TCPPacket, 10000) + l.messagesChan = make(chan *TCPMessage, 10000) + l.messageDelChan = make(chan *TCPMessage, 10000) + l.quit = make(chan bool) - rawListener.messages = make(map[string]*TCPMessage) - rawListener.ackAliases = make(map[uint32]uint32) - rawListener.seqWithData = make(map[uint32]uint32) - rawListener.respAliases = make(map[uint32]*request) + l.messages = make(map[string]*TCPMessage) + l.ackAliases = make(map[uint32]uint32) + l.seqWithData = make(map[uint32]uint32) + l.respAliases = make(map[uint32]*request) - rawListener.addr = addr + l.addr = addr _port, _ := strconv.Atoi(port) - rawListener.port = uint16(_port) + l.port = uint16(_port) if expire.Nanoseconds() == 0 { expire = 2000 * time.Millisecond } - rawListener.messageExpire = expire + l.messageExpire = expire - go rawListener.listen() - go rawListener.readRAWSocket() + go l.listen() + go l.readRAWSocket() return } @@ -89,6 +94,9 @@ func NewListener(addr string, port string, expire time.Duration, captureResponse func (t *Listener) listen() { for { select { + case <-t.quit: + t.conn.Close() + return // If message ready for deletion it means that its also complete or expired by timeout case message := <-t.messageDelChan: delete(t.ackAliases, message.Ack) @@ -109,21 +117,26 @@ func (t *Listener) listen() { } func (t *Listener) readRAWSocket() { conn, e := net.ListenPacket("ip4:tcp", t.addr) + t.conn = conn if e != nil { log.Fatal(e) } - defer conn.Close() + defer t.conn.Close() for { buf := make([]byte, 64*1024) // 64kb // Note: ReadFrom receive messages without IP header - n, addr, err := conn.ReadFrom(buf) + n, addr, err := t.conn.ReadFrom(buf) if err != nil { - log.Println("Error:", err) - continue + if strings.HasSuffix(err.Error(), "closed network connection") { + return + } else { + log.Println("Raw listener error:", err) + continue + } } if n > 0 { @@ -240,3 +253,9 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) { func (t *Listener) Receive() *TCPMessage { return <-t.messagesChan } + +func (t *Listener) Close() { + close(t.quit) + t.conn.Close() + return +}