webtorrent/
extensions.rs

1use crate::error::Result;
2use bytes::Bytes;
3use std::collections::HashMap;
4
5/// BitTorrent extension protocol handler
6pub struct ExtensionProtocol {
7    extensions: HashMap<u8, ExtensionInfo>,
8    #[allow(dead_code)]
9    handshake: Option<Bytes>,
10}
11
12#[derive(Debug, Clone)]
13pub struct ExtensionInfo {
14    pub id: u8,
15    pub name: String,
16    pub metadata_size: Option<usize>,
17}
18
19/// ut_metadata extension
20pub struct UtMetadata {
21    metadata: Option<Bytes>,
22    metadata_size: Option<usize>,
23}
24
25impl UtMetadata {
26    pub fn new() -> Self {
27        Self {
28            metadata: None,
29            metadata_size: None,
30        }
31    }
32
33    pub fn set_metadata(&mut self, metadata: Bytes) {
34        let len = metadata.len();
35        self.metadata = Some(metadata);
36        self.metadata_size = Some(len);
37    }
38
39    pub fn get_metadata(&self) -> Option<&Bytes> {
40        self.metadata.as_ref()
41    }
42
43    pub fn fetch(&self) -> Result<()> {
44        // Request metadata from peer
45        Ok(())
46    }
47
48    pub fn handle_request(&self, piece: usize) -> Result<Option<Bytes>> {
49        if let Some(ref metadata) = self.metadata {
50            let piece_size = 16 * 1024; // 16 KB per piece
51            let start = piece * piece_size;
52            let end = (start + piece_size).min(metadata.len());
53            if start < metadata.len() {
54                return Ok(Some(metadata.slice(start..end)));
55            }
56        }
57        Ok(None)
58    }
59}
60
61/// ut_pex extension (Peer Exchange)
62pub struct UtPex {
63    added: Vec<(String, u16)>,
64    dropped: Vec<(String, u16)>,
65}
66
67impl UtPex {
68    pub fn new() -> Self {
69        Self {
70            added: Vec::new(),
71            dropped: Vec::new(),
72        }
73    }
74
75    pub fn add_peer(&mut self, ip: String, port: u16) {
76        self.added.push((ip, port));
77    }
78
79    pub fn drop_peer(&mut self, ip: String, port: u16) {
80        self.dropped.push((ip, port));
81    }
82
83    pub fn get_added(&self) -> &[(String, u16)] {
84        &self.added
85    }
86
87    pub fn get_dropped(&self) -> &[(String, u16)] {
88        &self.dropped
89    }
90
91    pub fn reset(&mut self) {
92        self.added.clear();
93        self.dropped.clear();
94    }
95
96    pub fn encode(&self) -> Bytes {
97        // Encode peer list in compact format
98        let mut buf = Vec::new();
99        for (ip, port) in &self.added {
100            if let Ok(ip_bytes) = ip.parse::<std::net::Ipv4Addr>() {
101                buf.extend_from_slice(&ip_bytes.octets());
102                buf.extend_from_slice(&port.to_be_bytes());
103            }
104        }
105        Bytes::from(buf)
106    }
107
108    pub fn decode(data: &[u8]) -> Result<Vec<(String, u16)>> {
109        let mut peers = Vec::new();
110        for chunk in data.chunks_exact(6) {
111            let ip = format!("{}.{}.{}.{}", chunk[0], chunk[1], chunk[2], chunk[3]);
112            let port = u16::from_be_bytes([chunk[4], chunk[5]]);
113            peers.push((ip, port));
114        }
115        Ok(peers)
116    }
117}
118
119impl ExtensionProtocol {
120    pub fn new() -> Self {
121        Self {
122            extensions: HashMap::new(),
123            handshake: None,
124        }
125    }
126
127    pub fn register_extension(&mut self, id: u8, name: String) {
128        self.extensions.insert(id, ExtensionInfo {
129            id,
130            name,
131            metadata_size: None,
132        });
133    }
134
135    pub fn get_extension_id(&self, name: &str) -> Option<u8> {
136        self.extensions.values()
137            .find(|ext| ext.name == name)
138            .map(|ext| ext.id)
139    }
140}
141
webtorrent/extensions.rs

webtorrent/
extensions.rs
