From 70e7d8dd292839c43f4c120c58385ee2ae012d37 Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Fri, 14 Aug 2015 23:07:46 +0300 Subject: [PATCH] Reduce delay of raw input --- examples/token_modifier.go | 136 ++++++++++++++--------------- input_raw_test.go | 2 +- middleware.go | 2 +- middleware_test.go | 2 +- output_http.go | 2 +- raw_socket_listener/listener.go | 1 + raw_socket_listener/tcp_message.go | 77 ++++++++++++++-- 7 files changed, 141 insertions(+), 81 deletions(-) diff --git a/examples/token_modifier.go b/examples/token_modifier.go index 7ca37dd..89ec995 100644 --- a/examples/token_modifier.go +++ b/examples/token_modifier.go @@ -1,12 +1,12 @@ package main import ( - "os" - "bufio" - "encoding/hex" - "github.com/buger/gor/proto" - "bytes" - "fmt" + "bufio" + "bytes" + "encoding/hex" + "fmt" + "github.com/buger/gor/proto" + "os" ) // requestID -> originalToken @@ -16,82 +16,82 @@ var originalTokens map[string][]byte var tokenAliases map[string][]byte func main() { - originalTokens = make(map[string][]byte) - tokenAliases = make(map[string][]byte) + originalTokens = make(map[string][]byte) + tokenAliases = make(map[string][]byte) - scanner := bufio.NewScanner(os.Stdin) + scanner := bufio.NewScanner(os.Stdin) - for scanner.Scan() { - encoded := scanner.Bytes() - buf := make([]byte, len(encoded)/2) - hex.Decode(buf, encoded) + for scanner.Scan() { + encoded := scanner.Bytes() + buf := make([]byte, len(encoded)/2) + hex.Decode(buf, encoded) - go process(buf) - } + go process(buf) + } } func process(buf []byte) { - // First byte indicate payload type, possible values: - // 1 - Request - // 2 - Response - // 3 - ReplayedResponse - payloadType := buf[0] - headerSize := 42 - header := buf[:headerSize] - // For each request you should receive 3 payloads (request, response, replayed response) with same request id - reqID := string(header[2:headerSize]) - payload := buf[headerSize:] + // First byte indicate payload type, possible values: + // 1 - Request + // 2 - Response + // 3 - ReplayedResponse + payloadType := buf[0] + headerSize := 42 + header := buf[:headerSize] + // For each request you should receive 3 payloads (request, response, replayed response) with same request id + reqID := string(header[2:headerSize]) + payload := buf[headerSize:] - Debug("Received payload:", string(buf)) + Debug("Received payload:", string(buf)) - switch payloadType { - case '1': - if bytes.Equal(proto.Path(payload), []byte("/token")) { - originalTokens[reqID] = []byte{} - Debug("Found token request:", reqID) - } else { - tokenVal, vs, _ := proto.PathParam(payload, []byte("token")) + switch payloadType { + case '1': + if bytes.Equal(proto.Path(payload), []byte("/token")) { + originalTokens[reqID] = []byte{} + Debug("Found token request:", reqID) + } else { + tokenVal, vs, _ := proto.PathParam(payload, []byte("token")) - if vs != -1 { // If there is GET token param - if alias, ok := tokenAliases[string(tokenVal)]; ok { - // Rewrite original token to alias - payload = proto.SetPathParam(payload, []byte("token"), alias) + if vs != -1 { // If there is GET token param + if alias, ok := tokenAliases[string(tokenVal)]; ok { + // Rewrite original token to alias + payload = proto.SetPathParam(payload, []byte("token"), alias) - // Copy modified payload to our buffer - copy(buf[headerSize:], payload) - } - } - } + // Copy modified payload to our buffer + copy(buf[headerSize:], payload) + } + } + } - // Re-compute length in case if payload was modified - bufLen := len(header) + len(payload) - // Encoding request and sending it back - dst := make([]byte, bufLen*2+1) - hex.Encode(dst, buf[:bufLen]) - dst[len(dst)-1] = '\n' + // Re-compute length in case if payload was modified + bufLen := len(header) + len(payload) + // Encoding request and sending it back + dst := make([]byte, bufLen*2+1) + hex.Encode(dst, buf[:bufLen]) + dst[len(dst)-1] = '\n' - os.Stdout.Write(dst) + os.Stdout.Write(dst) - return - case '2': // Original response - if _, ok := originalTokens[reqID]; ok { - // Token is inside response body - secureToken := proto.Body(payload) - originalTokens[reqID] = secureToken - Debug("Remember origial token:", string(secureToken)) - } - case '3': // Replayed response - if originalToken, ok := originalTokens[reqID]; ok { - delete(originalTokens, reqID) - secureToken := proto.Body(payload) - tokenAliases[string(originalToken)] = secureToken + return + case '2': // Original response + if _, ok := originalTokens[reqID]; ok { + // Token is inside response body + secureToken := proto.Body(payload) + originalTokens[reqID] = secureToken + Debug("Remember origial token:", string(secureToken)) + } + case '3': // Replayed response + if originalToken, ok := originalTokens[reqID]; ok { + delete(originalTokens, reqID) + secureToken := proto.Body(payload) + tokenAliases[string(originalToken)] = secureToken - Debug("Create alias for new token token, was:", string(originalToken), "now:", string(secureToken)) - } - } + Debug("Create alias for new token token, was:", string(originalToken), "now:", string(secureToken)) + } + } } func Debug(args ...interface{}) { - fmt.Fprint(os.Stderr, "[DEBUG][TOKEN-MOD] ") - fmt.Fprintln(os.Stderr, args...) -} \ No newline at end of file + fmt.Fprint(os.Stderr, "[DEBUG][TOKEN-MOD] ") + fmt.Fprintln(os.Stderr, args...) +} diff --git a/input_raw_test.go b/input_raw_test.go index 6a2b117..3276e7a 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -138,7 +138,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) { }) replayAddr := listener.Addr().String() - httpOutput := NewHTTPOutput(replayAddr, &HTTPOutputConfig{Debug: true}) + httpOutput := NewHTTPOutput(replayAddr, &HTTPOutputConfig{Debug: false}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{httpOutput} diff --git a/middleware.go b/middleware.go index f1d072d..1057173 100644 --- a/middleware.go +++ b/middleware.go @@ -17,7 +17,7 @@ type Middleware struct { data chan []byte - mu sync.Mutex + mu sync.Mutex Stdin io.Writer Stdout io.Reader diff --git a/middleware_test.go b/middleware_test.go index c691e58..f80e330 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -187,7 +187,7 @@ func TestTokenMiddleware(t *testing.T) { resp, _ = client.Get("/token") token = proto.Body(resp) - time.Sleep(50*time.Millisecond) + time.Sleep(10 * time.Millisecond) resp, _ = client.Get("/secure?token=" + string(token)) if !bytes.Equal(proto.Status(resp), []byte("202")) { diff --git a/output_http.go b/output_http.go index 38c423f..c3e0450 100644 --- a/output_http.go +++ b/output_http.go @@ -11,7 +11,7 @@ const initialDynamicWorkers = 10 type response struct { payload []byte - uuid []byte + uuid []byte } // HTTPOutputConfig struct for holding http output configuration diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index 3c7af46..eee05f4 100644 --- a/raw_socket_listener/listener.go +++ b/raw_socket_listener/listener.go @@ -204,6 +204,7 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) { // Handling Expect: 100-continue requests if bytes.Equal(packet.Data[0:4], bPOST) { + // reading last 20 bytes (not counting CRLF): last header value (if no body presented) if bytes.Equal(packet.Data[len(packet.Data)-24:len(packet.Data)-4], bExpect100ContinueCheck) { t.seqWithData[packet.Seq+uint32(len(packet.Data))] = packet.Ack diff --git a/raw_socket_listener/tcp_message.go b/raw_socket_listener/tcp_message.go index 14e70d9..bba9d29 100644 --- a/raw_socket_listener/tcp_message.go +++ b/raw_socket_listener/tcp_message.go @@ -1,12 +1,14 @@ package rawSocket import ( - "log" - "sort" - "time" + "bytes" "crypto/sha1" "encoding/hex" + "github.com/buger/gor/proto" + "log" + "sort" "strconv" + "time" ) // TCPMessage ensure that all TCP packets for given request is received, and processed in right sequence @@ -22,7 +24,8 @@ type TCPMessage struct { RequestAck uint32 Start int64 IsIncoming bool - packets []*TCPPacket + + packets []*TCPPacket timer *time.Timer // Used for expire check @@ -40,9 +43,6 @@ func NewTCPMessage(ID string, delChan chan *TCPMessage, Ack uint32, expire *time msg.packetsChan = make(chan *TCPPacket) msg.delChan = delChan // used for notifying that message completed or expired - // Every time we receive packet we reset this timer - msg.timer = time.AfterFunc(*msg.expire, msg.Timeout) - go msg.listen() return @@ -64,6 +64,10 @@ 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 @@ -117,8 +121,63 @@ func (t *TCPMessage) AddPacket(packet *TCPPacket) { t.packets = append(t.packets, packet) } - // Reset message timeout timer - t.timer.Reset(*t.expire) + if !t.isMultipart() { + log.Println("MESSAGE NOT MULTIPART", string(packet.Data)) + t.Timeout() + } else { + log.Println("MESSAGE MULTIPART", string(packet.Data)) + // If more then 1 packet, wait for more, and set expiration + if len(t.packets) == 1 { + // Every time we receive packet we reset this timer + t.timer = time.AfterFunc(*t.expire, t.Timeout) + } else { + // Reset message timeout timer + t.timer.Reset(*t.expire) + } + } +} + +// isMultipart returns true if message contains from multiple tcp packets +func (t *TCPMessage) isMultipart() bool { + if len(t.packets) > 1 { + return true + } + + payload := t.packets[0].Data + m := payload[:3] + + if t.IsIncoming { + // If one GET, OPTIONS, or HEAD request + if bytes.Equal(m, []byte("GET")) || bytes.Equal(m, []byte("OPT")) || bytes.Equal(m, []byte("HEA")) { + return false + } else { + if length := proto.Header(payload, []byte("Content-Length")); len(length) > 0 { + l, _ := strconv.Atoi(string(length)) + + log.Println("Content-Length", l, "Body length:", len(proto.Body(payload))) + // If content-length equal current body length + if l > 0 && l == len(proto.Body(payload)) { + return false + } + } + } + } else { + if length := proto.Header(payload, []byte("Content-Length")); len(length) > 0 { + if length[0] == '0' { + return false + } + + l, _ := strconv.Atoi(string(length)) + + log.Println("Content-Length", l, "Body length:", len(proto.Body(payload))) + // If content-length equal current body length + if l > 0 && l == len(proto.Body(payload)) { + return false + } + } + } + + return true } func (t *TCPMessage) UUID() []byte {