From 99e6fdfd605f5ded3e399ef52f4f2408b8bd5716 Mon Sep 17 00:00:00 2001 From: Dima Golomozy Date: Tue, 10 Jan 2023 08:10:23 +0200 Subject: [PATCH] output-ws (#1136) * add websocket output * add tests * update mod --- go.mod | 3 +- go.sum | 4 +- output_ws.go | 154 ++++++++++++++++++++++++++++++++++++++++++++++ output_ws_test.go | 64 +++++++++++++++++++ plugins.go | 4 ++ settings.go | 10 +++ 6 files changed, 237 insertions(+), 2 deletions(-) create mode 100644 output_ws.go create mode 100644 output_ws_test.go diff --git a/go.mod b/go.mod index 0f3f626..43eab40 100644 --- a/go.mod +++ b/go.mod @@ -7,6 +7,8 @@ require ( github.com/aws/aws-sdk-go v1.33.2 github.com/coocood/freecache v1.2.0 github.com/google/gopacket v1.1.20-0.20210429153827-3eaba0894325 + github.com/gorilla/websocket v1.4.2 + github.com/klauspost/compress v1.11.13 // indirect github.com/mattbaird/elastigo v0.0.0-20170123220020-2fe47fd29e4b github.com/stretchr/testify v1.7.0 github.com/xdg-go/scram v1.1.1 @@ -36,7 +38,6 @@ require ( github.com/jcmturner/gofork v1.0.0 // indirect github.com/jmespath/go-jmespath v0.3.0 // indirect github.com/json-iterator/go v1.1.12 // indirect - github.com/klauspost/compress v1.9.8 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect github.com/pierrec/lz4 v2.4.1+incompatible // indirect diff --git a/go.sum b/go.sum index b632611..521c6a9 100644 --- a/go.sum +++ b/go.sum @@ -197,6 +197,7 @@ github.com/googleapis/gnostic v0.5.5/go.mod h1:7+EbHbldMins07ALC74bsA81Ovc97Dwqy github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1 h1:EGx4pi6eqNxGaHF6qqu48+N2wcFQ5qg5FXgOdqsJ5d8= github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY= github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So= +github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc= github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/gregjones/httpcache v0.0.0-20180305231024-9cad4c3443a7/go.mod h1:FecbI9+v66THATjSRHfNgh1IVFe/9kFxbXtjV0ctIMA= github.com/hashicorp/go-uuid v1.0.2 h1:cfejS+Tpcp13yd5nYHWDI6qVCny6wyX2Mt5SGur2IGE= @@ -220,8 +221,9 @@ github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7 github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= -github.com/klauspost/compress v1.9.8 h1:VMAMUUOh+gaxKTMk+zqbjsSjsIcUcL/LF4o63i82QyA= github.com/klauspost/compress v1.9.8/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A= +github.com/klauspost/compress v1.11.13 h1:eSvu8Tmq6j2psUJqJrLcWH6K3w5Dwc+qipbaA6eVEN4= +github.com/klauspost/compress v1.11.13/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= github.com/kr/pretty v0.2.0 h1:s5hAObm+yFO5uHYt5dYjxi2rXrsnmRpJx4OYvIWUaQs= github.com/kr/pretty v0.2.0/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= diff --git a/output_ws.go b/output_ws.go new file mode 100644 index 0000000..be7d33b --- /dev/null +++ b/output_ws.go @@ -0,0 +1,154 @@ +package main + +import ( + "context" + "crypto/tls" + "encoding/base64" + "fmt" + "github.com/gorilla/websocket" + "hash/fnv" + "log" + "net/http" + "net/url" + "strings" + "time" +) + +// WebSocketOutput used for sending raw tcp payloads +// Can be used for transferring binary payloads like protocol buffers +type WebSocketOutput struct { + address string + limit int + buf []chan *Message + bufStats *GorStat + config *WebSocketOutputConfig + workerIndex uint32 + headers http.Header + + close bool +} + +// WebSocketOutputConfig WebSocket output configuration +type WebSocketOutputConfig struct { + Sticky bool `json:"output-ws-sticky"` + SkipVerify bool `json:"output-ws-skip-verify"` + Workers int `json:"output-ws-workers"` +} + +// NewWebSocketOutput constructor for WebSocketOutput +// Initialize X workers which hold keep-alive connection +func NewWebSocketOutput(address string, config *WebSocketOutputConfig) PluginWriter { + o := new(WebSocketOutput) + + u, err := url.Parse(address) + if err != nil { + log.Fatal(fmt.Sprintf("[OUTPUT-WS] parse WS output URL error[%q]", err)) + } + + o.config = config + o.headers = http.Header{ + "Authorization": []string{"Basic " + base64.StdEncoding.EncodeToString([]byte(u.User.String()))}, + } + + u.User = nil // must be after creating the headers + o.address = u.String() + + if Settings.OutputWebSocketStats { + o.bufStats = NewGorStat("output_ws", 5000) + } + + // create X buffers and send the buffer index to the worker + o.buf = make([]chan *Message, o.config.Workers) + for i := 0; i < o.config.Workers; i++ { + o.buf[i] = make(chan *Message, 100) + go o.worker(i) + } + + return o +} + +func (o *WebSocketOutput) worker(bufferIndex int) { + retries := 0 + conn, err := o.connect(o.address) + for { + if o.close { + return + } + + if err == nil { + break + } + + Debug(1, fmt.Sprintf("Can't connect to aggregator instance, reconnecting in 1 second. Retries:%d", retries)) + time.Sleep(1 * time.Second) + + conn, err = o.connect(o.address) + retries++ + } + + if retries > 0 { + Debug(2, fmt.Sprintf("Connected to aggregator instance after %d retries", retries)) + } + + defer conn.Close() + + for { + msg := <-o.buf[bufferIndex] + err = conn.WriteMessage(websocket.BinaryMessage, append(msg.Meta, msg.Data...)) + if err != nil { + Debug(2, "INFO: WebSocket output connection closed, reconnecting "+err.Error()) + o.buf[bufferIndex] <- msg + go o.worker(bufferIndex) + break + } + } +} + +func (o *WebSocketOutput) getBufferIndex(msg *Message) int { + if !o.config.Sticky { + o.workerIndex++ + return int(o.workerIndex) % o.config.Workers + } + + hasher := fnv.New32a() + hasher.Write(payloadID(msg.Meta)) + return int(hasher.Sum32()) % o.config.Workers +} + +// PluginWrite writes message to this plugin +func (o *WebSocketOutput) PluginWrite(msg *Message) (n int, err error) { + if !isOriginPayload(msg.Meta) { + return len(msg.Data), nil + } + + bufferIndex := o.getBufferIndex(msg) + o.buf[bufferIndex] <- msg + + if Settings.OutputTCPStats { + o.bufStats.Write(len(o.buf[bufferIndex])) + } + + return len(msg.Data) + len(msg.Meta), nil +} + +func (o *WebSocketOutput) connect(address string) (conn *websocket.Conn, err error) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + d := websocket.DefaultDialer + if strings.HasPrefix(address, "wss://") { + d.TLSClientConfig = &tls.Config{InsecureSkipVerify: o.config.SkipVerify} + } + + conn, _, err = d.DialContext(ctx, address, o.headers) + return +} + +func (o *WebSocketOutput) String() string { + return fmt.Sprintf("WebSocket output %s, limit: %d", o.address, o.limit) +} + +// Close closes the output +func (o *WebSocketOutput) Close() { + o.close = true +} diff --git a/output_ws_test.go b/output_ws_test.go new file mode 100644 index 0000000..ffce5b0 --- /dev/null +++ b/output_ws_test.go @@ -0,0 +1,64 @@ +package main + +import ( + "github.com/gorilla/websocket" + "log" + "net/http" + "sync" + "testing" +) + +func TestWebSocketOutput(t *testing.T) { + wg := new(sync.WaitGroup) + + wsAddr := startWebsocket(func(data []byte) { + wg.Done() + }) + input := NewTestInput() + output := NewWebSocketOutput(wsAddr, &WebSocketOutputConfig{Workers: 1}) + + plugins := &InOutPlugins{ + Inputs: []PluginReader{input}, + Outputs: []PluginWriter{output}, + } + + emitter := NewEmitter() + go emitter.Start(plugins, Settings.Middleware) + + for i := 0; i < 10; i++ { + wg.Add(1) + input.EmitGET() + } + + wg.Wait() + emitter.Close() +} + +func startWebsocket(cb func([]byte)) string { + upgrader := websocket.Upgrader{} + + http.HandleFunc("/test", func(w http.ResponseWriter, r *http.Request) { + c, err := upgrader.Upgrade(w, r, nil) + if err != nil { + log.Print("upgrade:", err) + return + } + + go func(conn *websocket.Conn) { + defer conn.Close() + for { + _, msg, _ := conn.ReadMessage() + cb(msg) + } + }(c) + }) + + go func() { + err := http.ListenAndServe("localhost:8081", nil) + if err != nil { + log.Fatal("Can't start:", err) + } + }() + + return "ws://localhost:8081/test" +} diff --git a/plugins.go b/plugins.go index 72ffff7..be2aab3 100644 --- a/plugins.go +++ b/plugins.go @@ -117,6 +117,10 @@ func NewPlugins() *InOutPlugins { plugins.registerPlugin(NewTCPOutput, options, &Settings.OutputTCPConfig) } + for _, options := range Settings.OutputWebSocket { + plugins.registerPlugin(NewWebSocketOutput, options, &Settings.OutputWebSocketConfig) + } + for _, options := range Settings.InputFile { plugins.registerPlugin(NewFileInput, options, Settings.InputFileLoop, Settings.InputFileReadDepth, Settings.InputFileMaxWait, Settings.InputFileDryRun) } diff --git a/settings.go b/settings.go index e3d08ac..6bb032f 100644 --- a/settings.go +++ b/settings.go @@ -83,6 +83,10 @@ type AppSettings struct { OutputTCPConfig TCPOutputConfig OutputTCPStats bool `json:"output-tcp-stats"` + OutputWebSocket []string `json:"output-ws"` + OutputWebSocketConfig WebSocketOutputConfig + OutputWebSocketStats bool `json:"output-ws-stats"` + InputFile []string `json:"input-file"` InputFileLoop bool `json:"input-file-loop"` InputFileReadDepth int `json:"input-file-read-depth"` @@ -152,6 +156,12 @@ func init() { flag.IntVar(&Settings.OutputTCPConfig.Workers, "output-tcp-workers", 10, "Number of parallel tcp connections, default is 10") flag.BoolVar(&Settings.OutputTCPStats, "output-tcp-stats", false, "Report TCP output queue stats to console every 5 seconds.") + flag.Var(&MultiOption{&Settings.OutputWebSocket}, "output-ws", "Just like output tcp, just with WebSocket. Example: \n\t# Listen for requests on 80 port and forward them to other Gor instance on 28020 port\n\tgor --input-raw :80 --output-ws wss://replay.local:28020/endpoint") + flag.BoolVar(&Settings.OutputWebSocketConfig.SkipVerify, "output-ws-skip-verify", false, "Don't verify hostname on TLS secure connection.") + flag.BoolVar(&Settings.OutputWebSocketConfig.Sticky, "output-ws-sticky", false, "Use Sticky connection. Request/Response with same ID will be sent to the same connection.") + flag.IntVar(&Settings.OutputWebSocketConfig.Workers, "output-ws-workers", 10, "Number of parallel ws connections, default is 10") + flag.BoolVar(&Settings.OutputWebSocketStats, "output-ws-stats", false, "Report WebSocket output queue stats to console every 5 seconds.") + flag.Var(&MultiOption{&Settings.InputFile}, "input-file", "Read requests from file: \n\tgor --input-file ./requests.gor --output-http staging.com") flag.BoolVar(&Settings.InputFileLoop, "input-file-loop", false, "Loop input files, useful for performance testing.") flag.IntVar(&Settings.InputFileReadDepth, "input-file-read-depth", 100, "GoReplay tries to read and cache multiple records, in advance. In parallel it also perform sorting of requests, if they came out of order. Since it needs hold this buffer in memory, bigger values can cause worse performance")