diff --git a/include/libtorrent/aux_/session_impl.hpp b/include/libtorrent/aux_/session_impl.hpp index 05eeb3307..dccfb86c0 100644 --- a/include/libtorrent/aux_/session_impl.hpp +++ b/include/libtorrent/aux_/session_impl.hpp @@ -292,6 +292,7 @@ namespace libtorrent void add_dht_router(std::pair const& node); void set_dht_settings(dht_settings const& s); dht_settings const& get_dht_settings() const { return m_dht_settings; } + void set_dht_storage(dht::dht_storage_constructor_type sc); void start_dht(); void stop_dht(); void start_dht(entry const& startup_state); @@ -994,6 +995,7 @@ namespace libtorrent #ifndef TORRENT_DISABLE_DHT boost::shared_ptr m_dht; dht_settings m_dht_settings; + dht::dht_storage_constructor_type m_dht_storage_constructor; // these are used when starting the DHT // (and bootstrapping it), and then erased diff --git a/include/libtorrent/kademlia/dht_storage.hpp b/include/libtorrent/kademlia/dht_storage.hpp index c6edecdf7..8d948a5db 100644 --- a/include/libtorrent/kademlia/dht_storage.hpp +++ b/include/libtorrent/kademlia/dht_storage.hpp @@ -30,22 +30,20 @@ POSSIBILITY OF SUCH DAMAGE. */ -#ifndef DHT_STORAGE_HPP -#define DHT_STORAGE_HPP +#ifndef TORRENT_DHT_STORAGE_HPP +#define TORRENT_DHT_STORAGE_HPP -#include -#include -#include -#include +#include -#include -#include #include #include #include -#include -#include -#include + +namespace libtorrent +{ + struct dht_settings; + class entry; +} namespace libtorrent { namespace dht @@ -200,9 +198,12 @@ namespace dht virtual ~dht_storage_interface() {} }; + typedef boost::function dht_storage_constructor_type; + TORRENT_EXPORT dht_storage_interface* dht_default_storage_constructor(sha1_hash const& id , dht_settings const& settings); } } // namespace libtorrent::dht -#endif //DHT_STORAGE_HPP +#endif //TORRENT_DHT_STORAGE_HPP diff --git a/include/libtorrent/kademlia/dht_tracker.hpp b/include/libtorrent/kademlia/dht_tracker.hpp index c57f6857e..004901b30 100644 --- a/include/libtorrent/kademlia/dht_tracker.hpp +++ b/include/libtorrent/kademlia/dht_tracker.hpp @@ -72,7 +72,9 @@ namespace libtorrent { namespace dht , boost::enable_shared_from_this { dht_tracker(dht_observer* observer, rate_limited_udp_socket& sock - , dht_settings const& settings, counters& cnt, entry const* state = 0); + , dht_settings const& settings, counters& cnt + , dht_storage_constructor_type storage_constructor + , entry const* state = 0); virtual ~dht_tracker(); void start(entry const& bootstrap diff --git a/include/libtorrent/kademlia/node.hpp b/include/libtorrent/kademlia/node.hpp index 1ef2a3b4a..625d476a1 100644 --- a/include/libtorrent/kademlia/node.hpp +++ b/include/libtorrent/kademlia/node.hpp @@ -129,7 +129,8 @@ class TORRENT_EXTRA_EXPORT node : boost::noncopyable public: node(udp_socket_interface* sock , libtorrent::dht_settings const& settings, node_id nid - , dht_observer* observer, counters& cnt); + , dht_observer* observer, counters& cnt + , dht_storage_constructor_type storage_constructor = dht_default_storage_constructor); virtual ~node(); diff --git a/include/libtorrent/session_handle.hpp b/include/libtorrent/session_handle.hpp index 609b182c8..4386eddce 100644 --- a/include/libtorrent/session_handle.hpp +++ b/include/libtorrent/session_handle.hpp @@ -43,6 +43,8 @@ POSSIBILITY OF SUCH DAMAGE. #include "libtorrent/peer_class_type_filter.hpp" #include "libtorrent/session_settings.hpp" +#include "libtorrent/kademlia/dht_storage.hpp" + #ifndef TORRENT_NO_DEPRECATE #include "libtorrent/rss.hpp" #endif @@ -338,7 +340,7 @@ namespace libtorrent void stop_dht(); #endif - // ``set_dht_settings`` sets some parameters availavle to the dht node. + // ``set_dht_settings`` sets some parameters available to the dht node. // See dht_settings for more information. // // ``is_dht_running()`` returns true if the DHT support has been started @@ -350,6 +352,20 @@ namespace libtorrent bool is_dht_running() const; dht_settings get_dht_settings() const; + // ``set_dht_storage`` set a dht custom storage constructor function + // to be used internally when the dht is created. + // + // Since the dht storage is a critical component for the dht behavior, + // this function will only be effective the next time the dht is started. + // If you never touch this feature, a default map-memory based storage + // is used. + // + // If you want to make sure the dht is initially created with your + // custom storage, create a session with the setting + // ``settings_pack::enable_dht`` to false, set your constructor function + // and call ``apply_settings`` with ``settings_pack::enable_dht`` to true. + void set_dht_storage(dht::dht_storage_constructor_type sc); + // ``add_dht_node`` takes a host name and port pair. That endpoint will be // pinged, and if a valid DHT reply is received, the node will be added to // the routing table. diff --git a/src/kademlia/dht_tracker.cpp b/src/kademlia/dht_tracker.cpp index 498ef872c..0b0826bbd 100644 --- a/src/kademlia/dht_tracker.cpp +++ b/src/kademlia/dht_tracker.cpp @@ -98,9 +98,10 @@ namespace libtorrent { namespace dht , rate_limited_udp_socket& sock , dht_settings const& settings , counters& cnt + , dht_storage_constructor_type storage_constructor , entry const* state) : m_counters(cnt) - , m_dht(this, settings, extract_node_id(state), observer, cnt) + , m_dht(this, settings, extract_node_id(state), observer, cnt, storage_constructor) , m_sock(sock) , m_log(observer) , m_last_new_key(clock_type::now() - minutes(int(key_refresh))) diff --git a/src/kademlia/node.cpp b/src/kademlia/node.cpp index c976c4ca3..0c3461537 100644 --- a/src/kademlia/node.cpp +++ b/src/kademlia/node.cpp @@ -88,7 +88,8 @@ node_id calculate_node_id(node_id const& nid, dht_observer* observer) node::node(udp_socket_interface* sock , dht_settings const& settings, node_id nid , dht_observer* observer - , struct counters& cnt) + , struct counters& cnt + , dht_storage_constructor_type storage_constructor) : m_settings(settings) , m_id(calculate_node_id(nid, observer)) , m_table(m_id, 8, settings, observer) @@ -98,7 +99,7 @@ node::node(udp_socket_interface* sock , m_last_self_refresh(min_time()) , m_sock(sock) , m_counters(cnt) - , m_storage(dht_default_storage_constructor(m_id, m_settings)) + , m_storage(storage_constructor(m_id, m_settings)) { m_secret[0] = random(); m_secret[1] = random(); diff --git a/src/session_handle.cpp b/src/session_handle.cpp index ab7b67d10..ea76f96bb 100644 --- a/src/session_handle.cpp +++ b/src/session_handle.cpp @@ -351,6 +351,15 @@ namespace libtorrent #endif } + void session_handle::set_dht_storage(dht::dht_storage_constructor_type sc) + { +#ifndef TORRENT_DISABLE_DHT + TORRENT_ASYNC_CALL1(set_dht_storage, sc); +#else + TORRENT_UNUSED(sc); +#endif + } + void session_handle::add_dht_node(std::pair const& node) { #ifndef TORRENT_DISABLE_DHT diff --git a/src/session_impl.cpp b/src/session_impl.cpp index 69963b3c5..378882b83 100644 --- a/src/session_impl.cpp +++ b/src/session_impl.cpp @@ -403,6 +403,7 @@ namespace aux { , m_last_auto_manage(m_created) , m_next_port(0) #ifndef TORRENT_DISABLE_DHT + , m_dht_storage_constructor(dht::dht_default_storage_constructor) , m_dht_announce_timer(m_io_service) , m_dht_interval_update_torrents(0) , m_outstanding_router_lookups(0) @@ -5483,7 +5484,9 @@ retry: m_dht = boost::make_shared(static_cast(this) , boost::ref(m_udp_socket), boost::cref(m_dht_settings) - , boost::ref(m_stats_counters), &startup_state); + , boost::ref(m_stats_counters) + , m_dht_storage_constructor + , &startup_state); for (std::vector::iterator i = m_dht_router_nodes.begin() , end(m_dht_router_nodes.end()); i != end; ++i) @@ -5516,6 +5519,11 @@ retry: m_dht_settings = settings; } + void session_impl::set_dht_storage(dht::dht_storage_constructor_type sc) + { + m_dht_storage_constructor = sc; + } + #ifndef TORRENT_NO_DEPRECATE entry session_impl::dht_state() const { diff --git a/test/test_dht_storage.cpp b/test/test_dht_storage.cpp index 2e8ea4ff9..64b7c02a0 100644 --- a/test/test_dht_storage.cpp +++ b/test/test_dht_storage.cpp @@ -55,6 +55,7 @@ POSSIBILITY OF SUCH DAMAGE. using namespace libtorrent; using namespace libtorrent::dht; +namespace lt = libtorrent; namespace { @@ -71,6 +72,15 @@ namespace from_hex(s, 40, (char *) &ret[0]); return ret; } + + bool g_storage_constructor_invoked = false; + + dht_storage_interface* dht_custom_storage_constructor(sha1_hash const& id + , dht_settings const& settings) + { + g_storage_constructor_invoked = true; + return dht_default_storage_constructor(id, settings); + } } TORRENT_TEST(dht_storage) @@ -177,4 +187,50 @@ TORRENT_TEST(dht_storage_counters) TEST_EQUAL(s->counters().mutable_data, 1); } +TORRENT_TEST(dht_storage_set_custom) +{ + g_storage_constructor_invoked = false; + settings_pack p; + p.set_bool(settings_pack::enable_dht, false); + lt::session ses(p); + + bool r = ses.is_dht_running(); + TEST_CHECK(!r); + + ses.set_dht_storage(dht_custom_storage_constructor); + + p.set_bool(settings_pack::enable_dht, true); + ses.apply_settings(p); // async with dispatch + r = ses.is_dht_running(); + TEST_CHECK(r); + TEST_EQUAL(g_storage_constructor_invoked, true); +} + +TORRENT_TEST(dht_storage_default_set_custom) +{ + g_storage_constructor_invoked = false; + settings_pack p; + p.set_bool(settings_pack::enable_dht, true); + lt::session ses(p); + + bool r = ses.is_dht_running(); + TEST_CHECK(r); + + ses.set_dht_storage(dht_custom_storage_constructor); + + p.set_bool(settings_pack::enable_dht, false); + ses.apply_settings(p); // async with dispatch + r = ses.is_dht_running(); + TEST_CHECK(!r); + TEST_EQUAL(g_storage_constructor_invoked, false); + + ses.set_dht_storage(dht_custom_storage_constructor); + + p.set_bool(settings_pack::enable_dht, true); + ses.apply_settings(p); // async with dispatch + r = ses.is_dht_running(); + TEST_CHECK(r); + TEST_EQUAL(g_storage_constructor_invoked, true); +} + #endif