github/goreplay
4
0
Fork 0
mirror of https://github.com/buger/goreplay.git synced 2024-04-21 12:32:02 +00:00
Code Issues Releases Wiki Activity

Add request ID and token modifier middleware

Browse Source
This commit is contained in:
Leonid Bugaev
2015-08-14 20:42:23 +03:00
parent 491541c2df
commit bb319b5fd2
9 changed files with 220 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/echo_modifier.sh
+2 -2
View File
@@ -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")
examples/token_modifier.go
+97
View File
@@ -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...)
}
gor.go
+8 -10
View File
@@ -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()
input_raw.go
+30 -6
View File
@@ -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)
middleware.go
+20 -9
View File
@@ -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 {
middleware_test.go
+20 -11
View File
@@ -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 = ""
}
output_http.go
+20 -9
View File
@@ -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 {
raw_socket_listener/listener.go
+1
View File
@@ -198,6 +198,7 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
if !isIncoming && responseRequest != nil {
message.RequestStart = responseRequest.start
message.RequestAck = responseRequest.ack
}
}
raw_socket_listener/tcp_message.go
+22
View File
@@ -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
}