From ae230cf46b97369e01825f14825c146a97c56b10 Mon Sep 17 00:00:00 2001 From: undostres321 Date: Fri, 3 Oct 2025 01:00:11 +0200 Subject: [PATCH] Add ability to send application data with WebSocket ping frames Closes #4593 --- examples/websocket.js | 8 +++ .../k6/experimental/websockets/websockets.go | 41 ++++++++++++- .../websockets/websockets_test.go | 56 +++++++++++++++++ internal/js/modules/k6/ws/ws.go | 37 +++++++++++- internal/js/modules/k6/ws/ws_test.go | 60 +++++++++++++++++++ 5 files changed, 196 insertions(+), 6 deletions(-) diff --git a/examples/websocket.js b/examples/websocket.js index ecca709b9..01dcd1545 100644 --- a/examples/websocket.js +++ b/examples/websocket.js @@ -10,6 +10,14 @@ export default function () { console.log('connected'); socket.send(Date.now()); + // Send a regular ping + socket.ping(); + console.log("Sent a ping without application data"); + + // Send a ping with application data + socket.ping("application-data"); + console.log("Sent a ping with application data"); + socket.setInterval(function timeout() { socket.ping(); console.log("Pinging every 1sec (setInterval test)"); diff --git a/internal/js/modules/k6/experimental/websockets/websockets.go b/internal/js/modules/k6/experimental/websockets/websockets.go index 4baa6c372..33fdae947 100644 --- a/internal/js/modules/k6/experimental/websockets/websockets.go +++ b/internal/js/modules/k6/experimental/websockets/websockets.go @@ -10,6 +10,7 @@ import ( "net/http" "net/url" "strconv" + "strings" "sync" "time" @@ -678,14 +679,27 @@ func isArrayBufferView(rt *sobek.Runtime, v sobek.Value) (bool, error) { } // Ping sends a ping message over the websocket. -func (w *webSocket) ping() { +// It can optionally include application data as per RFC 6455 section 5.5.2. +func (w *webSocket) ping(args ...sobek.Value) { w.assertStateOpen() pingID := strconv.Itoa(w.sendPings.counter) + var data []byte + if len(args) > 0 && !sobek.IsUndefined(args[0]) && !sobek.IsNull(args[0]) { + if args[0].String() != "" { + // Use a special delimiter "|" to separate ping ID from application data + data = []byte(pingID + "|" + args[0].String()) + } else { + data = []byte(pingID) + } + } else { + data = []byte(pingID) + } + w.writeQueueCh <- message{ mtype: websocket.PingMessage, - data: []byte(pingID), + data: data, t: time.Now(), } @@ -693,9 +707,30 @@ func (w *webSocket) ping() { w.sendPings.counter++ } -func (w *webSocket) trackPong(pingID string) { +func (w *webSocket) trackPong(pongData string) { pongTimestamp := time.Now() + var pingID string + if parts := strings.SplitN(pongData, "|", 2); len(parts) > 0 { + pingID = parts[0] + } else { + for i, c := range pongData { + if c < '0' || c > '9' { + // Found non-numeric character, extract the ID + pingID = pongData[:i] + break + } + if i == len(pongData)-1 { + pingID = pongData + } + } + } + + if pingID == "" { + w.vu.State().Logger.Warnf("received pong with invalid ping ID format") + return + } + pingTimestamp, ok := w.sendPings.timestamps[pingID] if !ok { // We received a pong for a ping we didn't send; ignore diff --git a/internal/js/modules/k6/experimental/websockets/websockets_test.go b/internal/js/modules/k6/experimental/websockets/websockets_test.go index 0772da11d..761cbbd86 100644 --- a/internal/js/modules/k6/experimental/websockets/websockets_test.go +++ b/internal/js/modules/k6/experimental/websockets/websockets_test.go @@ -1546,3 +1546,59 @@ func TestReadyStateSwitch(t *testing.T) { logs := hook.Drain() require.Len(t, logs, 0) } + +func TestSessionPingWithApplicationData(t *testing.T) { + t.Parallel() + tb := httpmultibin.NewHTTPMultiBin(t) + sr := tb.Replacer.Replace + + ts := newTestState(t) + + _, err := ts.runtime.RunOnEventLoop(sr(` + var ws = new WebSocket("WSBIN_URL/ws-echo") + var applicationData = "hello-ping-data"; + ws.onopen = () => { + ws.ping(applicationData) + } + + ws.onpong = () => { + call("from onpong") + ws.close() + } + ws.onerror = (e) => { throw JSON.stringify(e) } + `)) + + require.NoError(t, err) + + samplesBuf := metrics.GetBufferedSamples(ts.samples) + assertSessionMetricsEmitted(t, samplesBuf, "", sr("WSBIN_URL/ws-echo"), http.StatusSwitchingProtocols, "") + assert.Equal(t, []string{"from onpong"}, ts.callRecorder.Recorded()) +} + +func TestSessionPingWithNumericApplicationData(t *testing.T) { + t.Parallel() + tb := httpmultibin.NewHTTPMultiBin(t) + sr := tb.Replacer.Replace + + ts := newTestState(t) + + _, err := ts.runtime.RunOnEventLoop(sr(` + var ws = new WebSocket("WSBIN_URL/ws-echo") + var applicationData = "123456-numeric-ping-data"; + ws.onopen = () => { + ws.ping(applicationData) + } + + ws.onpong = () => { + call("from onpong") + ws.close() + } + ws.onerror = (e) => { throw JSON.stringify(e) } + `)) + + require.NoError(t, err) + + samplesBuf := metrics.GetBufferedSamples(ts.samples) + assertSessionMetricsEmitted(t, samplesBuf, "", sr("WSBIN_URL/ws-echo"), http.StatusSwitchingProtocols, "") + assert.Equal(t, []string{"from onpong"}, ts.callRecorder.Recorded()) +} diff --git a/internal/js/modules/k6/ws/ws.go b/internal/js/modules/k6/ws/ws.go index 1ad5c9ab8..9bb965a6b 100644 --- a/internal/js/modules/k6/ws/ws.go +++ b/internal/js/modules/k6/ws/ws.go @@ -376,10 +376,22 @@ func (s *Socket) SendBinary(message sobek.Value) { } // Ping sends a ping message over the websocket. -func (s *Socket) Ping() { +// It can optionally include application data as per RFC 6455 section 5.5.2. +func (s *Socket) Ping(args ...sobek.Value) { deadline := time.Now().Add(writeWait) pingID := strconv.Itoa(s.pingSendCounter) - data := []byte(pingID) + + var data []byte + if len(args) > 0 && !sobek.IsUndefined(args[0]) && !sobek.IsNull(args[0]) { + if args[0].String() != "" { + // Add the pingID with a delimiter before the application data + data = []byte(pingID + "|" + args[0].String()) + } else { + data = []byte(pingID) + } + } else { + data = []byte(pingID) + } err := s.conn.WriteControl(websocket.PingMessage, data, deadline) if err != nil { @@ -391,9 +403,28 @@ func (s *Socket) Ping() { s.pingSendCounter++ } -func (s *Socket) trackPong(pingID string) { +func (s *Socket) trackPong(pongData string) { pongTimestamp := time.Now() + var pingID string + if parts := strings.SplitN(pongData, "|", 2); len(parts) > 0 { + pingID = parts[0] + } else { + for i, c := range pongData { + if c < '0' || c > '9' { + pingID = pongData[:i] + break + } + if i == len(pongData)-1 { + pingID = pongData + } + } + } + + if pingID == "" { + return + } + if _, ok := s.pingSendTimestamps[pingID]; !ok { // We received a pong for a ping we didn't send; ignore // (this shouldn't happen with a compliant server) diff --git a/internal/js/modules/k6/ws/ws_test.go b/internal/js/modules/k6/ws/ws_test.go index 36a36bf85..d0454d6c1 100644 --- a/internal/js/modules/k6/ws/ws_test.go +++ b/internal/js/modules/k6/ws/ws_test.go @@ -1226,3 +1226,63 @@ func TestWSConnectDisableThrowErrorOption(t *testing.T) { entries := logHook.Drain() assert.Empty(t, entries) } + +func TestSessionPingWithApplicationData(t *testing.T) { + t.Parallel() + tb := httpmultibin.NewHTTPMultiBin(t) + sr := tb.Replacer.Replace + + test := newTestState(t) + _, err := test.VU.Runtime().RunString(sr(` + var pongReceived = false; + var applicationData = "hello-ping-data"; + var res = ws.connect("WSBIN_URL/ws-echo", function(socket){ + socket.on("open", function(data) { + // Send ping with application data + socket.ping(applicationData); + }); + socket.on("pong", function() { + pongReceived = true; + socket.close(); + }); + socket.setTimeout(function (){socket.close();}, 3000); + }); + if (!pongReceived) { + throw new Error ("sent ping with application data but didn't get pong back"); + } + `)) + require.NoError(t, err) + samplesBuf := metrics.GetBufferedSamples(test.samples) + assertSessionMetricsEmitted(t, samplesBuf, "", sr("WSBIN_URL/ws-echo"), statusProtocolSwitch, "") + assertMetricEmittedCount(t, metrics.WSPingName, samplesBuf, sr("WSBIN_URL/ws-echo"), 1) +} + +func TestSessionPingWithNumericApplicationData(t *testing.T) { + t.Parallel() + tb := httpmultibin.NewHTTPMultiBin(t) + sr := tb.Replacer.Replace + + test := newTestState(t) + _, err := test.VU.Runtime().RunString(sr(` + var pongReceived = false; + var applicationData = "123456-numeric-ping-data"; + var res = ws.connect("WSBIN_URL/ws-echo", function(socket){ + socket.on("open", function(data) { + // Send ping with application data that starts with numbers + socket.ping(applicationData); + }); + socket.on("pong", function() { + pongReceived = true; + socket.close(); + }); + socket.setTimeout(function (){socket.close();}, 3000); + }); + if (!pongReceived) { + throw new Error ("sent ping with numeric application data but didn't get pong back"); + } + `)) + require.NoError(t, err) + samplesBuf := metrics.GetBufferedSamples(test.samples) + assertSessionMetricsEmitted(t, samplesBuf, "", sr("WSBIN_URL/ws-echo"), statusProtocolSwitch, "") + assertMetricEmittedCount(t, metrics.WSPingName, samplesBuf, sr("WSBIN_URL/ws-echo"), 1) +}