Add ability to follow UDP stream by index

-z "follow,udp" tshark cli command now supports a stream index
It is now possible to select the UDP stream displayed in Qt GUI (like for TCP)

Change-Id: Ia367f36ea4f60db0fddb997a7e0903c09e172f2d
Reviewed-on: https://code.wireshark.org/review/6083
Reviewed-by: Pascal Quantin <pascal.quantin@gmail.com>

5 files changed, 125 insertions, 65 deletions

diff --git a/epan/dissectors/packet-tcp.h b/epan/dissectors/packet-tcp.h
index 13d3172cda..8f6e0450dd 100644
--- a/epan/dissectors/packet-tcp.h
+++ b/epan/dissectors/packet-tcp.h
@@ -28,10 +28,7 @@ extern "C" {
 
 #include "ws_symbol_export.h"
 
-#ifndef __CONVERSATION_H__
 #include <epan/conversation.h>
-#endif
-
 #include <epan/wmem/wmem.h>
 
 /* TCP flags */
diff --git a/epan/dissectors/packet-udp.c b/epan/dissectors/packet-udp.c
index 693e009d35..35301677a8 100644
--- a/epan/dissectors/packet-udp.c
+++ b/epan/dissectors/packet-udp.c
@@ -232,7 +232,7 @@ init_udp_conversation_data(void)
   return udpd;
 }
 
-static struct udp_analysis *
+struct udp_analysis *
 get_udp_conversation_data(conversation_t *conv, packet_info *pinfo)
 {
   int direction;
diff --git a/epan/dissectors/packet-udp.h b/epan/dissectors/packet-udp.h
index b90de18b6d..a11a0afa17 100644
--- a/epan/dissectors/packet-udp.h
+++ b/epan/dissectors/packet-udp.h
@@ -23,6 +23,14 @@
 #ifndef __PACKET_UDP_H__
 #define __PACKET_UDP_H__
 
+#ifdef __cplusplus
+extern "C" {
+#endif /* __cplusplus */
+
+#include "ws_symbol_export.h"
+
+#include <epan/conversation.h>
+
 /* UDP structs and definitions */
 typedef struct _e_udphdr {
   guint16 uh_sport;
@@ -98,4 +106,11 @@ WS_DLL_PUBLIC guint32 get_udp_stream_count(void);
 WS_DLL_PUBLIC void decode_udp_ports(tvbuff_t *, int, packet_info *,
 	proto_tree *, int, int, int);
 
+WS_DLL_PUBLIC struct udp_analysis *get_udp_conversation_data(conversation_t *,
+	packet_info *);
+
+#ifdef __cplusplus
+}
+#endif /* __cplusplus */
+
 #endif
diff --git a/epan/follow.c b/epan/follow.c
index 4b894393d9..61f73b596f 100644
--- a/epan/follow.c
+++ b/epan/follow.c
@@ -36,8 +36,10 @@
 #include <epan/to_str.h>
 #include <epan/emem.h>
 #include <epan/dissectors/packet-tcp.h>
+#include <epan/dissectors/packet-udp.h>
 #include "follow.h"
 #include <epan/conversation.h>
+#include <epan/tap.h>
 
 #define MAX_IPADDR_LEN  16
 
@@ -55,7 +57,10 @@ FILE* data_out_file = NULL;
 gboolean empty_tcp_stream;
 gboolean incomplete_tcp_stream;
 
-static guint32 tcp_stream_to_follow = 0;
+static guint32 stream_to_follow[MAX_STREAM] = {0};
+static gboolean find_addr[MAX_STREAM] = {FALSE};
+static gboolean find_index[MAX_STREAM] = {FALSE};
+static address tcp_addr[2];
 static guint8  ip_address[2][MAX_IPADDR_LEN];
 static guint   port[2];
 static guint   bytes_written[2];
@@ -87,6 +92,7 @@ build_follow_conv_filter( packet_info *pi ) {
   int len;
   conversation_t *conv=NULL;
   struct tcp_analysis *tcpd;
+  struct udp_analysis *udpd;
   wmem_list_frame_t* protos;
   int proto_id;
   const char* proto_name;
@@ -114,11 +120,11 @@ build_follow_conv_filter( packet_info *pi ) {
        (pi->net_src.type == AT_IPv6 && pi->net_dst.type == AT_IPv6))
        && is_tcp && (conv=find_conversation(pi->fd->num, &pi->src, &pi->dst, pi->ptype,
               pi->srcport, pi->destport, 0)) != NULL ) {
-    /* TCP over IPv4 */
+    /* TCP over IPv4/6 */
     tcpd=get_tcp_conversation_data(conv, pi);
     if (tcpd) {
       buf = g_strdup_printf("tcp.stream eq %d", tcpd->stream);
-      tcp_stream_to_follow = tcpd->stream;
+      stream_to_follow[TCP_STREAM] = tcpd->stream;
       if (pi->net_src.type == AT_IPv4) {
         len = 4;
         is_ipv6 = FALSE;
@@ -130,27 +136,25 @@ build_follow_conv_filter( packet_info *pi ) {
       return NULL;
     }
   }
-  else if( pi->net_src.type == AT_IPv4 && pi->net_dst.type == AT_IPv4
-          && is_udp ) {
-  /* UDP over IPv4 */
-    buf = g_strdup_printf(
-            "(ip.addr eq %s and ip.addr eq %s) and (udp.port eq %d and udp.port eq %d)",
-            address_to_str(pi->pool, &pi->net_src),
-            address_to_str(pi->pool, &pi->net_dst),
-            pi->srcport, pi->destport );
-    len = 4;
-    is_ipv6 = FALSE;
-  }
-  else if( pi->net_src.type == AT_IPv6 && pi->net_dst.type == AT_IPv6
-       && is_udp ) {
-    /* UDP over IPv6 */
-    buf = g_strdup_printf(
-            "(ipv6.addr eq %s and ipv6.addr eq %s) and (udp.port eq %d and udp.port eq %d)",
-            address_to_str(pi->pool, &pi->net_src),
-            address_to_str(pi->pool, &pi->net_dst),
-            pi->srcport, pi->destport );
-    len = 16;
-    is_ipv6 = TRUE;
+  else if( ((pi->net_src.type == AT_IPv4 && pi->net_dst.type == AT_IPv4) ||
+            (pi->net_src.type == AT_IPv6 && pi->net_dst.type == AT_IPv6))
+          && is_udp && (conv=find_conversation(pi->fd->num, &pi->src, &pi->dst, pi->ptype,
+              pi->srcport, pi->destport, 0)) != NULL ) {
+    /* UDP over IPv4/6 */
+    udpd=get_udp_conversation_data(conv, pi);
+    if (udpd) {
+      buf = g_strdup_printf("udp.stream eq %d", udpd->stream);
+      stream_to_follow[UDP_STREAM] = udpd->stream;
+      if (pi->net_src.type == AT_IPv4) {
+        len = 4;
+        is_ipv6 = FALSE;
+      } else {
+        len = 16;
+        is_ipv6 = TRUE;
+      }
+    } else {
+      return NULL;
+    }
   }
   else {
     return NULL;
@@ -162,30 +166,62 @@ build_follow_conv_filter( packet_info *pi ) {
   return buf;
 }
 
-static gboolean         find_tcp_addr;
-static address          tcp_addr[2];
-static gboolean         find_tcp_index;
+static gboolean
+udp_follow_packet(void *tapdata _U_, packet_info *pinfo,
+                  epan_dissect_t *edt _U_, const void *data _U_)
+{
+  if (find_addr[UDP_STREAM]) {
+    if (pinfo->net_src.type == AT_IPv6) {
+      is_ipv6 = TRUE;
+    } else {
+      is_ipv6 = FALSE;
+    }
+    memcpy(ip_address[0], pinfo->net_src.data, pinfo->net_src.len);
+    memcpy(ip_address[1], pinfo->net_dst.data, pinfo->net_dst.len);
+    port[0] = pinfo->srcport;
+    port[1] = pinfo->destport;
+    find_addr[UDP_STREAM] = FALSE;
+  }
+
+  return FALSE;
+}
+
+void
+reset_udp_follow(void) {
+  remove_tap_listener(&stream_to_follow[UDP_STREAM]);
+  find_addr[UDP_STREAM] = FALSE;
+  find_addr[UDP_STREAM] = FALSE;
+}
 
 gchar*
-build_follow_index_filter(void) {
+build_follow_index_filter(stream_type stream) {
   gchar *buf;
 
-  find_tcp_addr = TRUE;
-  buf = g_strdup_printf("tcp.stream eq %d", tcp_stream_to_follow);
+  find_addr[stream] = TRUE;
+  if (stream == TCP_STREAM) {
+    buf = g_strdup_printf("tcp.stream eq %d", stream_to_follow[TCP_STREAM]);
+  } else {
+    GString * error_string;
+    buf = g_strdup_printf("udp.stream eq %d", stream_to_follow[UDP_STREAM]);
+    error_string = register_tap_listener("udp_follow", &stream_to_follow[UDP_STREAM], buf, 0, NULL, udp_follow_packet, NULL);
+    if (error_string) {
+      g_string_free(error_string, TRUE);
+    }
+  }
   return buf;
 }
 
 /* select a tcp stream to follow via it's address/port pairs */
 gboolean
-follow_tcp_addr(const address *addr0, guint port0,
-                const address *addr1, guint port1)
+follow_addr(stream_type stream, const address *addr0, guint port0,
+            const address *addr1, guint port1)
 {
   if (addr0 == NULL || addr1 == NULL || addr0->type != addr1->type ||
       port0 > G_MAXUINT16 || port1 > G_MAXUINT16 )  {
     return FALSE;
   }
 
-  if (find_tcp_index || find_tcp_addr) {
+  if (find_index[stream] || find_addr[stream]) {
     return FALSE;
   }
 
@@ -198,29 +234,32 @@ follow_tcp_addr(const address *addr0, guint port0,
     break;
   }
 
-  find_tcp_index = TRUE;
 
   memcpy(ip_address[0], addr0->data, addr0->len);
-  SET_ADDRESS(&tcp_addr[0], addr0->type, addr0->len, ip_address[0]);
   port[0] = port0;
 
   memcpy(ip_address[1], addr1->data, addr1->len);
-  SET_ADDRESS(&tcp_addr[1], addr1->type, addr1->len, ip_address[1]);
   port[1] = port1;
 
+  if (stream == TCP_STREAM) {
+    find_index[TCP_STREAM] = TRUE;
+    SET_ADDRESS(&tcp_addr[0], addr0->type, addr0->len, ip_address[0]);
+    SET_ADDRESS(&tcp_addr[1], addr1->type, addr1->len, ip_address[1]);
+  }
+
   return TRUE;
 }
 
-/* select a tcp stream to follow via its index */
+/* select a stream to follow via its index */
 gboolean
-follow_tcp_index(guint32 indx)
+follow_index(stream_type stream, guint32 indx)
 {
-  if (find_tcp_index || find_tcp_addr) {
+  if (find_index[stream] || find_addr[stream]) {
     return FALSE;
   }
 
-  find_tcp_addr = TRUE;
-  tcp_stream_to_follow = indx;
+  find_addr[stream] = TRUE;
+  stream_to_follow[stream] = indx;
   memset(ip_address, 0, sizeof ip_address);
   port[0] = port[1] = 0;
 
@@ -228,8 +267,8 @@ follow_tcp_index(guint32 indx)
 }
 
 guint32
-get_follow_tcp_index(void) {
-  return tcp_stream_to_follow;
+get_follow_index(stream_type stream) {
+  return stream_to_follow[stream];
 }
 
 /* here we are going to try and reconstruct the data portion of a TCP
@@ -255,7 +294,7 @@ reassemble_tcp( guint32 tcp_stream, guint32 sequence, guint32 acknowledgement,
   src_index = -1;
 
   /* First, check if this packet should be processed. */
-  if (find_tcp_index) {
+  if (find_index[TCP_STREAM]) {
     if ((port[0] == srcport && port[1] == dstport &&
          ADDRESSES_EQUAL(&tcp_addr[0], net_src) &&
          ADDRESSES_EQUAL(&tcp_addr[1], net_dst))
@@ -263,14 +302,14 @@ reassemble_tcp( guint32 tcp_stream, guint32 sequence, guint32 acknowledgement,
         (port[1] == srcport && port[0] == dstport &&
          ADDRESSES_EQUAL(&tcp_addr[1], net_src) &&
          ADDRESSES_EQUAL(&tcp_addr[0], net_dst))) {
-      find_tcp_index = FALSE;
-      tcp_stream_to_follow = tcp_stream;
+      find_index[TCP_STREAM] = FALSE;
+      stream_to_follow[TCP_STREAM] = tcp_stream;
     }
     else {
       return;
     }
   }
-  else if ( tcp_stream != tcp_stream_to_follow )
+  else if ( tcp_stream != stream_to_follow[TCP_STREAM] )
     return;
 
   if ((net_src->type != AT_IPv4 && net_src->type != AT_IPv6) ||
@@ -286,8 +325,8 @@ reassemble_tcp( guint32 tcp_stream, guint32 sequence, guint32 acknowledgement,
   memcpy(dstx, net_dst->data, len);
 
   /* follow_tcp_index() needs to learn address/port pairs */
-  if (find_tcp_addr) {
-    find_tcp_addr = FALSE;
+  if (find_addr[TCP_STREAM]) {
+    find_addr[TCP_STREAM] = FALSE;
     memcpy(ip_address[0], net_src->data, net_src->len);
     port[0] = srcport;
     memcpy(ip_address[1], net_dst->data, net_dst->len);
@@ -518,8 +557,8 @@ reset_tcp_reassembly(void)
 
   empty_tcp_stream = TRUE;
   incomplete_tcp_stream = FALSE;
-  find_tcp_addr = FALSE;
-  find_tcp_index = FALSE;
+  find_addr[TCP_STREAM] = FALSE;
+  find_index[TCP_STREAM] = FALSE;
   for( i=0; i<2; i++ ) {
     seq[i] = 0;
     memset(src_addr[i], '\0', MAX_IPADDR_LEN);
diff --git a/epan/follow.h b/epan/follow.h
index 173762279f..e883446949 100644
--- a/epan/follow.h
+++ b/epan/follow.h
@@ -34,6 +34,12 @@ extern "C" {
 
 #define MAX_IPADDR_LEN	16
 
+typedef enum {
+  TCP_STREAM = 0,
+  UDP_STREAM,
+  MAX_STREAM
+} stream_type;
+
 /* With MSVC and a libwireshark.dll, we need a special declaration. */
 WS_DLL_PUBLIC gboolean empty_tcp_stream;
 WS_DLL_PUBLIC gboolean incomplete_tcp_stream;
@@ -54,39 +60,42 @@ typedef struct _tcp_stream_chunk {
 WS_DLL_PUBLIC
 gchar* build_follow_conv_filter( packet_info * packet_info);
 
-/** Build a follow filter based on the current TCP stream index.
- * follow_tcp_index() must be called prior to calling this.
+/** Build a follow filter based on the current TCP/UDP stream index.
+ * follow_index() must be called prior to calling this.
  *
  * @return A filter that specifies the current stream. Must be g_free()d
  * the caller.
  */
 WS_DLL_PUBLIC
-gchar* build_follow_index_filter(void);
+gchar* build_follow_index_filter(stream_type stream);
 
 WS_DLL_PUBLIC
-gboolean follow_tcp_addr( const address *, guint, const address *, guint );
+gboolean follow_addr(stream_type, const address *, guint, const address *, guint );
 
-/** Select a TCP stream to follow via its index.
+/** Select a TCP/UDP stream to follow via its index.
  *
  * @param addr [in] The stream index to follow.
  * @return TRUE on success, FALSE on failure.
  */
 WS_DLL_PUBLIC
-gboolean follow_tcp_index( guint32 addr);
+gboolean follow_index(stream_type stream, guint32 addr);
 
-/** Get the current TCP index being followed.
+/** Get the current TCP/UDP index being followed.
  *
- * @return The current TCP index. The behavior is undefined
- * if no TCP stream is being followed.
+ * @return The current TCP/UDP index. The behavior is undefined
+ * if no TCP/UDP stream is being followed.
  */
 WS_DLL_PUBLIC
-guint32 get_follow_tcp_index(void);
+guint32 get_follow_index(stream_type stream);
 
 void reassemble_tcp( guint32, guint32, guint32, guint32, const char*, guint32,
                      int, address *, address *, guint, guint, guint32 );
 WS_DLL_PUBLIC
 void  reset_tcp_reassembly( void );
 
+WS_DLL_PUBLIC
+void reset_udp_follow(void);
+
 typedef struct {
 	guint8		ip_address[2][MAX_IPADDR_LEN];
 	guint32		port[2];