lokinet/llarp/dns/server.hpp

#pragma once

#include "message.hpp"
#include "platform.hpp"

#include <llarp/config/config.hpp>
#include <llarp/ev/ev.hpp>
#include <llarp/net/net.hpp>
#include <llarp/util/compare_ptr.hpp>
#include <llarp/util/fs.hpp>

#include <set>

namespace llarp::dns
{
    /// a job handling 1 dns query
    class QueryJob_Base
    {
       protected:
        /// the original dns query
        Message m_Query;

        /// True if we've sent a reply (including via a call to cancel)
        std::atomic_flag m_Done = ATOMIC_FLAG_INIT;

       public:
        explicit QueryJob_Base(Message query) : m_Query{std::move(query)}
        {}

        virtual ~QueryJob_Base() = default;

        Message& Underlying()
        {
            return m_Query;
        }

        const Message& Underlying() const
        {
            return m_Query;
        }

        /// cancel this operation and inform anyone who cares
        void Cancel();

        /// send a raw buffer back to the querier
        virtual void SendReply(llarp::OwnedBuffer replyBuf) = 0;
    };

    class PacketSource_Base
    {
       public:
        virtual ~PacketSource_Base() = default;

        /// return true if traffic with source and dest addresses would cause a
        /// loop in resolution and thus should not be sent to query handlers
        virtual bool WouldLoop(const SockAddr& to, const SockAddr& from) const = 0;

        /// send packet with src and dst address containing buf on this packet source
        virtual void SendTo(const SockAddr& to, const SockAddr& from, OwnedBuffer buf) const = 0;

        /// stop reading packets and end operation
        virtual void Stop() = 0;

        /// returns the sockaddr we are bound on if applicable
        virtual std::optional<SockAddr> BoundOn() const = 0;
    };

    /// a packet source which will override the sendto function of an wrapped packet source to
    /// construct a raw ip packet as a reply
    class PacketSource_Wrapper : public PacketSource_Base
    {
        std::weak_ptr<PacketSource_Base> m_Wrapped;
        std::function<void(net::IPPacket)> m_WritePacket;

       public:
        explicit PacketSource_Wrapper(
            std::weak_ptr<PacketSource_Base> wrapped,
            std::function<void(net::IPPacket)> write_packet)
            : m_Wrapped{wrapped}, m_WritePacket{write_packet}
        {}

        bool WouldLoop(const SockAddr& to, const SockAddr& from) const override
        {
            if (auto ptr = m_Wrapped.lock())
                return ptr->WouldLoop(to, from);
            return true;
        }

        void SendTo(const SockAddr& to, const SockAddr& from, OwnedBuffer buf) const override
        {
            m_WritePacket(net::IPPacket::make_udp(from, to, std::move(buf)));
        }

        /// stop reading packets and end operation
        void Stop() override
        {
            if (auto ptr = m_Wrapped.lock())
                ptr->Stop();
        }

        /// returns the sockaddr we are bound on if applicable
        std::optional<SockAddr> BoundOn() const override
        {
            if (auto ptr = m_Wrapped.lock())
                return ptr->BoundOn();
            return std::nullopt;
        }
    };

    /// non complex implementation of QueryJob_Base for use in things that
    /// only ever called on the mainloop thread
    class QueryJob : public QueryJob_Base, std::enable_shared_from_this<QueryJob>
    {
        std::shared_ptr<PacketSource_Base> src;
        const SockAddr resolver;
        const SockAddr asker;

       public:
        explicit QueryJob(
            std::shared_ptr<PacketSource_Base> source,
            const Message& query,
            const SockAddr& to_,
            const SockAddr& from_)
            : QueryJob_Base{query}, src{source}, resolver{to_}, asker{from_}
        {}

        void SendReply(llarp::OwnedBuffer replyBuf) override
        {
            src->SendTo(asker, resolver, std::move(replyBuf));
        }
    };

    /// handler of dns query hooking
    /// intercepts dns for internal processing
    class Resolver_Base
    {
       protected:
        /// return the sorting order for this resolver
        /// lower means it will be tried first
        virtual int Rank() const = 0;

       public:
        virtual ~Resolver_Base() = default;

        /// less than via rank
        bool operator<(const Resolver_Base& other) const
        {
            return Rank() < other.Rank();
        }

        /// greater than via rank
        bool operator>(const Resolver_Base& other) const
        {
            return Rank() > other.Rank();
        }

        /// get local socket address that queries are sent from
        virtual std::optional<SockAddr> GetLocalAddr() const
        {
            return std::nullopt;
        }

        /// get printable name
        virtual std::string_view ResolverName() const = 0;

        /// reset the resolver state, optionally replace upstream info with new info.  The default
        /// base implementation does nothing.
        virtual void ResetResolver(
            [[maybe_unused]] std::optional<std::vector<SockAddr>> replace_upstream = std::nullopt)
        {}

        /// cancel all pending requests and cease further operation.  Default operation is a no-op.
        virtual void Down()
        {}

        /// attempt to handle a dns message
        /// returns true if we consumed this query and it should not be processed again
        virtual bool MaybeHookDNS(
            std::shared_ptr<PacketSource_Base> source,
            const Message& query,
            const SockAddr& to,
            const SockAddr& from) = 0;
    };

    // Base class for DNS proxy
    class Server : public std::enable_shared_from_this<Server>
    {
       protected:
        /// add a packet source to this server, does share ownership
        void AddPacketSource(std::shared_ptr<PacketSource_Base> resolver);
        /// add a resolver to this packet handler, does share ownership
        void AddResolver(std::shared_ptr<Resolver_Base> resolver);

        /// create the platform dependant dns stuff
        virtual std::shared_ptr<I_Platform> CreatePlatform() const;

       public:
        virtual ~Server() = default;

        explicit Server(EventLoop_ptr loop, llarp::DnsConfig conf, unsigned int netif_index);

        /// returns all sockaddr we have from all of our PacketSources
        std::vector<SockAddr> BoundPacketSourceAddrs() const;

        /// returns the first sockaddr we have on our packet sources if we have one
        std::optional<SockAddr> FirstBoundPacketSourceAddr() const;

        /// add a resolver to this packet handler, does not share ownership
        void AddResolver(std::weak_ptr<Resolver_Base> resolver);

        /// add a packet source to this server, does not share ownership
        void AddPacketSource(std::weak_ptr<PacketSource_Base> resolver);

        /// create a packet source bound on bindaddr but does not add it
        virtual std::shared_ptr<PacketSource_Base> MakePacketSourceOn(
            const SockAddr& bindaddr, const llarp::DnsConfig& conf);

        /// sets up all internal binds and such and begins operation
        virtual void Start();

        /// stops all operation
        virtual void Stop();

        /// reset the internal state
        virtual void Reset();

        /// create the default resolver for out config
        virtual std::shared_ptr<Resolver_Base> MakeDefaultResolver();

        std::vector<std::weak_ptr<Resolver_Base>> GetAllResolvers() const;

        /// feed a packet buffer from a packet source.
        /// returns true if we decided to process the packet and consumed it
        /// returns false if we dont want to process the packet
        bool MaybeHandlePacket(
            std::shared_ptr<PacketSource_Base> pktsource,
            const SockAddr& resolver,
            const SockAddr& from,
            llarp::OwnedBuffer buf);
        /// set which dns mode we are in.
        /// true for intercepting all queries. false for just .loki and .snode
        void SetDNSMode(bool all_queries);

       protected:
        EventLoop_ptr m_Loop;
        llarp::DnsConfig m_Config;
        std::shared_ptr<I_Platform> m_Platform;

       private:
        const unsigned int m_NetIfIndex;
        std::set<std::shared_ptr<Resolver_Base>, ComparePtr<std::shared_ptr<Resolver_Base>>>
            m_OwnedResolvers;
        std::set<std::weak_ptr<Resolver_Base>, CompareWeakPtr<Resolver_Base>> m_Resolvers;

        std::vector<std::weak_ptr<PacketSource_Base>> m_PacketSources;
        std::vector<std::shared_ptr<PacketSource_Base>> m_OwnedPacketSources;
    };

}  // namespace llarp::dns