From 00f76a397e3a021cab0feda8ae8c1c7c5d4ab647 Mon Sep 17 00:00:00 2001
From: Mihail Stoykov <M.Stoikov@gmail.com>
Date: Wed, 24 Apr 2024 18:29:13 +0300
Subject: [PATCH] Warn on not setting binarytype

This is the first part of #35.

But even if #35 is not finished after this warning is an error we can
safely move this extension in k6.
---
 go.mod                        |  2 +-
 websockets/websockets.go      | 23 +++++++++++++++++++----
 websockets/websockets_test.go | 22 ++++++++++++++++++----
 3 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index 3b5b33a..2a85091 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/dop251/goja v0.0.0-20230919151941-fc55792775de
 	github.com/gorilla/websocket v1.5.0
 	github.com/mstoykov/k6-taskqueue-lib v0.1.0
+	github.com/sirupsen/logrus v1.9.3
 	github.com/stretchr/testify v1.8.4
 	go.k6.io/k6 v0.47.1-0.20231012091148-de19a6a1bc53
 	go.uber.org/goleak v1.2.1
@@ -37,7 +38,6 @@ require (
 	github.com/onsi/gomega v1.20.2 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/serenize/snaker v0.0.0-20201027110005-a7ad2135616e // indirect
-	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/spf13/afero v1.1.2 // indirect
 	github.com/tidwall/gjson v1.17.0 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
diff --git a/websockets/websockets.go b/websockets/websockets.go
index 3312dc4..f0a0ef6 100644
--- a/websockets/websockets.go
+++ b/websockets/websockets.go
@@ -83,6 +83,7 @@ type webSocket struct {
 	// fields that should be seen by js only be updated on the event loop
 	readyState     ReadyState
 	bufferedAmount int
+	binaryType     string
 }
 
 type ping struct {
@@ -168,10 +169,17 @@ func defineWebsocket(rt *goja.Runtime, w *webSocket) {
 	// protocol
 	must(rt, w.obj.DefineAccessorProperty(
 		"binaryType", rt.ToValue(func() goja.Value {
-			return rt.ToValue("ArrayBuffer")
-		}), rt.ToValue(func() goja.Value {
-			common.Throw(rt, errors.New("binaryType is not settable in k6 as it doesn't support Blob"))
-			return nil // it never gets to here
+			return rt.ToValue(w.binaryType)
+		}), rt.ToValue(func(s string) error {
+			switch s {
+			case "blob":
+				return fmt.Errorf("blob is currently not supported, only arraybuffer is.")
+			case "arraybuffer":
+				w.binaryType = s
+				return nil
+			default:
+				return fmt.Errorf("unknown binaryType %s, the supported one is arraybuffer.", s)
+			}
 		}), goja.FLAG_FALSE, goja.FLAG_TRUE))
 
 	setOn := func(property string, el *eventListener) {
@@ -382,6 +390,9 @@ func (w *webSocket) loop() {
 	}
 }
 
+const binarytypeWarning = `You have not set a Websocket binaryType to "arraybuffer", but you got a binary response. ` +
+	`This has been done automatically now, but in the future this will not work.`
+
 func (w *webSocket) queueMessage(msg *message) {
 	w.tq.Queue(func() error {
 		if w.readyState != OPEN {
@@ -402,6 +413,10 @@ func (w *webSocket) queueMessage(msg *message) {
 		ev := w.newEvent(events.MESSAGE, msg.t)
 
 		if msg.mtype == websocket.BinaryMessage {
+			if w.binaryType == "" {
+				w.binaryType = "arraybuffer"
+				w.vu.State().Logger.Warn(binarytypeWarning)
+			}
 			// TODO this technically could be BLOB , but we don't support that
 			ab := rt.NewArrayBuffer(msg.data)
 			must(rt, ev.DefineDataProperty("data", rt.ToValue(ab), goja.FLAG_FALSE, goja.FLAG_FALSE, goja.FLAG_TRUE))
diff --git a/websockets/websockets_test.go b/websockets/websockets_test.go
index f819c77..763731f 100644
--- a/websockets/websockets_test.go
+++ b/websockets/websockets_test.go
@@ -11,6 +11,7 @@ import (
 	"testing"
 
 	"github.com/gorilla/websocket"
+	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"gopkg.in/guregu/null.v3"
@@ -18,6 +19,7 @@ import (
 	httpModule "go.k6.io/k6/js/modules/k6/http"
 	"go.k6.io/k6/js/modulestest"
 	"go.k6.io/k6/lib"
+	"go.k6.io/k6/lib/testutils"
 	"go.k6.io/k6/lib/testutils/httpmultibin"
 	"go.k6.io/k6/metrics"
 )
@@ -264,12 +266,21 @@ func TestReadyState(t *testing.T) {
 func TestBinaryState(t *testing.T) {
 	t.Parallel()
 	ts := newTestState(t)
+	logger, hook := testutils.NewLoggerWithHook(t, logrus.WarnLevel)
+	ts.runtime.VU.StateField.Logger = logger
 	_, err := ts.runtime.RunOnEventLoop(ts.tb.Replacer.Replace(`
-		var ws = new WebSocket("WSBIN_URL/ws-echo")
-		ws.addEventListener("open", () => ws.close())
+		var ws = new WebSocket("WSBIN_URL/ws-echo-invalid")
+		ws.addEventListener("open", () => {
+			ws.send(new Uint8Array([164,41]).buffer)
+			ws.send("k6")
+			ws.onmessage = (e) => {
+				ws.close()
+				call(JSON.stringify(e))
+			}
+		})
 
-		if (ws.binaryType != "ArrayBuffer") {
-			throw new Error("Wrong binaryType value, expected ArrayBuffer got "+ ws.binaryType)
+		if (ws.binaryType != "") {
+			throw new Error("Wrong binaryType value, expected empty got "+ ws.binaryType)
 		}
 
 		var thrown = false;
@@ -283,6 +294,9 @@ func TestBinaryState(t *testing.T) {
 		}
 	`))
 	require.NoError(t, err)
+	logs := hook.Drain()
+	require.Len(t, logs, 1)
+	require.Contains(t, logs[0].Message, binarytypeWarning)
 }
 
 func TestExceptionDontPanic(t *testing.T) {