From 09e01adf32d674734dbaecec9bd85dd589ceb616 Mon Sep 17 00:00:00 2001
From: Jeff Becker <jeff@i2p.rocks>
Date: Wed, 12 Dec 2018 13:37:03 -0500
Subject: [PATCH] flush snode traffic queues

---
 CMakeLists.txt                |  3 ++-
 llarp/handlers/tun.cpp        |  4 +++-
 llarp/service/endpoint.cpp    | 14 ++++++++++++--
 llarp/service/endpoint.hpp    |  3 +++
 test/test_service_address.cpp | 24 ++++++++++++++++++++++++
 5 files changed, 44 insertions(+), 4 deletions(-)
 create mode 100644 test/test_service_address.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index fddfb0f47..228f0fc73 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -547,6 +547,7 @@ set(TEST_SRC
   test/test_llarp_queue.cpp
   test/test_llarp_queue_manager.cpp
   test/test_llarp_thread_pool.cpp
+  test/test_service_address.cpp
 )
 
 set(TEST_EXE testAll)
@@ -661,4 +662,4 @@ if(WITH_STATIC)
     #target_link_libraries(${RC_EXE} ${SHARED_LIB})
     #target_link_libraries(${DNS_EXE} ${SHARED_LIB} Threads::Threads)
   endif(WITH_SHARED)
-endif(SHADOW)
\ No newline at end of file
+endif(SHADOW)
diff --git a/llarp/handlers/tun.cpp b/llarp/handlers/tun.cpp
index c80f31fda..23fc4a6e3 100644
--- a/llarp/handlers/tun.cpp
+++ b/llarp/handlers/tun.cpp
@@ -205,7 +205,7 @@ namespace llarp
         {
           if(HasAddress(addr.data()))
           {
-            huint32_t ip = ObtainIPForAddr(addr.data(), true);
+            huint32_t ip = ObtainIPForAddr(addr.data(), false);
             msg.AddINReply(ip);
           }
           else
@@ -460,6 +460,7 @@ namespace llarp
 
         if(m_SNodes.at(itr->second))
         {
+          llarp::LogInfo(Name(), " send to service node");
           sendFunc = std::bind(&TunEndpoint::SendToSNodeOrQueue, this,
                                itr->second.data(), std::placeholders::_1);
         }
@@ -480,6 +481,7 @@ namespace llarp
       });
       if(m_Exit)
         m_Exit->FlushUpstreamTraffic();
+      FlushSNodeTraffic();
     }
 
     bool
diff --git a/llarp/service/endpoint.cpp b/llarp/service/endpoint.cpp
index c2c50f572..2c6f2c102 100644
--- a/llarp/service/endpoint.cpp
+++ b/llarp/service/endpoint.cpp
@@ -123,6 +123,17 @@ namespace llarp
       }
     }
 
+    void
+    Endpoint::FlushSNodeTraffic()
+    {
+      auto itr = m_SNodeSessions.begin();
+      while(itr != m_SNodeSessions.end())
+      {
+        itr->second->FlushUpstreamTraffic();
+        ++itr;
+      }
+    }
+
     void
     Endpoint::Tick(llarp_time_t now)
     {
@@ -1068,8 +1079,7 @@ namespace llarp
     void
     Endpoint::EnsurePathToSNode(const RouterID& snode)
     {
-      auto range = m_SNodeSessions.equal_range(snode);
-      if(range.first == range.second)
+      if(m_SNodeSessions.count(snode) == 0)
       {
         auto themIP = ObtainIPForAddr(snode, true);
         m_SNodeSessions.emplace(std::make_pair(
diff --git a/llarp/service/endpoint.hpp b/llarp/service/endpoint.hpp
index ff319c394..951436790 100644
--- a/llarp/service/endpoint.hpp
+++ b/llarp/service/endpoint.hpp
@@ -162,6 +162,9 @@ namespace llarp
       bool
       SendToSNodeOrQueue(const byte_t* addr, llarp_buffer_t payload);
 
+      void
+      FlushSNodeTraffic();
+
       struct PendingBuffer
       {
         std::vector< byte_t > payload;
diff --git a/test/test_service_address.cpp b/test/test_service_address.cpp
new file mode 100644
index 000000000..e800f86c7
--- /dev/null
+++ b/test/test_service_address.cpp
@@ -0,0 +1,24 @@
+#include <gtest/gtest.h>
+#include <service/address.hpp>
+
+struct ServiceAddressTest : public ::testing::Test
+{
+  const std::string snode =
+      "8zfiwpgonsu5zpddpxwdurxyb19x6r96xy4qbikff99jwsziws9y.snode";
+  const std::string loki =
+      "7okic5x5do3uh3usttnqz9ek3uuoemdrwzto1hciwim9f947or6y.loki";
+};
+
+TEST_F(ServiceAddressTest, TestParseSNodeNotLoki)
+{
+  llarp::service::Address addr;
+  ASSERT_TRUE(addr.FromString(snode, ".snode"));
+  ASSERT_FALSE(addr.FromString(snode, ".loki"));
+}
+
+TEST_F(ServiceAddressTest, TestParseLokiNotSNode)
+{
+  llarp::service::Address addr;
+  ASSERT_FALSE(addr.FromString(loki, ".snode"));
+  ASSERT_TRUE(addr.FromString(loki, ".loki"));
+}