irontide_engine/
extension.rs1use std::collections::BTreeMap;
8use std::net::SocketAddr;
9
10use irontide_bencode::BencodeValue;
11use irontide_core::Id20;
12use irontide_wire::ExtHandshake;
13
14pub trait ExtensionPlugin: Send + Sync + 'static {
36 fn name(&self) -> &str;
40
41 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 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 fn on_peer_connected(&self, _info_hash: &Id20, _peer_addr: SocketAddr) {}
68
69 fn on_peer_disconnected(&self, _info_hash: &Id20, _peer_addr: SocketAddr) {}
71}
72
73#[cfg(test)]
74mod tests {
75 use super::*;
76
77 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 plugin.on_peer_connected(&info_hash, addr);
118 plugin.on_peer_disconnected(&info_hash, addr);
119
120 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}