output-ws (#1136)

* add websocket output

* add tests

* update mod

Author: DimaGolomozy

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

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=

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
+}

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"
+}

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)
 	}

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")