From 743bc2433ae21be699d152d79846b18f8836fcce Mon Sep 17 00:00:00 2001
From: Jeff Becker <jeff@i2p.rocks>
Date: Wed, 13 Oct 2021 08:54:19 -0400
Subject: [PATCH] resolve race condition in udp flow and packet handling

---
 llarp/lokinet_shared.cpp | 60 +++++++++++++++++++---------------------
 1 file changed, 28 insertions(+), 32 deletions(-)

diff --git a/llarp/lokinet_shared.cpp b/llarp/lokinet_shared.cpp
index e85df7389..a900b801e 100644
--- a/llarp/lokinet_shared.cpp
+++ b/llarp/lokinet_shared.cpp
@@ -151,7 +151,8 @@ namespace
         const AddressVariant_t& from,
         const lokinet_udp_flowinfo& flow_addr,
         void* flow_userdata,
-        int flow_timeoutseconds)
+        int flow_timeoutseconds,
+        std::optional<llarp::net::IPPacket> firstPacket = std::nullopt)
     {
       std::unique_lock lock{m_Access};
       auto& flow = m_Flows[from];
@@ -159,6 +160,8 @@ namespace
       flow.m_FlowTimeout = std::chrono::seconds{flow_timeoutseconds};
       flow.m_FlowUserData = flow_userdata;
       flow.m_Recv = m_Recv;
+      if (firstPacket)
+        flow.HandlePacket(*firstPacket);
     }
 
     void
@@ -180,44 +183,37 @@ namespace
     void
     HandlePacketFrom(AddressVariant_t from, llarp::net::IPPacket pkt)
     {
-      bool isNewFlow{false};
       {
         std::unique_lock lock{m_Access};
-        isNewFlow = m_Flows.count(from) == 0;
-      }
-      if (isNewFlow)
-      {
-        lokinet_udp_flowinfo flow_addr{};
-        // set flow remote address
-        var::visit(
-            [&flow_addr](auto&& from) {
-              const auto addr = from.ToString();
-              std::copy_n(
-                  addr.data(),
-                  std::min(addr.size(), sizeof(flow_addr.remote_host)),
-                  flow_addr.remote_host);
-            },
-            from);
-        // set socket id
-        flow_addr.socket_id = m_SocketID;
-        // get source port
-        if (auto srcport = pkt.SrcPort())
+        if (m_Flows.count(from))
         {
-          flow_addr.remote_port = ToHost(*srcport).h;
-        }
-        else
-          return;  // invalid data so we bail
-        void* flow_userdata = nullptr;
-        int flow_timeoutseconds{};
-        // got a new flow, let's check if we want it
-        if (m_Filter(m_User, &flow_addr, &flow_userdata, &flow_timeoutseconds))
+          m_Flows[from].HandlePacket(pkt);
           return;
-        AddFlow(from, flow_addr, flow_userdata, flow_timeoutseconds);
+        }
       }
+      lokinet_udp_flowinfo flow_addr{};
+      // set flow remote address
+      std::string addrstr = var::visit([&flow_addr](auto&& from) { return from.ToString(); }, from);
+
+      std::copy_n(
+          addrstr.data(),
+          std::min(addrstr.size(), sizeof(flow_addr.remote_host)),
+          flow_addr.remote_host);
+      // set socket id
+      flow_addr.socket_id = m_SocketID;
+      // get source port
+      if (const auto srcport = pkt.SrcPort())
       {
-        std::unique_lock lock{m_Access};
-        m_Flows[from].HandlePacket(pkt);
+        flow_addr.remote_port = ToHost(*srcport).h;
       }
+      else
+        return;  // invalid data so we bail
+      void* flow_userdata = nullptr;
+      int flow_timeoutseconds{};
+      // got a new flow, let's check if we want it
+      if (m_Filter(m_User, &flow_addr, &flow_userdata, &flow_timeoutseconds))
+        return;
+      AddFlow(from, flow_addr, flow_userdata, flow_timeoutseconds, pkt);
     }
   };
 }  // namespace