Skip to main content

irontide_engine/
extension.rs

1//! Extension plugin interface for custom BEP 10 extensions.
2//!
3//! Plugins implement [`ExtensionPlugin`] to handle custom extension messages
4//! without modifying torrent internals. Built-in extensions (`ut_metadata`,
5//! `ut_pex`, `lt_trackers`) are hard-coded; plugins receive IDs starting at 10.
6
7use std::collections::BTreeMap;
8use std::net::SocketAddr;
9
10use irontide_bencode::BencodeValue;
11use irontide_core::Id20;
12use irontide_wire::ExtHandshake;
13
14/// A custom BEP 10 extension message handler.
15///
16/// Implement this trait to add custom extension protocol support to a torrent
17/// session. Plugins are registered via `ClientBuilder::add_extension()` and
18/// are immutable after session start.
19///
20/// # Extension ID allocation
21///
22/// Built-in extensions occupy IDs 1-3:
23/// - `ut_metadata` = 1
24/// - `ut_pex` = 2
25/// - `lt_trackers` = 3
26///
27/// Plugins are assigned IDs starting at 10, in registration order.
28///
29/// # Constraints
30///
31/// - Plugins cannot access torrent internals (piece state, peer list).
32/// - Plugins cannot initiate unsolicited messages -- respond only.
33/// - Plugins cannot override built-in extensions.
34/// - Callbacks are synchronous -- spawn tasks internally for async work.
35pub trait ExtensionPlugin: Send + Sync + 'static {
36    /// Extension name for BEP 10 handshake negotiation (e.g. `"ut_comment"`).
37    ///
38    /// This name is advertised in the extension handshake `m` dictionary.
39    fn name(&self) -> &str;
40
41    /// Called when a peer's extension handshake arrives.
42    ///
43    /// Return extra key-value pairs to merge into our handshake response,
44    /// or `None` to add nothing.
45    fn on_handshake(
46        &self,
47        _info_hash: &Id20,
48        _peer_addr: SocketAddr,
49        _handshake: &ExtHandshake,
50    ) -> Option<BTreeMap<String, BencodeValue>> {
51        None
52    }
53
54    /// Called when an extension message for this plugin arrives.
55    ///
56    /// Return an optional response payload to send back to the peer.
57    fn on_message(
58        &self,
59        _info_hash: &Id20,
60        _peer_addr: SocketAddr,
61        _payload: &[u8],
62    ) -> Option<Vec<u8>> {
63        None
64    }
65
66    /// Peer connected (after BT handshake, before extension handshake).
67    fn on_peer_connected(&self, _info_hash: &Id20, _peer_addr: SocketAddr) {}
68
69    /// Peer disconnected.
70    fn on_peer_disconnected(&self, _info_hash: &Id20, _peer_addr: SocketAddr) {}
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76
77    /// A minimal test plugin that echoes messages back.
78    struct EchoPlugin;
79
80    impl ExtensionPlugin for EchoPlugin {
81        fn name(&self) -> &'static str {
82            "ut_echo"
83        }
84
85        fn on_message(
86            &self,
87            _info_hash: &Id20,
88            _peer_addr: SocketAddr,
89            payload: &[u8],
90        ) -> Option<Vec<u8>> {
91            Some(payload.to_vec())
92        }
93    }
94
95    #[test]
96    fn plugin_as_trait_object() {
97        let plugin: Box<dyn ExtensionPlugin> = Box::new(EchoPlugin);
98        assert_eq!(plugin.name(), "ut_echo");
99    }
100
101    #[test]
102    fn plugin_echo_response() {
103        let plugin = EchoPlugin;
104        let info_hash = Id20::from_hex("aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d").unwrap();
105        let addr: SocketAddr = "127.0.0.1:6881".parse().unwrap();
106        let response = plugin.on_message(&info_hash, addr, b"hello");
107        assert_eq!(response, Some(b"hello".to_vec()));
108    }
109
110    #[test]
111    fn default_lifecycle_hooks_are_noops() {
112        let plugin = EchoPlugin;
113        let info_hash = Id20::from_hex("aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d").unwrap();
114        let addr: SocketAddr = "127.0.0.1:6881".parse().unwrap();
115
116        // These should not panic
117        plugin.on_peer_connected(&info_hash, addr);
118        plugin.on_peer_disconnected(&info_hash, addr);
119
120        // Default handshake returns None
121        let hs = ExtHandshake::new();
122        assert!(plugin.on_handshake(&info_hash, addr, &hs).is_none());
123    }
124
125    #[test]
126    fn plugin_vec_in_arc() {
127        use std::sync::Arc;
128        let plugins: Arc<Vec<Box<dyn ExtensionPlugin>>> = Arc::new(vec![Box::new(EchoPlugin)]);
129        assert_eq!(plugins.len(), 1);
130        assert_eq!(plugins[0].name(), "ut_echo");
131    }
132}