gor 引流测试中间件golang代码添加注释

最新推荐文章于 2025-09-05 14:44:44 发布

转载最新推荐文章于 2025-09-05 14:44:44 发布 · 368 阅读

0 ·

CC 4.0 BY-SA版权

原文链接：http://blog.51cto.com/zhongwen/2121000

文章标签：

#golang

本文介绍了一个使用Golang编写的工具，该工具用于在网络流量回放过程中替换敏感的token值，确保了生产环境中token的安全性，并实现了从生产环境到测试环境的平滑过渡。

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 id
reqID := 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'