securedht.h

// Copyright (c) 2014-2026 Savoir-faire Linux Inc.
// SPDX-License-Identifier: MIT
#pragma once
 
#include "dht.h"
#include "crypto.h"
 
#include <map>
#include <vector>
#include <memory>
#include <random>
 
namespace dht {
 
class OPENDHT_PUBLIC SecureDht final : public DhtInterface
{
public:
    typedef std::function<void(bool)> SignatureCheckCallback;
 
    using Config = SecureDhtConfig;
 
    static dht::Config getConfig(const SecureDht::Config& conf)
    {
        auto c = conf.node_config;
        if (not c.node_id and conf.id.second)
            c.node_id = InfoHash::get("node:" + conf.id.second->getId().toString());
        return c;
    }

    SecureDht(std::unique_ptr<DhtInterface> dht,
              Config config,
              IdentityAnnouncedCb iacb = {},
              const std::shared_ptr<Logger>& l = {});
 
    virtual ~SecureDht();
 
    InfoHash getId() const { return key_ ? key_->getPublicKey().getId() : InfoHash(); }
    PkId getLongId() const { return key_ ? key_->getPublicKey().getLongId() : PkId(); }
    Sp<crypto::PublicKey> getPublicKey() const { return key_ ? key_->getSharedPublicKey() : Sp<crypto::PublicKey> {}; }
 
    ValueType secureType(ValueType&& type);
 
    ValueType secureType(const ValueType& type)
    {
        ValueType tmp_type = type;
        return secureType(std::move(tmp_type));
    }
 
    void registerType(const ValueType& type) override
    {
        if (dht_)
            dht_->registerType(secureType(type));
    }
    void registerType(ValueType&& type)
    {
        if (dht_)
            dht_->registerType(secureType(std::forward<ValueType>(type)));
    }
    void registerInsecureType(const ValueType& type)
    {
        if (dht_)
            dht_->registerType(type);
    }

    void get(const InfoHash& id, GetCallback cb, DoneCallback donecb = {}, Value::Filter&& = {}, Where&& w = {}) override;
    void get(const InfoHash& id,
             GetCallback cb,
             DoneCallbackSimple donecb = {},
             Value::Filter&& f = {},
             Where&& w = {}) override
    {
        get(id, cb, bindDoneCb(donecb), std::forward<Value::Filter>(f), std::forward<Where>(w));
    }
    void get(const InfoHash& key,
             GetCallbackSimple cb,
             DoneCallback donecb = {},
             Value::Filter&& f = {},
             Where&& w = {}) override
    {
        get(key, bindGetCb(cb), donecb, std::forward<Value::Filter>(f), std::forward<Where>(w));
    }
    void get(const InfoHash& key,
             GetCallbackSimple cb,
             DoneCallbackSimple donecb,
             Value::Filter&& f = {},
             Where&& w = {}) override
    {
        get(key, bindGetCb(cb), bindDoneCb(donecb), std::forward<Value::Filter>(f), std::forward<Where>(w));
    }

    void putSigned(const InfoHash& hash, Sp<Value> val, DoneCallback callback, bool permanent = false);
    void putSigned(const InfoHash& hash, Value&& v, DoneCallback callback, bool permanent = false)
    {
        putSigned(hash, std::make_shared<Value>(std::move(v)), callback, permanent);
    }

    void putEncrypted(
        const InfoHash& hash, const InfoHash& to, Sp<Value> val, DoneCallback callback, bool permanent = false);
 
    [[deprecated("Use the shared_ptr version instead")]]
    void putEncrypted(const InfoHash& hash, const InfoHash& to, Value&& v, DoneCallback callback, bool permanent = false)
    {
        putEncrypted(hash, to, std::make_shared<Value>(std::move(v)), callback, permanent);
    }
    void putEncrypted(const InfoHash& hash,
                      const crypto::PublicKey& to,
                      Sp<Value> val,
                      DoneCallback callback,
                      bool permanent = false);
 
    [[deprecated("Use the shared_ptr version instead")]]
    void putEncrypted(
        const InfoHash& hash, const crypto::PublicKey& to, Value&& v, DoneCallback callback, bool permanent = false)
    {
        putEncrypted(hash, to, std::make_shared<Value>(std::move(v)), callback, permanent);
    }
 
    void putEncrypted(const InfoHash& hash, const PkId& to, Sp<Value> val, DoneCallback callback, bool permanent = false);

    void sign(Value& v) const;
 
    Value encrypt(Value& v, const crypto::PublicKey& to) const;
 
    Value decrypt(const Value& v);
 
    void findCertificate(const InfoHash& node, const std::function<void(const Sp<crypto::Certificate>)>& cb);
    void findPublicKey(const InfoHash& node, const std::function<void(const Sp<crypto::PublicKey>)>& cb);
 
    void findCertificate(const PkId& id, const std::function<void(const Sp<crypto::Certificate>)>& cb);
    void findPublicKey(const PkId& id, const std::function<void(const Sp<crypto::PublicKey>)>& cb);
 
    void registerCertificate(const Sp<crypto::Certificate>& cert);
 
    Sp<crypto::Certificate> getCertificate(const InfoHash& node) const;
    Sp<crypto::PublicKey> getPublicKey(const InfoHash& node) const;
 
    Sp<crypto::Certificate> getCertificate(const PkId& node) const;
    Sp<crypto::PublicKey> getPublicKey(const PkId& node) const;

    void setLocalCertificateStore(CertificateStoreQueryLegacy&& query_method)
    {
        localQueryMethodLegacy_ = std::move(query_method);
    }
    void setLocalCertificateStore(CertificateStoreQuery&& query_method) { localQueryMethod_ = std::move(query_method); }
    void setOnPublicAddressChanged(PublicAddressChangedCb cb) override { dht_->setOnPublicAddressChanged(cb); }

    void shutdown(ShutdownCallback cb, bool stop = false) override { dht_->shutdown(cb, stop); }
    void dumpTables() const override { dht_->dumpTables(); }
    inline const InfoHash& getNodeId() const override { return dht_->getNodeId(); }
 
    std::pair<size_t, size_t> getStoreSize() const override { return dht_->getStoreSize(); }
    std::pair<size_t, size_t> getLocalStoreSize() const override { return dht_->getLocalStoreSize(); }
    std::string getStorageLog() const override { return dht_->getStorageLog(); }
    std::string getStorageLog(const InfoHash& h) const override { return dht_->getStorageLog(h); }
    void setStorageLimit(size_t limit = 0) override { dht_->setStorageLimit(limit); }
    size_t getStorageLimit() const override { return dht_->getStorageLimit(); }
    void setLocalStorageLimit(size_t limit = 0) override { dht_->setLocalStorageLimit(limit); }
    size_t getLocalStorageLimit() const override { return dht_->getLocalStorageLimit(); }
 
    std::vector<NodeExport> exportNodes() const override { return dht_->exportNodes(); }
    std::vector<ValuesExport> exportValues() const override { return dht_->exportValues(); }
    void importValues(const std::vector<ValuesExport>& v) override { dht_->importValues(v); }
    NodeStats getNodesStats(sa_family_t af) const override { return dht_->getNodesStats(af); }
 
    NodeInfo getNodeInfo()
    {
        NodeInfo info;
        info.id = getId();
        info.node_id = getNodeId();
        info.ipv4 = getNodesStats(AF_INET);
        info.ipv6 = getNodesStats(AF_INET6);
        std::tie(info.storage_size, info.storage_values) = getStoreSize();
        std::tie(info.local_storage_size, info.local_storage_values) = getLocalStoreSize();
        if (auto sock = getSocket()) {
            info.bound4 = sock->getBoundRef(AF_INET).getPort();
            info.bound6 = sock->getBoundRef(AF_INET6).getPort();
        }
        return info;
    }
 
    std::vector<unsigned> getNodeMessageStats(bool in = false) override { return dht_->getNodeMessageStats(in); }
    std::string getRoutingTablesLog(sa_family_t af) const override { return dht_->getRoutingTablesLog(af); }
    std::string getSearchesLog(sa_family_t af) const override { return dht_->getSearchesLog(af); }
    std::string getSearchLog(const InfoHash& h, sa_family_t af = AF_UNSPEC) const override
    {
        return dht_->getSearchLog(h, af);
    }
    std::vector<SockAddr> getPublicAddress(sa_family_t family = 0) override { return dht_->getPublicAddress(family); }
    time_point periodic(const uint8_t* buf, size_t buflen, SockAddr sa, const time_point& now) override
    {
        return dht_->periodic(buf, buflen, std::move(sa), now);
    }
    time_point periodic(
        const uint8_t* buf, size_t buflen, const sockaddr* from, socklen_t fromlen, const time_point& now) override
    {
        return dht_->periodic(buf, buflen, from, fromlen, now);
    }
    NodeStatus updateStatus(sa_family_t af) override { return dht_->updateStatus(af); }
    NodeStatus getStatus(sa_family_t af) const override { return dht_->getStatus(af); }
    NodeStatus getStatus() const override { return dht_->getStatus(); }
    net::DatagramSocket* getSocket() const override { return dht_->getSocket(); };
    bool isRunning(sa_family_t af = 0) const override { return dht_->isRunning(af); }
    const ValueType& getType(ValueType::Id type_id) const override { return dht_->getType(type_id); }
    void addBootstrap(const std::string& host, const std::string& service) override
    {
        dht_->addBootstrap(host, service);
    }
    void clearBootstrap() override { dht_->clearBootstrap(); }
    void insertNode(const InfoHash& id, const SockAddr& sa) override { dht_->insertNode(id, sa); }
    void insertNode(const NodeExport& n) override { dht_->insertNode(n); }
    void pingNode(SockAddr sa, DoneCallbackSimple&& cb = {}) override { dht_->pingNode(std::move(sa), std::move(cb)); }
    void query(const InfoHash& key, QueryCallback cb, DoneCallback done_cb = {}, Query&& q = {}) override
    {
        dht_->query(key, cb, done_cb, std::move(q));
    }
    void query(const InfoHash& key, QueryCallback cb, DoneCallbackSimple done_cb = {}, Query&& q = {}) override
    {
        dht_->query(key, cb, done_cb, std::move(q));
    }
    std::vector<Sp<Value>> getLocal(const InfoHash& key, const Value::Filter& f = {}) const override
    {
        return dht_->getLocal(key, f);
    }
    Sp<Value> getLocalById(const InfoHash& key, Value::Id vid) const override { return dht_->getLocalById(key, vid); }
    void put(const InfoHash& key,
             Sp<Value> v,
             DoneCallback cb = nullptr,
             time_point created = time_point::max(),
             bool permanent = false) override
    {
        dht_->put(key, v, cb, created, permanent);
    }
    void put(const InfoHash& key,
             const Sp<Value>& v,
             DoneCallbackSimple cb,
             time_point created = time_point::max(),
             bool permanent = false) override
    {
        dht_->put(key, v, cb, created, permanent);
    }
 
    void put(const InfoHash& key,
             Value&& v,
             DoneCallback cb = nullptr,
             time_point created = time_point::max(),
             bool permanent = false) override
    {
        dht_->put(key, std::move(v), cb, created, permanent);
    }
    void put(const InfoHash& key,
             Value&& v,
             DoneCallbackSimple cb,
             time_point created = time_point::max(),
             bool permanent = false) override
    {
        dht_->put(key, std::move(v), cb, created, permanent);
    }
    std::vector<Sp<Value>> getPut(const InfoHash& h) const override { return dht_->getPut(h); }
    Sp<Value> getPut(const InfoHash& h, const Value::Id& vid) const override { return dht_->getPut(h, vid); }
    bool cancelPut(const InfoHash& h, const Value::Id& vid) override { return dht_->cancelPut(h, vid); }
 
    size_t listen(const InfoHash& key, ValueCallback, Value::Filter = {}, Where = {}) override;
    size_t listen(const InfoHash& key, GetCallback cb, Value::Filter = {}, Where w = {}) override;
    size_t listen(const InfoHash& key, GetCallbackSimple cb, Value::Filter f = {}, Where w = {}) override
    {
        return listen(key, bindGetCb(cb), f, w);
    }
    bool cancelListen(const InfoHash& h, size_t token) override { return dht_->cancelListen(h, token); }
    void connectivityChanged(sa_family_t af) override { dht_->connectivityChanged(af); }
    void connectivityChanged() override { dht_->connectivityChanged(); }
 
    void forwardAllMessages(bool forward) { forward_all_ = forward; }
 
    void setPushNotificationToken(const std::string& token = "") override { dht_->setPushNotificationToken(token); }

    PushNotificationResult pushNotificationReceived(const std::map<std::string, std::string>& notification) override
    {
        return dht_->pushNotificationReceived(notification);
    }
 
    void setLogger(const std::shared_ptr<Logger>& logger) override
    {
        DhtInterface::setLogger(logger);
        dht_->setLogger(logger);
    }

    void setLogFilter(const InfoHash& f) override
    {
        DhtInterface::setLogFilter(f);
        dht_->setLogFilter(f);
    }
 
private:
    std::unique_ptr<DhtInterface> dht_;
    // prevent copy
    SecureDht(const SecureDht&) = delete;
    SecureDht& operator=(const SecureDht&) = delete;
 
    Sp<Value> checkValue(const Sp<Value>& v);
    ValueCallback getCallbackFilter(const ValueCallback&, Value::Filter&&);
    GetCallback getCallbackFilter(const GetCallback&, Value::Filter&&);
 
    Sp<crypto::Certificate> registerCertificate(const InfoHash& node, const Blob& cert);
    Sp<crypto::Certificate> registerCertificate(const PkId& node, const Blob& cert);
 
    Sp<crypto::PrivateKey> key_ {};
    Sp<crypto::Certificate> certificate_ {};
 
    // method to query the local certificate store
    CertificateStoreQuery localQueryMethod_ {};
    CertificateStoreQueryLegacy localQueryMethodLegacy_ {};
 
    // our certificate cache
    std::map<InfoHash, Sp<crypto::Certificate>> nodesCertificates_ {};
    std::map<InfoHash, Sp<crypto::PublicKey>> nodesPubKeys_ {};
    std::map<PkId, Sp<crypto::Certificate>> nodesCertificatesLong_ {};
    std::map<PkId, Sp<crypto::PublicKey>> nodesPubKeysLong_ {};
 
    std::atomic_bool forward_all_ {false};
    bool enableCache_ {false};
};
 
const ValueType CERTIFICATE_TYPE = {8,
                                    "Certificate",
                                    std::chrono::hours(24 * 7),
                                    // A certificate can only be stored at its public key ID.
                                    [](InfoHash id, Sp<Value>& v, const InfoHash&, const SockAddr&) {
                                        try {
                                            crypto::Certificate crt(v->data);
                                            // TODO check certificate signature
                                            return crt.getPublicKey().getId() == id
                                                   || InfoHash::get(crt.getPublicKey().getLongId()) == id;
                                        } catch (const std::exception& e) {
                                        }
                                        return false;
                                    },
                                    [](InfoHash, const Sp<Value>& o, Sp<Value>& n, const InfoHash&, const SockAddr&) {
                                        try {
                                            return crypto::Certificate(o->data).getPublicKey().getLongId()
                                                   == crypto::Certificate(n->data).getPublicKey().getLongId();
                                        } catch (const std::exception& e) {
                                        }
                                        return false;
                                    }};
 
} // namespace dht