package main
import (
"bufio"
"bytes"
"encoding/hex"
"fmt"
"github.com/buger/goreplay/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) process(buf)}
}
func process(buf []byte) {
// First byte indicate payload type, possible values:
// 1 - Request
// 2 - Response
// 3 - ReplayedResponse
payloadType := buf[0]
headerSize := bytes.IndexByte(buf, '\n') + 1
header := buf[:headerSize-1]
// Header contains space separated values of: request type, request id, and request start time (or round-trip time for responses)meta := bytes.Split(header, []byte(" "))// For each request you should receive 3 payloads (request, response, replayed response) with same request idreqID := string(meta[1])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"))//取到生產伺服器的token值 if vs != -1 { // If there is GET token param if alias, ok := tokenAliases[string(token)]; ok { //檢查要替換的token值是否存在 // Rewrite original token to alias payload = proto.SetPathParam(payload, []byte("token"), alias)//將生產的token替換成測服的token // Copy modified payload to our buffer buf = append(buf[:headerSize], payload...) } } } // Emitting data back os.Stdout.Write(encode(buf))//重寫請求準備發往測試伺服器case '2': // Original response if _, ok := originalTokens[reqID]; ok { // Token is inside response body secureToken := proto.Body(payload) //取到生產伺服器中返回的token值 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//拿到測試伺服器的token值用來替換掉正服的token值 Debug("Create alias for new token token, was:", string(originalToken), "now:", string(secureToken)) }}
}
func encode(buf []byte) []byte {
dst := make([]byte, len(buf)*2+1)
hex.Encode(dst, buf)
dst[len(dst)-1] = '\n'
return dst
}
func Debug(args ...interface{}) {
fmt.Fprint(os.Stderr, "[DEBUG][TOKEN-MOD] ")
fmt.Fprintln(os.Stderr, args...)
}