From 7035a7fb42b9054657f5a99bb479fd4722d8ae7c Mon Sep 17 00:00:00 2001
From: Shulhan <ms@kilabit.info>
Date: Wed, 24 Mar 2021 01:56:39 +0700
Subject: 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.
---
 example/example.go | 113 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 112 insertions(+), 1 deletion(-)

(limited to 'example/example.go')

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
+}
-- 
cgit v1.3-5-g9baa