From 65475d3a86e5ff7159e0a1f694f0860658cdfcaa Mon Sep 17 00:00:00 2001
From: Stu Tomlinson <stu@nosnilmot.com>
Date: Sat, 13 Jan 2024 16:04:44 +0000
Subject: [PATCH] add websocket ping support

---
 examples/websocket.xml.in  |  8 ++++++++
 include/ts_websocket.hrl   |  3 ++-
 src/lib/websocket.erl      |  8 +++++++-
 src/tsung/ts_websocket.erl | 33 +++++++++++++++++++++++++++------
 tsung-1.0.dtd              |  4 ++--
 5 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/examples/websocket.xml.in b/examples/websocket.xml.in
index 1f01ccaf8..cfb000c54 100644
--- a/examples/websocket.xml.in
+++ b/examples/websocket.xml.in
@@ -31,10 +31,18 @@
             <websocket type="message">{"uid":"%%_uid%%", "data":"data"}</websocket>
         </request>
 
+        <request>
+            <websocket type="ping" ack="bidi_ack">hi</websocket>
+        </request>
+
         <request>
             <websocket type="message" ack="no_ack">{"key":"value"}</websocket>
         </request>
 
+        <for var="i" from="1" to="3" incr="1">
+            <thinktime value="30"/>
+        </for>
+
         <request>
             <websocket type="close"></websocket>
         </request>
diff --git a/include/ts_websocket.hrl b/include/ts_websocket.hrl
index 812b1ca2f..b12673758 100644
--- a/include/ts_websocket.hrl
+++ b/include/ts_websocket.hrl
@@ -44,7 +44,8 @@
 
 -record(websocket_session, {
           status, % status of handshake
-          accept  % Sec-Websocket-Accept header value
+          accept, % Sec-Websocket-Accept header value
+          ping    % ping payload, expected in pong response
          }).
 
 %% opcode of websocket
diff --git a/src/lib/websocket.erl b/src/lib/websocket.erl
index ead151b4c..d084c8fec 100644
--- a/src/lib/websocket.erl
+++ b/src/lib/websocket.erl
@@ -28,7 +28,7 @@
 -author('jzhihui521@gmail.com').
 
 -export([get_handshake/6, check_handshake/2, encode_binary/1, encode_text/1,
-         encode_close/1, encode/2, decode/1]).
+         encode_close/1, encode_ping/1, encode_pong/1, encode/2, decode/1]).
 
 -include("ts_profile.hrl").
 -include("ts_config.hrl").
@@ -113,6 +113,12 @@ encode_close(Reason) ->
     Data = <<StatusCode/binary, Reason/binary>>,
     encode(Data, ?OP_CLOSE).
 
+encode_ping(Data) ->
+    encode(Data, ?OP_PING).
+
+encode_pong(Data) ->
+    encode(Data, ?OP_PONG).
+
 encode(Data, Opcode) ->
     Key = crypto:strong_rand_bytes(4),
     PayloadLen = erlang:size(Data),
diff --git a/src/tsung/ts_websocket.erl b/src/tsung/ts_websocket.erl
index c632609c6..56640eb00 100644
--- a/src/tsung/ts_websocket.erl
+++ b/src/tsung/ts_websocket.erl
@@ -98,7 +98,11 @@ get_message(#websocket_request{type = message, data = Data, frame = Frame},
 get_message(#websocket_request{type = close},
             #state_rcv{session = WebsocketSession})
   when WebsocketSession#websocket_session.status == connected ->
-    {websocket:encode_close(<<"close">>), WebsocketSession}.
+    {websocket:encode_close(<<"close">>), WebsocketSession};
+get_message(#websocket_request{type = ping, data = Data},
+            #state_rcv{session = WebsocketSession})
+  when WebsocketSession#websocket_session.status == connected ->
+    {websocket:encode_ping(list_to_binary(Data)), WebsocketSession#websocket_session{ping = list_to_binary(Data)}}.
 
 %%----------------------------------------------------------------------
 %% Function: parse/2
@@ -137,10 +141,16 @@ parse(Data, State=#state_rcv{acc = [],
 %% normal websocket message
 parse(Data, State=#state_rcv{acc = [], session = WebsocketSession})
   when WebsocketSession#websocket_session.status == connected ->
+    Expected = WebsocketSession#websocket_session.ping,
     case websocket:decode(Data) of
-        {?OP_CLOSE, _Reason, _} ->
-            ?DebugF("receive close from server: ~p~n", [_Reason]),
-            {State#state_rcv{ack_done = true}, [], true};
+        {?OP_PONG, Expected, << >>} ->
+            ?DebugF("received pong: [~p] ~n", [Expected]),
+            ts_mon_cache:add({count, websocket_pong}),
+            {State#state_rcv{ack_done = true, session = WebsocketSession#websocket_session{ping = none}}, [], false};
+        {?OP_PONG, _Unexpected, << >>} ->
+            ts_mon_cache:add({count, error_websocket_pong}),
+            ?DebugF("unexpected pong: [~p] ~n", [_Unexpected]),
+            {State#state_rcv{ack_done = true}, [], false};
         {_Opcode, _Payload, Left} ->
             ?DebugF("receive from server: ~p ~p~n", [_Opcode, _Payload]),
             {State#state_rcv{ack_done = true, acc = Left}, [], false};
@@ -154,8 +164,19 @@ parse(Data, State=#state_rcv{acc = Acc, datasize = DataSize}) ->
     parse(<< Acc/binary, Data/binary >>,
           State#state_rcv{acc = [], datasize = NewSize}).
 
-parse_bidi(Data, State) ->
-    ts_plugin:parse_bidi(Data, State).
+parse_bidi(Data, State=#state_rcv{acc=[]}) ->
+    case websocket:decode(Data) of
+        {?OP_CLOSE, _Reason, _} ->
+            ?DebugF("receive close from server: ~p~n", [_Reason]),
+            {websocket:encode_close(<< >>), State, think};
+        {?OP_PING, Payload, Tail} ->
+            ts_mon_cache:add({count, websocket_ping}),
+            {websocket:encode_pong(Payload), State#state_rcv{acc = Tail}, think};
+        {_Opcode, _Payload, Tail} ->
+            {nodata, State#state_rcv{acc = Tail}, think}
+    end;
+parse_bidi(Data, State=#state_rcv{acc=Acc}) ->
+    parse_bidi(<< Acc/binary, Data/binary >>, State#state_rcv{acc = []}).
 
 %%----------------------------------------------------------------------
 %% Function: parse_config/2
diff --git a/tsung-1.0.dtd b/tsung-1.0.dtd
index 0b73e65dc..3bef4a881 100644
--- a/tsung-1.0.dtd
+++ b/tsung-1.0.dtd
@@ -376,8 +376,8 @@ repeat | if | change_type | foreach | set_option | interaction | abort )*>
     >
 <!ELEMENT websocket (#PCDATA | http_header)* >
 <!ATTLIST websocket
-    type (connect | message | close) #REQUIRED
-    ack  (no_ack | parse) #IMPLIED
+    type (connect | message | ping | close) #REQUIRED
+    ack  (no_ack | parse | bidi_ack) #IMPLIED
     frame (binary | text) #IMPLIED
     origin CDATA ""
     subprotocols CDATA ""