Add follow websocket stream support

The websocket protocol masking feature makes follow TCP stream
on websocket traffic show masked payload. To easily view unmasked
and reassembled websocket payload add follow websocket stream
support.

7 files changed, 34 insertions, 1 deletions

diff --git a/epan/dissectors/packet-websocket.c b/epan/dissectors/packet-websocket.c
index 22a56120f9..74ab98be67 100644
--- a/epan/dissectors/packet-websocket.c
+++ b/epan/dissectors/packet-websocket.c
@@ -13,7 +13,9 @@
 #include "config.h"
 #include <wsutil/wslog.h>
 
+#include <epan/addr_resolv.h>
 #include <epan/conversation.h>
+#include <epan/follow.h>
 #include <epan/proto_data.h>
 #include <epan/packet.h>
 #include <epan/expert.h>
@@ -75,6 +77,8 @@ typedef struct {
 } websocket_packet_t;
 #endif
 
+static int websocket_follow_tap = -1;
+
 /* Initialize the protocol and registered fields */
 static int proto_websocket = -1;
 static int proto_http = -1;
@@ -691,6 +695,10 @@ dissect_websocket_frame(tvbuff_t *tvb, packet_info *pinfo, proto_tree *tree, voi
       tvb_payload = tvb_new_subset_length_caplen(tvb, payload_offset, payload_length, payload_length);
     }
     dissect_websocket_payload(tvb_payload, pinfo, tree, ws_tree, fin, opcode, websocket_conv, pmc, tvb_raw_offset(tvb));
+
+    if (have_tap_listener(websocket_follow_tap)) {
+      tap_queue_packet(websocket_follow_tap, pinfo, tvb_payload);
+    }
   }
 
   return tvb_captured_length(tvb);
@@ -978,6 +986,11 @@ proto_register_websocket(void)
 
   reassembly_table_register(&ws_reassembly_table, &addresses_reassembly_table_functions);
 
+  websocket_follow_tap = register_tap("websocket_follow"); /* websocket follow tap */
+  register_follow_stream(proto_websocket, "websocket_follow", tcp_follow_conv_filter, tcp_follow_index_filter,
+                         tcp_follow_address_filter,	tcp_port_to_display, follow_tvb_tap_listener,
+                         get_tcp_stream_count, NULL);
+
   proto_register_field_array(proto_websocket, hf, array_length(hf));
   proto_register_subtree_array(ett, array_length(ett));
   expert_websocket = expert_register_protocol(proto_websocket);
diff --git a/epan/follow.h b/epan/follow.h
index f17b3f0c19..e23e387412 100644
--- a/epan/follow.h
+++ b/epan/follow.h
@@ -47,6 +47,7 @@ typedef enum {
     FOLLOW_HTTP2,
     FOLLOW_QUIC,
     FOLLOW_SIP,
+    FOLLOW_WEBSOCKET,
 } follow_type_t;
 
 /* Show Type */
diff --git a/ui/qt/follow_stream_dialog.cpp b/ui/qt/follow_stream_dialog.cpp
index b4ce678e5e..417761b79d 100644
--- a/ui/qt/follow_stream_dialog.cpp
+++ b/ui/qt/follow_stream_dialog.cpp
@@ -116,6 +116,9 @@ FollowStreamDialog::FollowStreamDialog(QWidget &parent, CaptureFile &cf, follow_
     case FOLLOW_SIP:
         follower_ = get_follow_by_name("SIP");
         break;
+    case FOLLOW_WEBSOCKET:
+        follower_ = get_follow_by_name("WebSocket");
+        break;
     default :
         ws_assert_not_reached();
     }
@@ -587,6 +590,7 @@ FollowStreamDialog::readStream()
     case FOLLOW_QUIC:
     case FOLLOW_TLS :
     case FOLLOW_SIP :
+    case FOLLOW_WEBSOCKET :
         ret = readFollowStream();
         break;
 
diff --git a/ui/qt/packet_list.cpp b/ui/qt/packet_list.cpp
index 227fc139c3..831121ac7d 100644
--- a/ui/qt/packet_list.cpp
+++ b/ui/qt/packet_list.cpp
@@ -699,6 +699,7 @@ void PacketList::contextMenuEvent(QContextMenuEvent *event)
     submenu->addAction(window()->findChild<QAction *>("actionAnalyzeFollowHTTPStream"));
     submenu->addAction(window()->findChild<QAction *>("actionAnalyzeFollowHTTP2Stream"));
     submenu->addAction(window()->findChild<QAction *>("actionAnalyzeFollowQUICStream"));
+    submenu->addAction(window()->findChild<QAction *>("actionAnalyzeFollowWebsocketStream"));
     submenu->addAction(window()->findChild<QAction *>("actionAnalyzeFollowSIPCall"));
 
     ctx_menu->addSeparator();
diff --git a/ui/qt/wireshark_main_window.cpp b/ui/qt/wireshark_main_window.cpp
index 7c47292014..ce1401e099 100644
--- a/ui/qt/wireshark_main_window.cpp
+++ b/ui/qt/wireshark_main_window.cpp
@@ -655,6 +655,9 @@ main_ui_->goToLineEdit->setValidator(goToLineQiv);
     connect(main_ui_->actionAnalyzeFollowQUICStream, &QAction::triggered, this,
             [this]() { this->openFollowStreamDialogForType(FOLLOW_QUIC); },
             Qt::QueuedConnection);
+    connect(main_ui_->actionAnalyzeFollowWebsocketStream, &QAction::triggered, this,
+            [this]() { this->openFollowStreamDialogForType(FOLLOW_WEBSOCKET); },
+            Qt::QueuedConnection);
     connect(main_ui_->actionAnalyzeFollowSIPCall, &QAction::triggered, this,
             [this]() { this->openFollowStreamDialogForType(FOLLOW_SIP); },
             Qt::QueuedConnection);
diff --git a/ui/qt/wireshark_main_window.ui b/ui/qt/wireshark_main_window.ui
index 963f795f61..cf33b3ba3b 100644
--- a/ui/qt/wireshark_main_window.ui
+++ b/ui/qt/wireshark_main_window.ui
@@ -424,6 +424,7 @@
      <addaction name="actionAnalyzeFollowHTTPStream"/>
      <addaction name="actionAnalyzeFollowHTTP2Stream"/>
      <addaction name="actionAnalyzeFollowQUICStream"/>
+     <addaction name="actionAnalyzeFollowWebsocketStream"/>
      <addaction name="actionAnalyzeFollowSIPCall"/>
     </widget>
     <widget class="QMenu" name="menuConversationFilter">
@@ -1785,6 +1786,14 @@
     <string>QUIC Stream</string>
    </property>
   </action>
+  <action name="actionAnalyzeFollowWebsocketStream">
+   <property name="enabled">
+    <bool>false</bool>
+   </property>
+   <property name="text">
+    <string>Websocket Stream</string>
+   </property>
+  </action>
   <action name="actionAnalyzeFollowSIPCall">
    <property name="enabled">
     <bool>false</bool>
diff --git a/ui/qt/wireshark_main_window_slots.cpp b/ui/qt/wireshark_main_window_slots.cpp
index 3e0fdc07ea..e6051a359b 100644
--- a/ui/qt/wireshark_main_window_slots.cpp
+++ b/ui/qt/wireshark_main_window_slots.cpp
@@ -1168,7 +1168,7 @@ void WiresharkMainWindow::setEditCommentsMenu()
 void WiresharkMainWindow::setMenusForSelectedPacket()
 {
     gboolean is_ip = FALSE, is_tcp = FALSE, is_udp = FALSE, is_dccp = FALSE, is_sctp = FALSE, is_tls = FALSE, is_rtp = FALSE, is_lte_rlc = FALSE,
-             is_http = FALSE, is_http2 = FALSE, is_quic = FALSE, is_sip = FALSE, is_exported_pdu = FALSE;
+             is_http = FALSE, is_http2 = FALSE, is_quic = FALSE, is_sip = FALSE, is_websocket = FALSE, is_exported_pdu = FALSE;
 
     /* Making the menu context-sensitive allows for easier selection of the
        desired item and has the added benefit, with large captures, of
@@ -1243,6 +1243,7 @@ void WiresharkMainWindow::setMenusForSelectedPacket()
             /* TODO: to follow a QUIC stream we need a *decrypted* QUIC connection, i.e. checking for "quic" in the protocol stack is not enough */
             is_quic = proto_is_frame_protocol(capture_file_.capFile()->edt->pi.layers, "quic");
             is_sip = proto_is_frame_protocol(capture_file_.capFile()->edt->pi.layers, "sip");
+            is_websocket = proto_is_frame_protocol(capture_file_.capFile()->edt->pi.layers, "websocket");
             is_exported_pdu = proto_is_frame_protocol(capture_file_.capFile()->edt->pi.layers, "exported_pdu");
             /* For Exported PDU there is a tag inserting IP addresses into the SRC and DST columns */
             if (is_exported_pdu &&
@@ -1297,6 +1298,7 @@ void WiresharkMainWindow::setMenusForSelectedPacket()
     main_ui_->actionAnalyzeFollowHTTPStream->setEnabled(is_http);
     main_ui_->actionAnalyzeFollowHTTP2Stream->setEnabled(is_http2);
     main_ui_->actionAnalyzeFollowQUICStream->setEnabled(is_quic);
+    main_ui_->actionAnalyzeFollowWebsocketStream->setEnabled(is_websocket);
     main_ui_->actionAnalyzeFollowSIPCall->setEnabled(is_sip);
 
     foreach(QAction *cc_action, cc_actions) {