From e44fb6d24007559a1729e5e45039e60db558cf77 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ivan=20Miri=C4=87?= <ivan@loadimpact.com>
Date: Thu, 1 Apr 2021 15:39:05 +0200
Subject: [PATCH] Handle binary WS message using a separate event handler

Resolves https://github.com/k6io/k6/pull/1841#discussion_r605035516
---
 js/modules/k6/ws/ws.go      |  7 +++++--
 js/modules/k6/ws/ws_test.go | 14 +++++++++++++-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/js/modules/k6/ws/ws.go b/js/modules/k6/ws/ws.go
index 492265656fa..8eebb29830f 100644
--- a/js/modules/k6/ws/ws.go
+++ b/js/modules/k6/ws/ws.go
@@ -292,8 +292,11 @@ func (*WS) Connect(ctx context.Context, url string, args ...goja.Value) (*WSHTTP
 				Tags:   socket.sampleTags,
 				Value:  1,
 			})
-			ab := rt.NewArrayBuffer(readData)
-			socket.handleEvent("message", rt.ToValue(string(readData)), rt.ToValue(&ab))
+			socket.handleEvent("message", rt.ToValue(string(readData)))
+			if _, ok := socket.eventHandlers["binaryMessage"]; ok {
+				ab := rt.NewArrayBuffer(readData)
+				socket.handleEvent("binaryMessage", rt.ToValue(&ab))
+			}
 
 		case readErr := <-readErrChan:
 			socket.handleEvent("error", rt.ToValue(readErr))
diff --git a/js/modules/k6/ws/ws_test.go b/js/modules/k6/ws/ws_test.go
index 7c97d4aaaff..37095c655af 100644
--- a/js/modules/k6/ws/ws_test.go
+++ b/js/modules/k6/ws/ws_test.go
@@ -382,15 +382,21 @@ func TestSocketSendBinary(t *testing.T) {
 	assert.NoError(t, err)
 
 	_, err = rt.RunString(sr(`
+		var gotMsg = false, gotBinMsg = false;
 		var res = ws.connect('WSBIN_URL/ws-echo', function(socket){
 			var data = new Uint8Array([104, 101, 108, 108, 111]); // 'hello'
+
 			socket.on('open', function() {
 				socket.sendBinary(data.buffer);
 			})
-			socket.on('message', function (msg, msgBin){
+			socket.on('message', function(msg) {
+				gotMsg = true;
 				if (msg !== 'hello') {
 					throw new Error('received unexpected message: ' + msg);
 				}
+			});
+			socket.on('binaryMessage', function(msgBin) {
+				gotBinMsg = true;
 				let decText = String.fromCharCode.apply(null, new Uint8Array(msgBin));
 				decText = decodeURIComponent(escape(decText));
 				if (decText !== 'hello') {
@@ -399,6 +405,12 @@ func TestSocketSendBinary(t *testing.T) {
 				socket.close()
 			});
 		});
+		if (!gotMsg) {
+			throw new Error("the 'message' handler wasn't called")
+		}
+		if (!gotBinMsg) {
+			throw new Error("the 'binaryMessage' handler wasn't called")
+		}
 		`))
 	assert.NoError(t, err)
 }