From 5909ad03864bf2f905a86d856896a46fd9d3101d Mon Sep 17 00:00:00 2001
From: Jeff Becker <jeff@i2p.rocks>
Date: Sat, 5 Jun 2021 09:06:17 -0400
Subject: [PATCH] add MarkAddressOutbound to plainquic

---
 llarp/endpoint_base.hpp    | 3 +++
 llarp/handlers/exit.hpp    | 2 ++
 llarp/quic/tunnel.cpp      | 4 ++++
 llarp/service/endpoint.cpp | 5 +++--
 llarp/service/endpoint.hpp | 3 +--
 llarp/service/handler.hpp  | 4 ----
 6 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/llarp/endpoint_base.hpp b/llarp/endpoint_base.hpp
index 2a33f7602..ce13786ac 100644
--- a/llarp/endpoint_base.hpp
+++ b/llarp/endpoint_base.hpp
@@ -130,6 +130,9 @@ namespace llarp
         std::string name,
         std::string service,
         std::function<void(std::vector<dns::SRVData>)> resultHandler) = 0;
+
+    virtual void
+    MarkAddressOutbound(AddressVariant_t remote) = 0;
   };
 
 }  // namespace llarp
diff --git a/llarp/handlers/exit.hpp b/llarp/handlers/exit.hpp
index 82e5c786b..9addb3ab4 100644
--- a/llarp/handlers/exit.hpp
+++ b/llarp/handlers/exit.hpp
@@ -41,6 +41,8 @@ namespace llarp
       void
       SRVRecordsChanged() override;
 
+      void MarkAddressOutbound(AddressVariant_t) override{};
+
       bool
       SendToOrQueue(
           service::ConvoTag tag, const llarp_buffer_t& payload, service::ProtocolType t) override;
diff --git a/llarp/quic/tunnel.cpp b/llarp/quic/tunnel.cpp
index b3da985b9..588dc4efc 100644
--- a/llarp/quic/tunnel.cpp
+++ b/llarp/quic/tunnel.cpp
@@ -528,6 +528,7 @@ namespace llarp::quic
             if (not continue_connecting(
                     pport, (bool)maybe_remote, "endpoint ONS lookup", remote_addr))
               return;
+            service_endpoint_.MarkAddressOutbound(*maybe_remote);
             service_endpoint_.EnsurePathTo(*maybe_remote, after_path, open_timeout);
           });
       return result;
@@ -539,7 +540,10 @@ namespace llarp::quic
     if (auto maybe_convo = service_endpoint_.GetBestConvoTagFor(remote))
       after_path(maybe_convo);
     else
+    {
+      service_endpoint_.MarkAddressOutbound(remote);
       service_endpoint_.EnsurePathTo(remote, after_path, open_timeout);
+    }
 
     return result;
   }
diff --git a/llarp/service/endpoint.cpp b/llarp/service/endpoint.cpp
index f11925d04..f995a3b84 100644
--- a/llarp/service/endpoint.cpp
+++ b/llarp/service/endpoint.cpp
@@ -1366,9 +1366,10 @@ namespace llarp
     }
 
     void
-    Endpoint::MarkAddressOutbound(const Address& addr)
+    Endpoint::MarkAddressOutbound(AddressVariant_t addr)
     {
-      m_state->m_OutboundSessions.insert(addr);
+      if (auto* ptr = std::get_if<Address>(&addr))
+        m_state->m_OutboundSessions.insert(*ptr);
     }
 
     bool
diff --git a/llarp/service/endpoint.hpp b/llarp/service/endpoint.hpp
index 740feec01..0fd8c625c 100644
--- a/llarp/service/endpoint.hpp
+++ b/llarp/service/endpoint.hpp
@@ -286,8 +286,7 @@ namespace llarp
       WantsOutboundSession(const Address&) const override;
 
       /// this MUST be called if you want to call EnsurePathTo on the given address
-      void
-      MarkAddressOutbound(const Address&) override;
+      void MarkAddressOutbound(AddressVariant_t) override;
 
       bool
       ShouldBundleRC() const override
diff --git a/llarp/service/handler.hpp b/llarp/service/handler.hpp
index 4ae29c8f3..74498c0ca 100644
--- a/llarp/service/handler.hpp
+++ b/llarp/service/handler.hpp
@@ -77,10 +77,6 @@ namespace llarp
       /// do we want a session outbound to addr
       virtual bool
       WantsOutboundSession(const Address& addr) const = 0;
-
-      virtual void
-      MarkAddressOutbound(const Address& addr) = 0;
-
       virtual void
       QueueRecvData(RecvDataEvent ev) = 0;
     };