From bb319b5fd24bd147d28b1604ca78da04d7091527 Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Fri, 14 Aug 2015 20:42:23 +0300 Subject: [PATCH] Add request ID and token modifier middleware --- examples/echo_modifier.sh | 4 +- examples/token_modifier.go | 97 ++++++++++++++++++++++++++++++ gor.go | 18 +++--- input_raw.go | 36 +++++++++-- middleware.go | 29 ++++++--- middleware_test.go | 31 ++++++---- output_http.go | 29 ++++++--- raw_socket_listener/listener.go | 1 + raw_socket_listener/tcp_message.go | 22 +++++++ 9 files changed, 220 insertions(+), 47 deletions(-) create mode 100644 examples/token_modifier.go diff --git a/examples/echo_modifier.sh b/examples/echo_modifier.sh index 105b4bb..13d2074 100755 --- a/examples/echo_modifier.sh +++ b/examples/echo_modifier.sh @@ -11,10 +11,10 @@ while read line; do >&2 echo "[DEBUG][ECHO] ===================================" case ${header:0:1} in - "3") + "2") >&2 echo "[DEBUG][ECHO] Request type: Original Response" ;; - "2") + "3") >&2 echo "[DEBUG][ECHO] Request type: Replayed Response" ;; "1") diff --git a/examples/token_modifier.go b/examples/token_modifier.go new file mode 100644 index 0000000..7ca37dd --- /dev/null +++ b/examples/token_modifier.go @@ -0,0 +1,97 @@ +package main + +import ( + "os" + "bufio" + "encoding/hex" + "github.com/buger/gor/proto" + "bytes" + "fmt" +) + +// requestID -> originalToken +var originalTokens map[string][]byte + +// originalToken -> replayedToken +var tokenAliases map[string][]byte + +func main() { + originalTokens = make(map[string][]byte) + tokenAliases = make(map[string][]byte) + + scanner := bufio.NewScanner(os.Stdin) + + for scanner.Scan() { + encoded := scanner.Bytes() + buf := make([]byte, len(encoded)/2) + hex.Decode(buf, encoded) + + 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:] + + 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")) + + 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) + } + } + } + + // 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) + + 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)) + } + } +} + +func Debug(args ...interface{}) { + fmt.Fprint(os.Stderr, "[DEBUG][TOKEN-MOD] ") + fmt.Fprintln(os.Stderr, args...) +} \ No newline at end of file diff --git a/gor.go b/gor.go index af9bf5f..4cd8596 100644 --- a/gor.go +++ b/gor.go @@ -8,7 +8,7 @@ import ( "log" "os" "runtime" - "runtime/debug" + _ "runtime/debug" "runtime/pprof" "time" ) @@ -20,20 +20,18 @@ var ( ) func main() { + // // Don't exit on panic + // defer func() { + // if r := recover(); r != nil { + // fmt.Printf("PANIC: pkg: %v %s \n", r, debug.Stack()) + // } + // }() + // If not set via env cariable if len(os.Getenv("GOMAXPROCS")) == 0 { runtime.GOMAXPROCS(runtime.NumCPU() * 2) } - // Don't exit on panic - defer func() { - if r := recover(); r != nil { - if _, ok := r.(error); !ok { - fmt.Printf("PANIC: pkg: %v %s \n", r, debug.Stack()) - } - } - }() - fmt.Println("Version:", VERSION) flag.Parse() diff --git a/input_raw.go b/input_raw.go index 5201f2d..9d786b3 100644 --- a/input_raw.go +++ b/input_raw.go @@ -29,19 +29,43 @@ func NewRAWInput(address string, expire time.Duration, captureResponse bool) (i return } +const ( + RequestPayload = 1 << iota + ResponsePayload + ReplayedResponsePayload +) + +func payloadHeader(payloadType int, uuid []byte) (header []byte) { + header = make([]byte, 43) + header[1] = ' ' + header[len(header)-1] = '\n' + + switch payloadType { + case RequestPayload: + header[0] = '1' + case ResponsePayload: + header[0] = '2' + case ReplayedResponsePayload: + header[0] = '3' + } + + copy(header[2:], uuid) + + return header +} + func (i *RAWInput) Read(data []byte) (int, error) { msg := <-i.data buf := msg.Bytes() if i.captureResponse { - var header []byte - - if msg.IsIncoming { - header = []byte("1\n") - } else { - header = []byte("3\n") + payloadType := RequestPayload + if !msg.IsIncoming { + payloadType = ResponsePayload } + header := payloadHeader(payloadType, msg.UUID()) + copy(data[0:len(header)], header) copy(data[len(header):], buf) diff --git a/middleware.go b/middleware.go index c642cd2..f1d072d 100644 --- a/middleware.go +++ b/middleware.go @@ -9,6 +9,7 @@ import ( "os" "os/exec" "strings" + "sync" ) type Middleware struct { @@ -16,6 +17,8 @@ type Middleware struct { data chan []byte + mu sync.Mutex + Stdin io.Writer Stdout io.Reader } @@ -48,6 +51,7 @@ func NewMiddleware(command string) *Middleware { } func (m *Middleware) ReadFrom(plugin io.Reader) { + Debug("[MIDDLEWARE-MASTER] Starting reading from", plugin) go m.copy(m.Stdin, plugin) } @@ -60,8 +64,11 @@ func (m *Middleware) copy(to io.Writer, from io.Reader) { if nr > 0 && len(buf) > nr { hex.Encode(dst, buf[0:nr]) - to.Write(dst[0 : nr*2]) - to.Write([]byte("\n")) + dst[nr*2] = '\n' + + m.mu.Lock() + to.Write(dst[0 : nr*2+1]) + m.mu.Unlock() if Settings.debug { Debug("[MIDDLEWARE-MASTER] Sending:", string(buf[0:nr]), "From:", from) @@ -71,19 +78,23 @@ func (m *Middleware) copy(to io.Writer, from io.Reader) { } func (m *Middleware) read(from io.Reader) { - buf := make([]byte, 5*1024*1024) - scanner := bufio.NewScanner(from) for scanner.Scan() { bytes := scanner.Bytes() - hex.Decode(buf, bytes) - - if Settings.debug { - Debug("[MIDDLEWARE-MASTER] Received:", string(buf[0:len(bytes)/2])) + buf := make([]byte, len(bytes)/2) + if _, err := hex.Decode(buf, bytes); err != nil { + fmt.Fprintln(os.Stderr, "Failed to decode input payload", err, len(bytes)) } - m.data <- buf[0 : len(bytes)/2] + if Settings.debug { + Debug("[MIDDLEWARE-MASTER] Received:", string(buf)) + } + + // We should accept only request payloads + if buf[0] == '1' { + m.data <- buf + } } if err := scanner.Err(); err != nil { diff --git a/middleware_test.go b/middleware_test.go index e323fd9..c691e58 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -96,23 +96,28 @@ func TestEchoMiddleware(t *testing.T) { wg := new(sync.WaitGroup) from := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Env", "prod") + w.Header().Set("RequestPath", r.URL.Path) wg.Done() })) to := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Env", "test") + w.Header().Set("RequestPath", r.URL.Path) wg.Done() })) quit := make(chan int) + Settings.middleware = "./examples/echo_modifier.sh" + // Catch traffic from one service input := NewRAWInput(from.Listener.Addr().String(), testRawExpire, true) // And redirect to another - output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true}) + output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: false}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} - Settings.middleware = "./examples/echo_modifier.sh" // Start Gor go Start(quit) @@ -125,14 +130,14 @@ func TestEchoMiddleware(t *testing.T) { client := NewHTTPClient(from.URL, &HTTPClientConfig{Debug: false}) // Request should be echoed - client.Get("/") - client.Get("/") + client.Get("/a") + client.Get("/b") wg.Wait() close(quit) - Settings.middleware = "" - time.Sleep(100 * time.Millisecond) + + Settings.middleware = "" } func TestTokenMiddleware(t *testing.T) { @@ -150,7 +155,7 @@ func TestTokenMiddleware(t *testing.T) { } case "/secure": if status != 202 { - // t.Error("Server should receive valid rewritten token") + t.Error("Server should receive valid rewritten token") } } }) @@ -161,26 +166,29 @@ func TestTokenMiddleware(t *testing.T) { input := NewRAWInput(from, testRawExpire, true) // And redirect to another - output := NewHTTPOutput(to, &HTTPOutputConfig{}) + output := NewHTTPOutput(to, &HTTPOutputConfig{Debug: true}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} - Settings.middleware = "./examples/echo_modifier.sh" + Settings.middleware = "go run ./examples/token_modifier.go" // Start Gor go Start(quit) - time.Sleep(time.Millisecond) + // Wait for middleware to initialize + time.Sleep(500 * time.Millisecond) // Should receive 2 requests from original + 2 from replayed wg.Add(4) - client := NewHTTPClient("http://"+from, &HTTPClientConfig{Debug: true}) + client := NewHTTPClient("http://"+from, &HTTPClientConfig{Debug: false}) // Sending traffic to original service resp, _ = client.Get("/token") token = proto.Body(resp) + time.Sleep(50*time.Millisecond) + resp, _ = client.Get("/secure?token=" + string(token)) if !bytes.Equal(proto.Status(resp), []byte("202")) { t.Error("Valid token should return 202:", proto.Status(resp)) @@ -188,5 +196,6 @@ func TestTokenMiddleware(t *testing.T) { wg.Wait() close(quit) + time.Sleep(100 * time.Millisecond) Settings.middleware = "" } diff --git a/output_http.go b/output_http.go index 7be477f..38c423f 100644 --- a/output_http.go +++ b/output_http.go @@ -9,6 +9,11 @@ import ( const initialDynamicWorkers = 10 +type response struct { + payload []byte + uuid []byte +} + // HTTPOutputConfig struct for holding http output configuration type HTTPOutputConfig struct { redirectLimit int @@ -36,7 +41,7 @@ type HTTPOutput struct { address string limit int queue chan []byte - responses chan []byte + responses chan response needWorker chan int @@ -60,7 +65,7 @@ func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer { } o.queue = make(chan []byte, 100) - o.responses = make(chan []byte, 100) + o.responses = make(chan response, 100) o.needWorker = make(chan int, 1) // Initial workers count @@ -154,17 +159,23 @@ func (o *HTTPOutput) Write(data []byte) (n int, err error) { } func (o *HTTPOutput) Read(data []byte) (int, error) { - buf := <-o.responses - header := []byte("2\n") - copy(data[0:2], header) - copy(data[2:], buf) + resp := <-o.responses - return len(buf) + len(header), nil + Debug("[OUTPUT-HTTP] Received response", string(resp.payload)) + + header := payloadHeader(ReplayedResponsePayload, resp.uuid) + copy(data[0:len(header)], header) + copy(data[len(header):], resp.payload) + + return len(resp.payload) + len(header), nil } func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) { + var uuid []byte + if len(Settings.middleware) > 0 { - request = request[2:] + uuid = request[2:42] + request = request[43:] } start := time.Now() @@ -176,7 +187,7 @@ func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) { } if len(Settings.middleware) > 0 { - o.responses <- resp + o.responses <- response{resp, uuid} } if o.elasticSearch != nil { diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index fa208d4..3c7af46 100644 --- a/raw_socket_listener/listener.go +++ b/raw_socket_listener/listener.go @@ -198,6 +198,7 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) { if !isIncoming && responseRequest != nil { message.RequestStart = responseRequest.start + message.RequestAck = responseRequest.ack } } diff --git a/raw_socket_listener/tcp_message.go b/raw_socket_listener/tcp_message.go index d06ad19..14e70d9 100644 --- a/raw_socket_listener/tcp_message.go +++ b/raw_socket_listener/tcp_message.go @@ -4,6 +4,9 @@ import ( "log" "sort" "time" + "crypto/sha1" + "encoding/hex" + "strconv" ) // TCPMessage ensure that all TCP packets for given request is received, and processed in right sequence @@ -16,6 +19,7 @@ type TCPMessage struct { ID string // Message ID Ack uint32 RequestStart int64 + RequestAck uint32 Start int64 IsIncoming bool packets []*TCPPacket @@ -116,3 +120,21 @@ func (t *TCPMessage) AddPacket(packet *TCPPacket) { // Reset message timeout timer t.timer.Reset(*t.expire) } + +func (t *TCPMessage) UUID() []byte { + var key []byte + + if t.IsIncoming { + key = strconv.AppendInt(key, t.Start, 10) + key = strconv.AppendUint(key, uint64(t.Ack), 10) + } else { + key = strconv.AppendInt(key, t.RequestStart, 10) + key = strconv.AppendUint(key, uint64(t.RequestAck), 10) + } + + uuid := make([]byte, 40) + sha := sha1.Sum(key) + hex.Encode(uuid, sha[:20]) + + return uuid +}