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