mirror of
https://github.com/buger/goreplay.git
synced 2024-04-21 12:32:02 +00:00
117 lines
3.7 KiB
Go
117 lines
3.7 KiB
Go
/*
|
|
This middleware made for auth system that randomly generate access tokens, which used later for accessing secure content. Since there is no pre-defined token value, naive approach without middleware (or if middleware use only request payloads) will fail, because replayed server have own tokens, not synced with origin. To fix this, our middleware should take in account responses of replayed and origin server, store `originalToken -> replayedToken` aliases and rewrite all requests using this token to use replayed alias. See `middleware_test.go#TestTokenMiddleware` test for examples of using this middleware.
|
|
|
|
How middleware works:
|
|
|
|
Original request +--------------+
|
|
+-------------+----------STDIN---------->+ |
|
|
| Gor input | | Middleware |
|
|
+-------------+----------STDIN---------->+ |
|
|
Original response +------+---+---+
|
|
| ^
|
|
+-------------+ Modified request v |
|
|
| Gor output +<---------STDOUT-----------------+ |
|
|
+-----+-------+ |
|
|
| |
|
|
| Replayed response |
|
|
+------------------STDIN----------------->----+
|
|
*/
|
|
|
|
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"github.com/buger/gor/proto"
|
|
"os"
|
|
)
|
|
|
|
// 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': // Request
|
|
if bytes.Equal(proto.Path(payload), []byte("/token")) {
|
|
originalTokens[reqID] = []byte{}
|
|
Debug("Found token request:", reqID)
|
|
} else {
|
|
token, vs, _ := proto.PathParam(payload, []byte("token"))
|
|
|
|
if vs != -1 { // If there is GET token param
|
|
if alias, ok := tokenAliases[string(token)]; 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...)
|
|
}
|