all: implement target for WebSocket

One can register WebSocketTarget just like HttpTarget, its have ID,
Name, Headers, and Params.

Unlike HTTP, the WebSocket target only able to execute Run, it does not
have "Attack", yet.

1 files changed, 112 insertions, 1 deletions

diff --git a/example/example.go b/example/example.go
index b2ef46c..851b614 100644
--- a/example/example.go
+++ b/example/example.go
@@ -5,12 +5,16 @@
 package example
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"sync"
 	"time"
 
 	libhttp "github.com/shuLhan/share/lib/http"
+	"github.com/shuLhan/share/lib/mlog"
+	"github.com/shuLhan/share/lib/websocket"
 	vegeta "github.com/tsenart/vegeta/v12/lib"
 
 	"git.sr.ht/~shulhan/trunks"
@@ -21,8 +25,14 @@ const (
 	pathExamplePostForm = "/example/post/form"
 )
 
+const (
+	websocketAddress = "127.0.0.1:28240"
+)
+
 type Example struct {
-	trunks                *trunks.Trunks
+	trunks   *trunks.Trunks
+	wsServer *websocket.Server
+
 	targetExampleGet      vegeta.Target
 	targetExamplePostForm vegeta.Target
 }
@@ -48,6 +58,19 @@ func New() (ex *Example, err error) {
 		return nil, fmt.Errorf("example: New: %w", err)
 	}
 
+	// Create and register endpoint for WebSocket server.
+	wsOpts := &websocket.ServerOptions{
+		Address: websocketAddress,
+	}
+
+	ex.wsServer = websocket.NewServer(wsOpts)
+
+	err = ex.registerWebSocketEndpoints()
+	if err != nil {
+		return nil, fmt.Errorf("example: New: %w", err)
+	}
+
+	// Register targets for testing HTTP and WebSocket endpoints.
 	err = ex.registerTargets()
 	if err != nil {
 		return nil, fmt.Errorf("example: New: %w", err)
@@ -57,10 +80,18 @@ func New() (ex *Example, err error) {
 }
 
 func (ex *Example) Start() (err error) {
+	go func() {
+		err = ex.wsServer.Start()
+		if err != nil {
+			mlog.Errf("example.Start: %s\n", err)
+		}
+	}()
+
 	return ex.trunks.Start()
 }
 
 func (ex *Example) Stop() {
+	ex.wsServer.Stop()
 	ex.trunks.Stop()
 }
 
@@ -90,6 +121,15 @@ func (ex *Example) registerEndpoints() (err error) {
 	return err
 }
 
+func (ex *Example) registerWebSocketEndpoints() (err error) {
+	err = ex.wsServer.RegisterTextHandler(http.MethodGet, pathExampleGet,
+		ex.handleWSExampleGet)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (ex *Example) registerTargets() (err error) {
 	targetHttp := &trunks.Target{
 		Name:    "Example HTTP",
@@ -137,6 +177,24 @@ func (ex *Example) registerTargets() (err error) {
 
 	ex.trunks.RegisterTarget(targetHttp)
 
+	targetWebSocket := &trunks.Target{
+		Name:    "Example WebSocket",
+		BaseUrl: fmt.Sprintf("ws://%s", websocketAddress),
+		Opts:    &trunks.AttackOptions{},
+		Vars: trunks.KeyValue{
+			"WebSocketVar": "hello",
+		},
+		WebSocketTargets: []*trunks.WebSocketTarget{{
+			Name: "Similar to HTTP GET",
+			Params: trunks.KeyValue{
+				"Param1": "123",
+			},
+			Run: ex.runWebSocketGet,
+		}},
+	}
+
+	ex.trunks.RegisterTarget(targetWebSocket)
+
 	return nil
 }
 
@@ -233,3 +291,56 @@ func (ex *Example) attackExamplePostForm(rr *trunks.RunRequest) vegeta.Targeter
 		return nil
 	}
 }
+
+func (ex *Example) handleWSExampleGet(ctx context.Context, req *websocket.Request) (res websocket.Response) {
+	res.ID = req.ID
+	res.Code = http.StatusOK
+	res.Body = req.Body
+	return res
+}
+
+func (ex *Example) runWebSocketGet(rr *trunks.RunRequest) (resbody []byte, err error) {
+	var wg sync.WaitGroup
+
+	wsc := &websocket.Client{
+		Endpoint: "ws://" + websocketAddress,
+		HandleText: func(cl *websocket.Client, frame *websocket.Frame) error {
+			resbody = frame.Payload()
+			wg.Done()
+			return nil
+		},
+	}
+
+	err = wsc.Connect()
+	if err != nil {
+		return nil, err
+	}
+
+	body, err := json.Marshal(rr.WebSocketTarget.Params)
+	if err != nil {
+		return nil, err
+	}
+
+	req := websocket.Request{
+		ID:     uint64(time.Now().UnixNano()),
+		Method: http.MethodGet,
+		Target: pathExampleGet,
+		Body:   string(body),
+	}
+
+	reqtext, err := json.Marshal(&req)
+	if err != nil {
+		return nil, err
+	}
+
+	err = wsc.SendText(reqtext)
+	if err != nil {
+		return nil, err
+	}
+	wg.Add(1)
+	wg.Wait()
+
+	_ = wsc.Close()
+
+	return resbody, err
+}