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.

7 files changed, 378 insertions, 24 deletions

diff --git a/_www/index.js b/_www/index.js
index fd0f4d7..1924c06 100644
--- a/_www/index.js
+++ b/_www/index.js
@@ -141,7 +141,23 @@ function renderTarget(targetID) {
 		w += "</div>"
 	}
 
-	w += "<div class='HttpTargets'>"
+	w += `
+		<div id="${targetID}.HttpTargets" class="HttpTargets"></div>
+		<div id="${targetID}.WebSocketTargets" class="WebSocketTargets"></div>
+	`
+
+	document.getElementById("main-content").innerHTML = w
+
+	renderHttpTargets(target)
+	renderWebSocketTargets(target)
+}
+
+function renderHttpTargets(target) {
+	let w = "";
+
+	if (!target.HttpTargets) {
+		return;
+	}
 
 	for (let x = 0; x < target.HttpTargets.length; x++) {
 		let http = target.HttpTargets[x]
@@ -179,9 +195,8 @@ function renderTarget(targetID) {
 			</div>
 		`
 	}
-	w += "</div>"
 
-	document.getElementById("main-content").innerHTML = w
+	document.getElementById(`${target.ID}.HttpTargets`).innerHTML = w
 
 	for (let x = 0; x < target.HttpTargets.length; x++) {
 		let http = target.HttpTargets[x]
@@ -200,6 +215,54 @@ function renderTarget(targetID) {
 	}
 }
 
+function renderWebSocketTargets(target) {
+	let w = "";
+
+	if (!target.WebSocketTargets) {
+		return;
+	}
+
+	for (let x = 0; x < target.WebSocketTargets.length; x++) {
+		let wst = target.WebSocketTargets[x]
+
+		w += `
+			<div id="${wst.ID}" class="WebSocketTarget">
+				<h3>
+					${wst.Name}
+					<span class="WebSocketTargetActions">
+						<button onclick="runWebSocket('${target.ID}', '${wst.ID}')">
+							Run
+						</button>
+					</span>
+				</h3>
+
+				<div id="${wst.ID}_headers" class="headers"></div>
+
+				<h4>Parameters</h4>
+				<div id="${wst.ID}_params" class="params"></div>
+
+				<h4>Run response</h4>
+				<pre id="${wst.ID}_response" class="response mono"></pre>
+			</div>
+		`
+	}
+
+	document.getElementById(`${target.ID}.WebSocketTargets`).innerHTML = w
+
+	for (let x = 0; x < target.WebSocketTargets.length; x++) {
+		let wst = target.WebSocketTargets[x]
+
+		if (wst.Headers && Object.keys(wst.Headers).length > 0) {
+			renderHttpTargetHeaders(target, wst)
+		}
+
+		if (wst.Params && Object.keys(wst.Params).length > 0) {
+			renderHttpTargetParams(target, wst)
+		}
+	}
+}
+
+
 function renderHttpTargetHeaders(target, http) {
 	let w = ""
 	for (const k in http.Headers) {
@@ -262,7 +325,7 @@ async function run(targetID, httpTargetID) {
 	req.Target = _targets[targetID]
 	req.HttpTarget = getHttpTargetByID(req.Target, httpTargetID)
 
-	let fres = await fetch("/_trunks/api/target/run", {
+	let fres = await fetch("/_trunks/api/target/run/http", {
 		method: "POST",
 		headers: {
 			"Content-Type": "application/json",
@@ -280,6 +343,29 @@ async function run(targetID, httpTargetID) {
 	elResponse.innerHTML = JSON.stringify(res, null, 2)
 }
 
+async function runWebSocket(targetID, wstID) {
+	let req = {}
+	req.Target = _targets[targetID]
+	req.WebSocketTarget = getWebSocketTargetByID(req.Target, wstID)
+
+	let fres = await fetch("/_trunks/api/target/run/websocket", {
+		method: "POST",
+		headers: {
+			"Content-Type": "application/json",
+		},
+		body: JSON.stringify(req),
+	})
+
+	let res = await fres.json()
+	if (res.code != 200) {
+		notifError(res.message)
+		return
+	}
+
+	let elResponse = document.getElementById(wstID + "_response")
+	elResponse.innerHTML = JSON.stringify(res, null, 2)
+}
+
 async function attack(targetID, httpTargetID) {
 	let target = _targets[targetID]
 	let httpTarget = getHttpTargetByID(target, httpTargetID)
@@ -410,6 +496,15 @@ function getHttpTargetByID(target, id) {
 	return null
 }
 
+function getWebSocketTargetByID(target, id) {
+	for (let x = 0; x < target.WebSocketTargets.length; x++) {
+		if (id == target.WebSocketTargets[x].ID) {
+			return target.WebSocketTargets[x]
+		}
+	}
+	return null
+}
+
 function onChangeHttpHeader(targetID, httpTargetID, key, val) {
 	let target = _targets[targetID]
 	let httpTarget = getHttpTargetByID(target, httpTargetID)
diff --git a/errors.go b/errors.go
index f5f04b4..2ff3194 100644
--- a/errors.go
+++ b/errors.go
@@ -79,3 +79,14 @@ func errInvalidHttpTarget(id string) error {
 	}
 	return res
 }
+
+func errInvalidWebSocketTarget(id string) error {
+	res := &libhttp.EndpointResponse{
+		E: liberrors.E{
+			Code:    http.StatusBadRequest,
+			Message: fmt.Sprintf("invalid or emtpy WebSocketTarget.ID: %q", id),
+			Name:    "ERR_INVALID_WEBSOCKET_TARGET",
+		},
+	}
+	return res
+}
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
+}
diff --git a/run_request.go b/run_request.go
index 27fe5fc..eccb202 100644
--- a/run_request.go
+++ b/run_request.go
@@ -13,10 +13,11 @@ import (
 )
 
 type RunRequest struct {
-	Locker     sync.Mutex
-	Target     *Target
-	HttpTarget *HttpTarget
-	result     *AttackResult
+	Locker          sync.Mutex
+	Target          *Target
+	HttpTarget      *HttpTarget
+	WebSocketTarget *WebSocketTarget
+	result          *AttackResult
 }
 
 func (rr *RunRequest) String() string {
@@ -24,9 +25,10 @@ func (rr *RunRequest) String() string {
 }
 
 //
-// merge the request parameter into original target and HTTP target.
+// mergeHttpTarget merge the request parameter into original target and HTTP
+// target.
 //
-func (rr *RunRequest) merge(env *Environment, origTarget *Target, origHttpTarget *HttpTarget) {
+func (rr *RunRequest) mergeHttpTarget(env *Environment, origTarget *Target, origHttpTarget *HttpTarget) {
 	if rr.Target.Opts.Duration > 0 &&
 		rr.Target.Opts.Duration <= env.MaxAttackDuration {
 		origTarget.Opts.Duration = rr.Target.Opts.Duration
@@ -53,3 +55,37 @@ func (rr *RunRequest) merge(env *Environment, origTarget *Target, origHttpTarget
 	origHttpTarget.Params = rr.HttpTarget.Params
 	rr.HttpTarget = origHttpTarget
 }
+
+//
+// mergeWebSocketTarget merge the request parameter into original target and
+// WebSocket target.
+//
+func (rr *RunRequest) mergeWebSocketTarget(env *Environment,
+	origTarget *Target, origWebSocketTarget *WebSocketTarget,
+) {
+	if rr.Target.Opts.Duration > 0 &&
+		rr.Target.Opts.Duration <= env.MaxAttackDuration {
+		origTarget.Opts.Duration = rr.Target.Opts.Duration
+	}
+
+	if rr.Target.Opts.RatePerSecond > 0 &&
+		rr.Target.Opts.RatePerSecond <= env.MaxAttackRate {
+		origTarget.Opts.RatePerSecond = rr.Target.Opts.RatePerSecond
+		origTarget.Opts.ratePerSecond = vegeta.Rate{
+			Freq: rr.Target.Opts.RatePerSecond,
+			Per:  time.Second,
+		}
+	}
+
+	if rr.Target.Opts.Timeout > 0 &&
+		rr.Target.Opts.Timeout <= DefaultAttackTimeout {
+		origTarget.Opts.Timeout = rr.Target.Opts.Timeout
+	}
+
+	origTarget.Vars = rr.Target.Vars
+	rr.Target = origTarget
+
+	origWebSocketTarget.Headers = rr.WebSocketTarget.Headers
+	origWebSocketTarget.Params = rr.WebSocketTarget.Params
+	rr.WebSocketTarget = origWebSocketTarget
+}
diff --git a/target.go b/target.go
index 3d0be3f..a0aca87 100644
--- a/target.go
+++ b/target.go
@@ -55,6 +55,10 @@ func (target *Target) init() (err error) {
 		ht.init()
 	}
 
+	for _, wst := range target.WebSocketTargets {
+		wst.init()
+	}
+
 	return nil
 }
 
diff --git a/trunks.go b/trunks.go
index 7f36e7d..5fcd938 100644
--- a/trunks.go
+++ b/trunks.go
@@ -36,7 +36,8 @@ const (
 	apiEnvironment        = "/_trunks/api/environment"
 	apiTargetAttack       = "/_trunks/api/target/attack"
 	apiTargetAttackResult = "/_trunks/api/target/attack/result"
-	apiTargetRun          = "/_trunks/api/target/run"
+	apiTargetRunHttp      = "/_trunks/api/target/run/http"
+	apiTargetRunWebSocket = "/_trunks/api/target/run/websocket"
 	apiTargets            = "/_trunks/api/targets"
 )
 
@@ -184,6 +185,9 @@ func (trunks *Trunks) registerHttpApis() (err error) {
 		ResponseType: libhttp.ResponseTypeJSON,
 		Call:         trunks.apiTargetAttackCancel,
 	})
+	if err != nil {
+		return err
+	}
 
 	err = trunks.Server.RegisterEndpoint(&libhttp.Endpoint{
 		Method:       libhttp.RequestMethodGet,
@@ -208,10 +212,21 @@ func (trunks *Trunks) registerHttpApis() (err error) {
 
 	err = trunks.Server.RegisterEndpoint(&libhttp.Endpoint{
 		Method:       libhttp.RequestMethodPost,
-		Path:         apiTargetRun,
+		Path:         apiTargetRunHttp,
 		RequestType:  libhttp.RequestTypeJSON,
 		ResponseType: libhttp.ResponseTypeJSON,
-		Call:         trunks.apiTargetRun,
+		Call:         trunks.apiTargetRunHttp,
+	})
+	if err != nil {
+		return err
+	}
+
+	err = trunks.Server.RegisterEndpoint(&libhttp.Endpoint{
+		Method:       libhttp.RequestMethodPost,
+		Path:         apiTargetRunWebSocket,
+		RequestType:  libhttp.RequestTypeJSON,
+		ResponseType: libhttp.ResponseTypeJSON,
+		Call:         trunks.apiTargetRunWebSocket,
 	})
 	if err != nil {
 		return err
@@ -275,7 +290,7 @@ func (trunks *Trunks) apiTargetAttack(epr *libhttp.EndpointRequest) (resbody []b
 		return nil, errAttackNotAllowed()
 	}
 
-	req.merge(trunks.Env, origTarget, origHttpTarget)
+	req.mergeHttpTarget(trunks.Env, origTarget, origHttpTarget)
 
 	req.result, err = newAttackResult(trunks.Env, req)
 	if err != nil {
@@ -365,7 +380,7 @@ func (trunks *Trunks) apiTargetAttackResultDelete(epr *libhttp.EndpointRequest)
 	return json.Marshal(&res)
 }
 
-func (trunks *Trunks) apiTargetRun(epr *libhttp.EndpointRequest) ([]byte, error) {
+func (trunks *Trunks) apiTargetRunHttp(epr *libhttp.EndpointRequest) ([]byte, error) {
 	req := &RunRequest{}
 	err := json.Unmarshal(epr.RequestBody, req)
 	if err != nil {
@@ -380,18 +395,47 @@ func (trunks *Trunks) apiTargetRun(epr *libhttp.EndpointRequest) ([]byte, error)
 		return nil, errInvalidTarget(req.Target.ID)
 	}
 
-	if req.HttpTarget != nil {
-		origHttpTarget := origTarget.getHttpTargetByID(req.HttpTarget.ID)
-		if origHttpTarget == nil {
-			return nil, errInvalidHttpTarget(req.HttpTarget.ID)
-		}
+	if req.HttpTarget == nil {
+		return nil, errInvalidHttpTarget("")
+	}
 
-		req.merge(trunks.Env, origTarget, origHttpTarget)
+	origHttpTarget := origTarget.getHttpTargetByID(req.HttpTarget.ID)
+	if origHttpTarget == nil {
+		return nil, errInvalidHttpTarget(req.HttpTarget.ID)
+	}
 
-		return req.HttpTarget.Run(req)
+	req.mergeHttpTarget(trunks.Env, origTarget, origHttpTarget)
+
+	return req.HttpTarget.Run(req)
+}
+
+func (trunks *Trunks) apiTargetRunWebSocket(epr *libhttp.EndpointRequest) ([]byte, error) {
+	req := &RunRequest{}
+	err := json.Unmarshal(epr.RequestBody, req)
+	if err != nil {
+		return nil, errInternal(err)
+	}
+	if req.Target == nil {
+		return nil, errInvalidTarget("")
+	}
+
+	origTarget := trunks.getTargetByID(req.Target.ID)
+	if origTarget == nil {
+		return nil, errInvalidTarget(req.Target.ID)
 	}
 
-	return nil, errInvalidHttpTarget("")
+	if req.WebSocketTarget == nil {
+		return nil, errInvalidWebSocketTarget("")
+	}
+
+	origWsTarget := origTarget.getWebSocketTargetByID(req.WebSocketTarget.ID)
+	if origWsTarget == nil {
+		return nil, errInvalidWebSocketTarget(req.WebSocketTarget.ID)
+	}
+
+	req.mergeWebSocketTarget(trunks.Env, origTarget, origWsTarget)
+
+	return req.WebSocketTarget.Run(req)
 }
 
 func (trunks *Trunks) apiTargets(epr *libhttp.EndpointRequest) (resbody []byte, err error) {
diff --git a/websocket_target.go b/websocket_target.go
new file mode 100644
index 0000000..4e17ab9
--- /dev/null
+++ b/websocket_target.go
@@ -0,0 +1,53 @@
+// Copyright 2021, Shulhan <ms@kilabit.info>. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package trunks
+
+import "fmt"
+
+//
+// WebSocketRunHandler define a function type that will be called to run the
+// WebSocket target.
+//
+type WebSocketRunHandler func(rr *RunRequest) ([]byte, error)
+
+//
+// WebSocketTarget define the target to test WebSocket service.
+//
+type WebSocketTarget struct {
+	// Name of target, required.
+	Name string
+
+	// ID of target, optional.
+	// If its empty, it will generated by normalized the value of Name.
+	ID string
+
+	Headers KeyValue
+
+	// Params contains additional parameters to be passed when running
+	// the WebSocket handler.
+	// It could be custom HTTP headers, query parameters, or arguments in
+	// the body. Its up to the user on how to use the Params inside the
+	// Run handler.
+	// This field is optional.
+	Params KeyValue
+
+	Run WebSocketRunHandler `json:"-"`
+}
+
+func (wst *WebSocketTarget) init() (err error) {
+	if len(wst.Name) == 0 {
+		return fmt.Errorf("empty WebSocketTarget Name")
+	}
+	if len(wst.ID) == 0 {
+		wst.ID = generateID(wst.Name)
+	}
+	if wst.Headers == nil {
+		wst.Headers = KeyValue{}
+	}
+	if wst.Params == nil {
+		wst.Params = KeyValue{}
+	}
+	return nil
+}