From fc9b09bdbcf8053e0755a0d9a0bda12bb0e01d41 Mon Sep 17 00:00:00 2001 From: Jeff Becker Date: Mon, 15 Mar 2021 12:01:19 -0400 Subject: [PATCH] clean up address/ip management code to use std::variant and std::optional --- llarp/handlers/null.hpp | 9 ++- llarp/handlers/tun.cpp | 93 ++++++++++++++--------------- llarp/handlers/tun.hpp | 26 ++------ llarp/service/endpoint.cpp | 19 +++--- llarp/service/endpoint.hpp | 18 +++--- pybind/llarp/handlers/pyhandler.hpp | 25 +++++--- 6 files changed, 92 insertions(+), 98 deletions(-) diff --git a/llarp/handlers/null.hpp b/llarp/handlers/null.hpp index 80f27d724..0e83a5595 100644 --- a/llarp/handlers/null.hpp +++ b/llarp/handlers/null.hpp @@ -41,11 +41,16 @@ namespace llarp void SendPacketToRemote(const llarp_buffer_t&) override{}; - huint128_t - ObtainIPForAddr(const AlignedBuffer<32>&, bool) override + huint128_t ObtainIPForAddr(std::variant) override { return {0}; } + + std::optional> ObtainAddrForIP( + huint128_t) const override + { + return std::nullopt; + } }; } // namespace handlers } // namespace llarp diff --git a/llarp/handlers/tun.cpp b/llarp/handlers/tun.cpp index 041d256a5..dd63cb7e7 100644 --- a/llarp/handlers/tun.cpp +++ b/llarp/handlers/tun.cpp @@ -1,5 +1,6 @@ #include #include +#include // harmless on other platforms #define __USE_MINGW_ANSI_STDIO 1 #include "tun.hpp" @@ -297,32 +298,6 @@ namespace llarp return msg.questions[0].IsLocalhost(); } - template <> - bool - TunEndpoint::FindAddrForIP(service::Address& addr, huint128_t ip) - { - auto itr = m_IPToAddr.find(ip); - if (itr != m_IPToAddr.end() and not m_SNodes[itr->second]) - { - addr = service::Address(itr->second.as_array()); - return true; - } - return false; - } - - template <> - bool - TunEndpoint::FindAddrForIP(RouterID& addr, huint128_t ip) - { - auto itr = m_IPToAddr.find(ip); - if (itr != m_IPToAddr.end() and m_SNodes[itr->second]) - { - addr = RouterID(itr->second.as_array()); - return true; - } - return false; - } - static dns::Message& clear_dns_message(dns::Message& msg) { @@ -333,13 +308,25 @@ namespace llarp return msg; } + std::optional> + TunEndpoint::ObtainAddrForIP(huint128_t ip) const + { + auto itr = m_IPToAddr.find(ip); + if (itr == m_IPToAddr.end()) + return std::nullopt; + if (m_SNodes.at(itr->second)) + return RouterID{itr->second.as_array()}; + else + return service::Address{itr->second.as_array()}; + } + bool TunEndpoint::HandleHookedDNSMessage(dns::Message msg, std::function reply) { auto ReplyToSNodeDNSWhenReady = [self = this, reply = reply]( RouterID snode, auto msg, bool isV6) -> bool { return self->EnsurePathToSNode(snode, [=](const RouterID&, exit::BaseSession_ptr s) { - self->SendDNSReply(snode, s, msg, reply, true, isV6); + self->SendDNSReply(snode, s, msg, reply, isV6); }); }; auto ReplyToLokiDNSWhenReady = [self = this, reply = reply]( @@ -349,7 +336,7 @@ namespace llarp return self->EnsurePathToService( addr, [=](const Address&, OutboundContext* ctx) { - self->SendDNSReply(addr, ctx, msg, reply, false, isV6); + self->SendDNSReply(addr, ctx, msg, reply, isV6); }, 2s); }; @@ -666,17 +653,10 @@ namespace llarp reply(msg); return true; } - RouterID snodeAddr; - if (FindAddrForIP(snodeAddr, ip)) - { - msg.AddAReply(snodeAddr.ToString()); - reply(msg); - return true; - } - service::Address lokiAddr; - if (FindAddrForIP(lokiAddr, ip)) + + if (auto maybe = ObtainAddrForIP(ip)) { - msg.AddAReply(lokiAddr.ToString()); + std::visit([&msg](auto&& result) { msg.AddAReply(result.ToString()); }, *maybe); reply(msg); return true; } @@ -1043,9 +1023,12 @@ namespace llarp if (t != service::ProtocolType::TrafficV4 && t != service::ProtocolType::TrafficV6 && t != service::ProtocolType::Exit) return false; - AlignedBuffer<32> addr; - bool snode = false; - if (!GetEndpointWithConvoTag(tag, addr, snode)) + std::variant addr; + if (auto maybe = GetEndpointWithConvoTag(tag)) + { + addr = *maybe; + } + else return false; huint128_t src, dst; @@ -1056,7 +1039,7 @@ namespace llarp if (m_state->m_ExitEnabled) { // exit side from exit - src = ObtainIPForAddr(addr, snode); + src = ObtainIPForAddr(addr); if (t == service::ProtocolType::Exit) { if (pkt.IsV4()) @@ -1088,16 +1071,22 @@ namespace llarp } // find what exit we think this should be for const auto mapped = m_ExitMap.FindAll(src); - if (mapped.count(service::Address{addr}) == 0 or IsBogon(src)) - { - // we got exit traffic from someone who we should not have gotten it from + if (IsBogon(src)) return false; + + if (const auto ptr = std::get_if(&addr)) + { + if (mapped.count(*ptr) == 0) + { + // we got exit traffic from someone who we should not have gotten it from + return false; + } } } else { // snapp traffic - src = ObtainIPForAddr(addr, snode); + src = ObtainIPForAddr(addr); dst = m_OurIP; } HandleWriteIPPacket(buf, src, dst, seqno); @@ -1136,10 +1125,20 @@ namespace llarp } huint128_t - TunEndpoint::ObtainIPForAddr(const AlignedBuffer<32>& ident, bool snode) + TunEndpoint::ObtainIPForAddr(std::variant addr) { llarp_time_t now = Now(); huint128_t nextIP = {0}; + AlignedBuffer<32> ident{}; + bool snode = false; + + std::visit([&ident](auto&& val) { ident = val.data(); }, addr); + + if (std::get_if(&addr)) + { + snode = true; + } + { // previously allocated address auto itr = m_AddrToIP.find(ident); diff --git a/llarp/handlers/tun.hpp b/llarp/handlers/tun.hpp index 31399e86a..5fb5facad 100644 --- a/llarp/handlers/tun.hpp +++ b/llarp/handlers/tun.hpp @@ -13,6 +13,8 @@ #include #include +#include +#include namespace llarp { @@ -121,23 +123,8 @@ namespace llarp HasLocalIP(const huint128_t& ip) const; /// get a key for ip address - template - Addr_t - ObtainAddrForIP(huint128_t ip, bool isSNode) - { - Addr_t addr; - auto itr = m_IPToAddr.find(ip); - if (itr != m_IPToAddr.end() and m_SNodes[itr->second] == isSNode) - { - addr = Addr_t(itr->second); - } - // found - return addr; - } - - template - bool - FindAddrForIP(Addr_t& addr, huint128_t ip); + std::optional> + ObtainAddrForIP(huint128_t ip) const override; bool HasAddress(const AlignedBuffer<32>& addr) const @@ -147,7 +134,7 @@ namespace llarp /// get ip address for key unconditionally huint128_t - ObtainIPForAddr(const AlignedBuffer<32>& addr, bool serviceNode) override; + ObtainIPForAddr(std::variant addr) override; /// flush network traffic void @@ -214,12 +201,11 @@ namespace llarp Endpoint_t ctx, std::shared_ptr query, std::function reply, - bool snode, bool sendIPv6) { if (ctx) { - huint128_t ip = ObtainIPForAddr(addr, snode); + huint128_t ip = ObtainIPForAddr(addr); query->answers.clear(); query->AddINReply(ip, sendIPv6); } diff --git a/llarp/service/endpoint.cpp b/llarp/service/endpoint.cpp index 10b899db4..ee56edfd5 100644 --- a/llarp/service/endpoint.cpp +++ b/llarp/service/endpoint.cpp @@ -31,6 +31,7 @@ #include #include +#include #include #include @@ -210,29 +211,23 @@ namespace llarp return routers.find(remote) != routers.end(); } - bool - Endpoint::GetEndpointWithConvoTag( - const ConvoTag tag, llarp::AlignedBuffer<32>& addr, bool& snode) const + std::optional> + Endpoint::GetEndpointWithConvoTag(ConvoTag tag) const { auto itr = Sessions().find(tag); if (itr != Sessions().end()) { - snode = false; - addr = itr->second.remote.Addr(); - return true; + return itr->second.remote.Addr(); } for (const auto& item : m_state->m_SNodeSessions) { if (item.second.second == tag) { - snode = true; - addr = item.first; - return true; + return item.first; } } - - return false; + return std::nullopt; } bool @@ -1345,7 +1340,7 @@ namespace llarp // some day :DDDDD tag.Randomize(); const auto src = xhtonl(net::TruncateV6(GetIfAddr())); - const auto dst = xhtonl(net::TruncateV6(ObtainIPForAddr(snode, true))); + const auto dst = xhtonl(net::TruncateV6(ObtainIPForAddr(snode))); auto session = std::make_shared( snode, diff --git a/llarp/service/endpoint.hpp b/llarp/service/endpoint.hpp index 956318561..0f60c16da 100644 --- a/llarp/service/endpoint.hpp +++ b/llarp/service/endpoint.hpp @@ -17,6 +17,7 @@ #include "lookup.hpp" #include #include +#include #include #include "endpoint_types.hpp" @@ -178,8 +179,11 @@ namespace llarp void SetAuthInfoForEndpoint(Address remote, AuthInfo info); - virtual huint128_t - ObtainIPForAddr(const AlignedBuffer<32>& addr, bool serviceNode) = 0; + virtual huint128_t ObtainIPForAddr(std::variant) = 0; + + /// get a key for ip address + virtual std::optional> + ObtainAddrForIP(huint128_t ip) const = 0; // virtual bool // HasServiceAddress(const AlignedBuffer< 32 >& addr) const = 0; @@ -273,13 +277,9 @@ namespace llarp void BlacklistSNode(const RouterID snode) override; - /// return true if we have a convotag as an exit session - /// or as a hidden service session - /// set addr and issnode - /// - /// return false if we don't have either - bool - GetEndpointWithConvoTag(const ConvoTag t, AlignedBuffer<32>& addr, bool& issnode) const; + /// maybe get an endpoint variant given its convo tag + std::optional> + GetEndpointWithConvoTag(ConvoTag t) const; bool HasConvoTag(const ConvoTag& t) const override; diff --git a/pybind/llarp/handlers/pyhandler.hpp b/pybind/llarp/handlers/pyhandler.hpp index 871981dc0..2ed934548 100644 --- a/pybind/llarp/handlers/pyhandler.hpp +++ b/pybind/llarp/handlers/pyhandler.hpp @@ -30,16 +30,20 @@ namespace llarp { if (handlePacket) { - AlignedBuffer<32> addr; - bool isSnode = false; - if (not GetEndpointWithConvoTag(tag, addr, isSnode)) + service::Address addr{}; + if (auto maybe = GetEndpointWithConvoTag(tag)) + { + if (auto ptr = std::get_if(&*maybe)) + addr = *ptr; + else + return false; + } + else return false; - if (isSnode) - return true; std::vector pkt; pkt.resize(pktbuf.sz); std::copy_n(pktbuf.base, pktbuf.sz, pkt.data()); - handlePacket(service::Address(addr), std::move(pkt), proto); + handlePacket(addr, std::move(pkt), proto); } return true; } @@ -56,12 +60,17 @@ namespace llarp return false; } - llarp::huint128_t - ObtainIPForAddr(const llarp::AlignedBuffer<32>&, bool) override + llarp::huint128_t ObtainIPForAddr(std::variant) override { return {0}; } + std::optional> ObtainAddrForIP( + huint128_t) const override + { + return std::nullopt; + } + std::string GetIfName() const override {