diff --git a/Makefile b/Makefile index c051a84..f2bba5c 100644 --- a/Makefile +++ b/Makefile @@ -20,7 +20,7 @@ drace: docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor go test ./... $(ARGS) -v -race -timeout 15s dtest: - docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go test ./... $(ARGS) -v -timeout 10s + docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go test ./... $(ARGS) -v -timeout 60s dcover: docker run -v `pwd`:$(SOURCE_PATH) -t -i --env GORACE="halt_on_error=1" gor go test $(ARGS) -race -v -timeout 15s -coverprofile=coverage.out diff --git a/http_client_test.go b/http_client_test.go index 2500c73..d970381 100644 --- a/http_client_test.go +++ b/http_client_test.go @@ -5,7 +5,7 @@ import ( "crypto/rand" "github.com/buger/gor/proto" "io/ioutil" - "log" + _ "log" "net" "net/http" "net/http/httptest" @@ -171,12 +171,16 @@ func TestHTTPClientServerInstantDisconnect(t *testing.T) { GETPayload := []byte("GET / HTTP/1.1\r\n\r\n") ln, _ := net.Listen("tcp", ":0") + defer ln.Close() go func() { for { - conn, _ := ln.Accept() - conn.Close() + conn, err := ln.Accept() + if err != nil { + break + } + conn.Close() wg.Done() } }() @@ -196,12 +200,13 @@ func TestHTTPClientServerNoKeepAlive(t *testing.T) { GETPayload := []byte("GET / HTTP/1.1\r\n\r\n") ln, _ := net.Listen("tcp", ":0") + defer ln.Close() go func() { for { conn, err := ln.Accept() if err != nil { - // handle error + break } buf := make([]byte, 4096) @@ -370,12 +375,7 @@ func TestHTTPClientErrors(t *testing.T) { ln2, _ := net.Listen("tcp", ":0") go func() { buf := make([]byte, 64*1024) - conn, err := ln2.Accept() - - if err != nil { - log.Println("Error while Accept()", err) - continue - } + conn, _ := ln2.Accept() conn.Read(buf) defer conn.Close() diff --git a/middleware_test.go b/middleware_test.go index bf8203c..aab936e 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -149,10 +149,6 @@ func TestTokenMiddleware(t *testing.T) { }) to := NewFakeSecureService(wg, func(path string, status int, tok []byte) { switch path { - case "/token": - if bytes.Equal(token, tok) { - t.Error("Tokens should not match") - } case "/secure": if status != 202 { t.Error("Server should receive valid rewritten token") @@ -164,6 +160,8 @@ func TestTokenMiddleware(t *testing.T) { quit := make(chan int) + Settings.middleware = "go run ./examples/middleware/token_modifier.go" + // Catch traffic from one service input := NewRAWInput(from, testRawExpire) @@ -172,7 +170,6 @@ func TestTokenMiddleware(t *testing.T) { Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} - Settings.middleware = "go run ./examples/middleware/token_modifier.go" // Start Gor go Start(quit) diff --git a/output_http.go b/output_http.go index cc75192..0f57ea9 100644 --- a/output_http.go +++ b/output_http.go @@ -28,6 +28,8 @@ type HTTPOutputConfig struct { OriginalHost bool Debug bool + + TrackResponses bool } // HTTPOutput plugin manage pool of workers which send request to replayed server @@ -42,6 +44,7 @@ type HTTPOutput struct { address string limit int queue chan []byte + responses chan response needWorker chan int @@ -81,6 +84,10 @@ func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer { o.elasticSearch.Init(o.config.elasticSearch) } + if len(Settings.middleware) > 0 { + o.config.TrackResponses = true + } + go o.workerMaster() return o @@ -187,7 +194,7 @@ func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) { log.Println("Request error:", err) } - if len(Settings.middleware) > 0 { + if o.config.TrackResponses { o.responses <- response{resp, uuid, stop.UnixNano() - start.UnixNano()} } diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index eefb3e0..fc3de75 100644 --- a/raw_socket_listener/listener.go +++ b/raw_socket_listener/listener.go @@ -102,6 +102,7 @@ func (t *Listener) listen() { // We need to use channels to process each packet to avoid data races case packet := <-t.packetsChan: + // log.Println("Received packet:", packet) t.processTCPPacket(packet) } } @@ -219,10 +220,12 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) { } if t.captureResponse && isIncoming { + message.mu.Lock() // If message have multiple packets, delete previous alias if len(message.packets) > 0 { delete(t.respAliases, message.ResponseAck) } + message.mu.Unlock() responseAck := packet.Seq + uint32(len(packet.Data)) t.respAliases[responseAck] = &request{message.Start, message.Ack} diff --git a/raw_socket_listener/tcp_message.go b/raw_socket_listener/tcp_message.go index a0b8f88..2ef8327 100644 --- a/raw_socket_listener/tcp_message.go +++ b/raw_socket_listener/tcp_message.go @@ -9,6 +9,7 @@ import ( "sort" "strconv" "time" + "sync" ) // TCPMessage ensure that all TCP packets for given request is received, and processed in right sequence @@ -36,6 +37,8 @@ type TCPMessage struct { delChan chan *TCPMessage expire *time.Duration + + mu sync.Mutex } // NewTCPMessage pointer created from a Acknowledgment number and a channel of messages readuy to be deleted @@ -44,6 +47,7 @@ func NewTCPMessage(ID string, delChan chan *TCPMessage, Ack uint32, expire *time msg.Start = time.Now().UnixNano() msg.packetsChan = make(chan *TCPPacket) msg.delChan = delChan // used for notifying that message completed or expired + msg.timer = time.NewTimer(0) go msg.listen() @@ -66,16 +70,13 @@ func (t *TCPMessage) listen() { // Timeout notifies message to stop listening, close channel and message ready to be sent func (t *TCPMessage) Timeout() { - if t.timer != nil { - t.timer.Stop() - } - select { // In some cases Timeout can be called multiple times (do not know how yet) // Ensure that we did not close channel 2 times case packet, ok := <-t.packetsChan: if ok { t.AddPacket(packet) + t.Timeout() } else { return } @@ -112,6 +113,9 @@ func (t *TCPMessage) Size() (size int) { // AddPacket to the message and ensure packet uniqueness // TCP allows that packet can be re-send multiple times func (t *TCPMessage) AddPacket(packet *TCPPacket) { + t.mu.Lock() + defer t.mu.Unlock() + packetFound := false for _, pkt := range t.packets { @@ -124,7 +128,13 @@ func (t *TCPMessage) AddPacket(packet *TCPPacket) { if packetFound { log.Println("Received packet with same sequence") } else { - t.packets = append(t.packets, packet) + // Packets not always captured in same Seq order, and sometimes we need to prepend + if len(t.packets) == 0 || packet.Seq > t.packets[len(t.packets)-1].Seq { + t.packets = append(t.packets, packet) + } else { + t.packets = append([]*TCPPacket{packet}, t.packets...) + } + t.End = time.Now().UnixNano() }