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>

3 files changed, 50 insertions, 17 deletions

diff --git a/ui/cli/tap-follow.c b/ui/cli/tap-follow.c
index 9c448e3255..6be940fe7b 100644
--- a/ui/cli/tap-follow.c
+++ b/ui/cli/tap-follow.c
@@ -150,6 +150,7 @@ followStrFilter(
   static char   filter[512];
   int           len     = 0;
   const gchar  *verp;
+  gchar        *udpfilter;
   gchar         ip0[MAX_IP6_STR_LEN];
   gchar         ip1[MAX_IP6_STR_LEN];
 
@@ -163,6 +164,10 @@ followStrFilter(
                      "tcp.stream eq %d", fp->index);
       break;
     case type_UDP:
+      udpfilter = build_follow_index_filter(UDP_STREAM);
+      len = g_snprintf(filter, sizeof filter,
+                     "%s", udpfilter);
+      g_free(udpfilter);
       break;
     }
   }
@@ -512,7 +517,7 @@ followDraw(
 
   g_assert(sizeof bin % BYTES_PER_LINE == 0);
 
-  if (fp->type == type_TCP)
+  if ((fp->type == type_TCP) || (fp->type == type_UDP))
   {
     static const guint8 ip_zero[MAX_IPADDR_LEN] = {0};
     follow_stats_t      stats;
@@ -888,17 +893,17 @@ followTcp(
   reset_tcp_reassembly();
   if (fp->index != G_MAXUINT32)
   {
-    if (!follow_tcp_index(fp->index))
+    if (!follow_index(TCP_STREAM, fp->index))
     {
-      followExit("Can't follow tcp index.");
+      followExit("Can't follow TCP index.");
     }
   }
   else
   {
-    if (!follow_tcp_addr(&fp->addr[0], fp->port[0],
-                         &fp->addr[1], fp->port[1]))
+    if (!follow_addr(TCP_STREAM, &fp->addr[0], fp->port[0],
+                     &fp->addr[1], fp->port[1]))
     {
-      followExit("Can't follow tcp address/port pairs.");
+      followExit("Can't follow TCP address/port pairs.");
     }
   }
 
@@ -910,7 +915,7 @@ followTcp(
   {
     followFree(fp);
     g_string_free(errp, TRUE);
-    followExit("Error registering tcp tap listener.");
+    followExit("Error registering TCP tap listener.");
   }
 }
 
@@ -932,9 +937,21 @@ followUdp(
   followArgRange(&opt_argp, fp);
   followArgDone(opt_argp);
 
+  reset_udp_follow();
   if (fp->index != G_MAXUINT32)
   {
-    followExit("UDP does not support index filters.");
+    if (!follow_index(UDP_STREAM, fp->index))
+    {
+      followExit("Can't follow UDP index.");
+    }
+  }
+  else
+  {
+    if (!follow_addr(UDP_STREAM, &fp->addr[0], fp->port[0],
+                     &fp->addr[1], fp->port[1]))
+    {
+      followExit("Can't follow UDP address/port pairs.");
+    }
   }
 
   followFileOpen(fp);
@@ -945,7 +962,7 @@ followUdp(
   {
     followFree(fp);
     g_string_free(errp, TRUE);
-    followExit("Error registering udp tap listner.");
+    followExit("Error registering UDP tap listner.");
   }
 }
 
@@ -967,6 +984,7 @@ followSsl(
   followArgRange(&opt_argp, fp);
   followArgDone(opt_argp);
 
+  reset_tcp_reassembly();
   if (fp->index == G_MAXUINT32)
   {
     followExit("SSL only supports index filters.");
@@ -980,7 +998,7 @@ followSsl(
   {
     followFree(fp);
     g_string_free(errp, TRUE);
-    followExit("Error registering ssl tap listener.");
+    followExit("Error registering SSL tap listener.");
   }
 }
 
diff --git a/ui/qt/follow_stream_dialog.cpp b/ui/qt/follow_stream_dialog.cpp
index 13c8f47d6b..a6f2638546 100644
--- a/ui/qt/follow_stream_dialog.cpp
+++ b/ui/qt/follow_stream_dialog.cpp
@@ -27,6 +27,7 @@
 
 #include "epan/follow.h"
 #include "epan/dissectors/packet-tcp.h"
+#include "epan/dissectors/packet-udp.h"
 #include "epan/prefs.h"
 #include "epan/addr_resolv.h"
 #include "epan/charsets.h"
@@ -290,7 +291,7 @@ void FollowStreamDialog::on_streamNumberSpinBox_valueChanged(int stream_num)
 {
     if (stream_num >= 0) {
         updateWidgets(false);
-        follow_tcp_index(stream_num);
+        follow_index((follow_type_ == FOLLOW_TCP) ? TCP_STREAM : UDP_STREAM, stream_num);
         follow(QString(), true);
         updateWidgets();
     }
@@ -327,6 +328,7 @@ void FollowStreamDialog::resetStream()
     }
     g_list_free(follow_info_.payload);
     follow_info_.payload = NULL;
+    follow_info_.client_port = 0;
 }
 
 frs_return_t
@@ -840,7 +842,7 @@ FollowStreamDialog::showBuffer(char *buffer, size_t nchars, gboolean is_from_ser
     return FRS_OK;
 }
 
-bool FollowStreamDialog::follow(QString previous_filter, bool use_tcp_index)
+bool FollowStreamDialog::follow(QString previous_filter, bool use_stream_index)
 {
     int                 tmp_fd;
     QString             follow_filter;
@@ -880,7 +882,6 @@ bool FollowStreamDialog::follow(QString previous_filter, bool use_tcp_index)
         }
         break;
     case FOLLOW_UDP:
-        removeStreamControls();
         if (!is_udp) {
             QMessageBox::warning(this, tr("Error following stream."), tr("Please make sure you have a UDP packet selected."));
             return false;
@@ -902,10 +903,13 @@ bool FollowStreamDialog::follow(QString previous_filter, bool use_tcp_index)
         /* Create a new filter that matches all packets in the TCP stream,
            and set the display filter entry accordingly */
         reset_tcp_reassembly();
+    } else {
+        reset_udp_follow();
     }
 
-    if (use_tcp_index) {
-        follow_filter = gchar_free_to_qstring(build_follow_index_filter());
+    if (use_stream_index) {
+        follow_filter = gchar_free_to_qstring(
+            build_follow_index_filter((follow_type_ == FOLLOW_TCP) ? TCP_STREAM : UDP_STREAM));
     } else {
         follow_filter = gchar_free_to_qstring(build_follow_conv_filter(&cap_file_->edt->pi));
     }
@@ -967,7 +971,7 @@ bool FollowStreamDialog::follow(QString previous_filter, bool use_tcp_index)
         int stream_count = get_tcp_stream_count() - 1;
         ui->streamNumberSpinBox->blockSignals(true);
         ui->streamNumberSpinBox->setMaximum(stream_count);
-        ui->streamNumberSpinBox->setValue(get_follow_tcp_index());
+        ui->streamNumberSpinBox->setValue(get_follow_index(TCP_STREAM));
         ui->streamNumberSpinBox->blockSignals(false);
         ui->streamNumberSpinBox->setToolTip(tr("%Ln total stream(s).", "", stream_count));
         ui->streamNumberLabel->setToolTip(ui->streamNumberSpinBox->toolTip());
@@ -975,6 +979,7 @@ bool FollowStreamDialog::follow(QString previous_filter, bool use_tcp_index)
         break;
     }
     case FOLLOW_UDP:
+    {
         /* data will be passed via tap callback*/
         msg = register_tap_listener("udp_follow", &follow_info_,
                                     follow_filter.toUtf8().constData(),
@@ -985,7 +990,17 @@ bool FollowStreamDialog::follow(QString previous_filter, bool use_tcp_index)
                                msg->str);
             return false;
         }
+
+        int stream_count = get_udp_stream_count() - 1;
+        ui->streamNumberSpinBox->blockSignals(true);
+        ui->streamNumberSpinBox->setMaximum(stream_count);
+        ui->streamNumberSpinBox->setValue(get_follow_index(UDP_STREAM));
+        ui->streamNumberSpinBox->blockSignals(false);
+        ui->streamNumberSpinBox->setToolTip(tr("%Ln total stream(s).", "", stream_count));
+        ui->streamNumberLabel->setToolTip(ui->streamNumberSpinBox->toolTip());
+
         break;
+    }
     case FOLLOW_SSL:
         /* we got ssl so we can follow */
         msg = register_tap_listener("ssl", &follow_info_,
diff --git a/ui/qt/follow_stream_dialog.h b/ui/qt/follow_stream_dialog.h
index 1e5969b3c3..9a2fa55e37 100644
--- a/ui/qt/follow_stream_dialog.h
+++ b/ui/qt/follow_stream_dialog.h
@@ -68,7 +68,7 @@ public:
     explicit FollowStreamDialog(QWidget *parent = 0, follow_type_t type = FOLLOW_TCP, capture_file *cf = NULL);
     ~FollowStreamDialog();
 
-    bool follow(QString previous_filter = QString(), bool use_tcp_index = false);
+    bool follow(QString previous_filter = QString(), bool use_stream_index = false);
 
 public slots:
     void setCaptureFile(capture_file *cf);