From 6310b7d9e3cbf8c82d97308f277a073d3b239d41 Mon Sep 17 00:00:00 2001
From: dmiller <dmiller@e0a8ed71-7df4-0310-8962-fdc924857419>
Date: Mon, 18 Jan 2021 22:23:55 +0000
Subject: [PATCH] Ensure all UDP payloads get sent before giving up on a port.

---
 payload.cc     | 14 ++++++++++++++
 payload.h      |  1 +
 scan_engine.cc |  7 ++++++-
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/payload.cc b/payload.cc
index 9025e2646..c7bfe8fb6 100644
--- a/payload.cc
+++ b/payload.cc
@@ -363,3 +363,17 @@ const char *get_udp_payload(u16 dport, size_t *length, u8 tryno) {
     return udp_port2payload(dport, length, tryno);
   }
 }
+
+size_t udp_payload_count(u16 dport) {
+  std::map<struct proto_dport, std::vector<struct payload> >::iterator portPayloadIterator;
+  proto_dport key(IPPROTO_UDP, dport);
+  size_t portPayloadVectorSize = 0;
+
+  portPayloadIterator = portPayloads.find(key);
+
+  if (portPayloadIterator != portPayloads.end()) {
+    portPayloadVectorSize = portPayloadIterator->second.size();
+  }
+
+  return portPayloadVectorSize;
+}
diff --git a/payload.h b/payload.h
index 322c33e92..ab1cba146 100644
--- a/payload.h
+++ b/payload.h
@@ -68,6 +68,7 @@
 
 const char *get_udp_payload(u16 dport, size_t *length, u8 tryno);
 const char *udp_port2payload(u16 dport, size_t *length, u8 tryno);
+size_t udp_payload_count(u16 dport);
 int init_payloads(void);
 
 #endif /* PAYLOAD_H */
diff --git a/scan_engine.cc b/scan_engine.cc
index cec3064cf..be775c1a7 100644
--- a/scan_engine.cc
+++ b/scan_engine.cc
@@ -2454,7 +2454,12 @@ static void doAnyOutstandingRetransmits(UltraScanInfo *USI) {
         probeI--;
         probe = *probeI;
         if (probe->timedout && !probe->retransmitted &&
-            maxtries > probe->tryno && !probe->isPing()) {
+            (maxtries > probe->tryno ||
+            // We may exceed maxtries if this is UDP...
+             ((USI->udp_scan || (USI->ping_scan && USI->ptech.rawudpscan))
+             // ...and there are more payloads we haven't tried.
+              && udp_payload_count(probe->dport()) > probe->tryno)
+            ) && !probe->isPing()) {
           /* For rate limit detection, we delay the first time a new tryno
              is seen, as long as we are scanning at least 2 ports */
           if (probe->tryno + 1 > (int) host->rld.max_tryno_sent &&