mtorrent_core/pwp/
message.rs

1use bitvec::prelude::*;
2use bytes::{Buf, BufMut};
3use derive_more::Display;
4use mtorrent_utils::{benc, net};
5use std::collections::{BTreeMap, HashMap, HashSet};
6use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
7use std::{fmt, io};
8
9/// Bitfield from the bitfield message.
10pub type Bitfield = BitVec<u8, Msb0>;
11
12#[derive(Debug)]
13pub(super) enum PeerMessage {
14    KeepAlive,
15    Choke,
16    Unchoke,
17    Interested,
18    NotInterested,
19    Have {
20        piece_index: u32,
21    },
22    Bitfield {
23        bitfield: Bitfield,
24    },
25    Request {
26        index: u32,
27        begin: u32,
28        length: u32,
29    },
30    Piece {
31        index: u32,
32        begin: u32,
33        block: Vec<u8>,
34    },
35    Cancel {
36        index: u32,
37        begin: u32,
38        length: u32,
39    },
40    DhtPort {
41        listen_port: u16,
42    },
43    Extended {
44        id: u8,
45        data: Vec<u8>,
46    },
47}
48
49impl PeerMessage {
50    /// Decode message after the first 4 bytes. `src` must contain at least `msg_len` bytes.
51    pub(super) fn decode_body<B: Buf>(msg_len: usize, src: &mut B) -> io::Result<Self> {
52        fn invalid_data_err(e: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> io::Error {
53            io::Error::new(io::ErrorKind::InvalidData, e)
54        }
55
56        if msg_len > src.remaining() {
57            return Err(invalid_data_err(format!(
58                "received incomplete msg ({} instead of {})",
59                src.remaining(),
60                msg_len
61            )));
62        }
63
64        if msg_len == 0 {
65            return Ok(Self::KeepAlive);
66        }
67
68        let id = src.get_u8();
69        match id {
70            ID_CHOKE => Ok(Self::Choke),
71            ID_UNCHOKE => Ok(Self::Unchoke),
72            ID_INTERESTED => Ok(Self::Interested),
73            ID_NOT_INTERESTED => Ok(Self::NotInterested),
74            ID_HAVE => Ok(Self::Have {
75                piece_index: src
76                    .try_get_u32()
77                    .map_err(|_| invalid_data_err("Have is too short"))?,
78            }),
79            ID_BITFIELD => {
80                let mut bitfield_bytes = vec![0u8; msg_len - 1];
81                src.copy_to_slice(&mut bitfield_bytes);
82                Ok(Self::Bitfield {
83                    bitfield: BitVec::try_from_vec(bitfield_bytes)
84                        .map_err(|_| invalid_data_err("Bitfield is too long"))?,
85                })
86            }
87            ID_REQUEST => {
88                if src.remaining() < 12 {
89                    Err(invalid_data_err("Request is too short"))
90                } else {
91                    let index = src.get_u32();
92                    let begin = src.get_u32();
93                    let length = src.get_u32();
94                    Ok(Self::Request {
95                        index,
96                        begin,
97                        length,
98                    })
99                }
100            }
101            ID_CANCEL => {
102                if src.remaining() < 12 {
103                    Err(invalid_data_err("Cancel is too short"))
104                } else {
105                    let index = src.get_u32();
106                    let begin = src.get_u32();
107                    let length = src.get_u32();
108                    Ok(Self::Cancel {
109                        index,
110                        begin,
111                        length,
112                    })
113                }
114            }
115            ID_PIECE => {
116                if msg_len < 9 {
117                    Err(invalid_data_err("invalid Piece len"))
118                } else {
119                    let index = src.get_u32();
120                    let begin = src.get_u32();
121                    let block_len = msg_len - 9;
122                    let data = src
123                        .chunk()
124                        .get(..block_len)
125                        .ok_or_else(|| invalid_data_err("Piece is shorter than expected"))?;
126                    let mut block = Vec::with_capacity(block_len);
127                    block.extend_from_slice(data);
128                    Ok(Self::Piece {
129                        index,
130                        begin,
131                        block,
132                    })
133                }
134            }
135            ID_PORT => Ok(Self::DhtPort {
136                listen_port: src
137                    .try_get_u16()
138                    .map_err(|_| invalid_data_err("DhtPort is too short"))?,
139            }),
140            ID_EXTENDED => {
141                if msg_len < 2 {
142                    Err(invalid_data_err("invalid Extended len"))
143                } else {
144                    let id = src.get_u8();
145                    let data_len = msg_len - 2;
146                    let data = src.chunk().get(..data_len).ok_or_else(|| {
147                        invalid_data_err("Extended data is shorter than expected")
148                    })?;
149                    let mut ext_data = Vec::with_capacity(data_len);
150                    ext_data.extend_from_slice(data);
151                    Ok(Self::Extended { id, data: ext_data })
152                }
153            }
154            _ => Err(io::Error::new(
155                io::ErrorKind::InvalidInput,
156                format!("Unknown message type: {id}"),
157            )),
158        }
159    }
160
161    /// Write message to `dest`. The destination buffer must be big enough.
162    pub(super) fn encode<B: BufMut>(&self, dest: &mut B) -> io::Result<()> {
163        let len = self.get_length();
164        if len + 4 > dest.remaining_mut() {
165            return Err(io::Error::new(
166                io::ErrorKind::OutOfMemory,
167                format!(
168                    "buffer too short to encode message ({} < {})",
169                    dest.remaining_mut(),
170                    len + 4
171                ),
172            ));
173        }
174
175        dest.put_u32(len as u32);
176
177        match self {
178            Self::KeepAlive => (),
179            Self::Choke => {
180                dest.put_u8(ID_CHOKE);
181            }
182            Self::Unchoke => {
183                dest.put_u8(ID_UNCHOKE);
184            }
185            Self::Interested => {
186                dest.put_u8(ID_INTERESTED);
187            }
188            Self::NotInterested => {
189                dest.put_u8(ID_NOT_INTERESTED);
190            }
191            Self::Have { piece_index } => {
192                dest.put_u8(ID_HAVE);
193                dest.put_u32(*piece_index);
194            }
195            Self::Bitfield { bitfield } => {
196                dest.put_u8(ID_BITFIELD);
197                dest.put_slice(bitfield.as_raw_slice());
198            }
199            Self::Request {
200                index,
201                begin,
202                length,
203            } => {
204                dest.put_u8(ID_REQUEST);
205                dest.put_u32(*index);
206                dest.put_u32(*begin);
207                dest.put_u32(*length);
208            }
209            Self::Piece {
210                index,
211                begin,
212                block,
213            } => {
214                dest.put_u8(ID_PIECE);
215                dest.put_u32(*index);
216                dest.put_u32(*begin);
217                dest.put_slice(block);
218            }
219            Self::Cancel {
220                index,
221                begin,
222                length,
223            } => {
224                dest.put_u8(ID_CANCEL);
225                dest.put_u32(*index);
226                dest.put_u32(*begin);
227                dest.put_u32(*length);
228            }
229            Self::DhtPort { listen_port } => {
230                dest.put_u8(ID_PORT);
231                dest.put_u16(*listen_port);
232            }
233            Self::Extended { id, data } => {
234                dest.put_u8(ID_EXTENDED);
235                dest.put_u8(*id);
236                dest.put_slice(data);
237            }
238        };
239        Ok(())
240    }
241
242    fn get_length(&self) -> usize {
243        use PeerMessage::*;
244
245        match self {
246            KeepAlive => 0,
247            Choke | Unchoke | Interested | NotInterested => 1,
248            Have { .. } => 5,
249            Bitfield { bitfield } => 1 + bitfield.as_raw_slice().len(),
250            Request { .. } | Cancel { .. } => 13,
251            Piece { block, .. } => 9 + block.len(),
252            DhtPort { .. } => 3,
253            Extended { data, .. } => 2 + data.len(),
254        }
255    }
256}
257
258const ID_CHOKE: u8 = 0;
259const ID_UNCHOKE: u8 = 1;
260const ID_INTERESTED: u8 = 2;
261const ID_NOT_INTERESTED: u8 = 3;
262const ID_HAVE: u8 = 4;
263const ID_BITFIELD: u8 = 5;
264const ID_REQUEST: u8 = 6;
265const ID_PIECE: u8 = 7;
266const ID_CANCEL: u8 = 8;
267const ID_PORT: u8 = 9;
268const ID_EXTENDED: u8 = 20;
269
270// ------
271
272/// Information about a block of data.
273#[derive(Clone, Eq, PartialEq, Hash, Debug, Display)]
274#[display("ind={piece_index} off={in_piece_offset} len={block_length}")]
275pub struct BlockInfo {
276    pub piece_index: usize,
277    pub in_piece_offset: usize,
278    pub block_length: usize,
279}
280
281/// Messages pertaining to the upload of data by the peer sending the messages.
282#[derive(Clone, Eq, PartialEq, Debug)]
283pub enum UploaderMessage {
284    Choke,
285    Unchoke,
286    Have {
287        piece_index: usize,
288    },
289    Bitfield(Bitfield),
290    Block(BlockInfo, Vec<u8>),
291}
292
293/// Messages pertaining to the download of data by the peer sending the messages.
294#[derive(Clone, Eq, PartialEq, Debug)]
295pub enum DownloaderMessage {
296    Interested,
297    NotInterested,
298    Request(BlockInfo),
299    Cancel(BlockInfo),
300}
301
302/// Types of extensions for the [extension protocol](https://www.bittorrent.org/beps/bep_0010.html).
303#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
304pub enum Extension {
305    Metadata,
306    PeerExchange,
307}
308
309/// Parsed extended handshake.
310#[derive(Default, Clone, Eq, PartialEq, Debug)]
311pub struct ExtendedHandshake {
312    pub extensions: HashMap<Extension, u8>,
313    pub listen_port: Option<u16>,
314    pub client_type: Option<String>,
315    pub yourip: Option<IpAddr>,
316    pub ipv4: Option<Ipv4Addr>,
317    pub ipv6: Option<Ipv6Addr>,
318    pub request_limit: Option<usize>,
319    pub metadata_size: Option<usize>,
320}
321
322/// Parsed PEX message.
323#[derive(Default, Clone, Eq, PartialEq, Debug, Display)]
324#[display("added={added:?} dropped={dropped:?}")]
325pub struct PeerExchangeData {
326    pub added: HashSet<SocketAddr>,
327    pub dropped: HashSet<SocketAddr>,
328}
329
330/// Parsed extended message as defined by the [extension protocol](https://www.bittorrent.org/beps/bep_0010.html).
331#[derive(Clone, Eq, PartialEq, Debug)]
332pub enum ExtendedMessage {
333    Handshake(Box<ExtendedHandshake>),
334    MetadataRequest {
335        piece: usize,
336    },
337    MetadataBlock {
338        piece: usize,
339        total_size: usize,
340        data: Vec<u8>,
341    },
342    MetadataReject {
343        piece: usize,
344    },
345    PeerExchange(Box<PeerExchangeData>),
346}
347
348impl fmt::Display for ExtendedHandshake {
349    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
350        write!(f, "ext=[{:?}]", self.extensions)?;
351        if let Some(port) = self.listen_port {
352            write!(f, " port={port}")?;
353        }
354        if let Some(client_type) = self.client_type.as_ref() {
355            write!(f, " client={client_type}")?;
356        }
357        if let Some(yourip) = self.yourip {
358            write!(f, " yourip={yourip}")?;
359        }
360        if let Some(ipv4) = self.ipv4 {
361            write!(f, " ipv4={ipv4}")?;
362        }
363        if let Some(ipv6) = self.ipv6 {
364            write!(f, " ipv6={ipv6}")?;
365        }
366        if let Some(reqq) = self.request_limit {
367            write!(f, " reqq={reqq}")?;
368        }
369        if let Some(metasize) = self.metadata_size {
370            write!(f, " metasize={metasize}")?;
371        }
372        Ok(())
373    }
374}
375
376// ------
377
378impl From<UploaderMessage> for PeerMessage {
379    fn from(msg: UploaderMessage) -> Self {
380        match msg {
381            UploaderMessage::Choke => PeerMessage::Choke,
382            UploaderMessage::Unchoke => PeerMessage::Unchoke,
383            UploaderMessage::Have { piece_index } => PeerMessage::Have {
384                piece_index: piece_index as u32,
385            },
386            UploaderMessage::Bitfield(bitfield) => PeerMessage::Bitfield { bitfield },
387            UploaderMessage::Block(info, data) => PeerMessage::Piece {
388                index: info.piece_index as u32,
389                begin: info.in_piece_offset as u32,
390                block: data,
391            },
392        }
393    }
394}
395
396impl TryFrom<PeerMessage> for UploaderMessage {
397    type Error = PeerMessage;
398
399    fn try_from(msg: PeerMessage) -> Result<Self, Self::Error> {
400        match msg {
401            PeerMessage::Choke => Ok(UploaderMessage::Choke),
402            PeerMessage::Unchoke => Ok(UploaderMessage::Unchoke),
403            PeerMessage::Have { piece_index } => Ok(UploaderMessage::Have {
404                piece_index: piece_index as usize,
405            }),
406            PeerMessage::Bitfield { bitfield } => Ok(UploaderMessage::Bitfield(bitfield)),
407            PeerMessage::Piece {
408                index,
409                begin,
410                block,
411            } => Ok(UploaderMessage::Block(
412                BlockInfo {
413                    piece_index: index as usize,
414                    in_piece_offset: begin as usize,
415                    block_length: block.len(),
416                },
417                block,
418            )),
419            _ => Err(msg),
420        }
421    }
422}
423
424impl fmt::Display for UploaderMessage {
425    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
426        match self {
427            UploaderMessage::Choke => {
428                write!(f, "Choke")
429            }
430            UploaderMessage::Unchoke => {
431                write!(f, "Unchoke")
432            }
433            UploaderMessage::Have { piece_index } => {
434                write!(f, "Have[ind={piece_index}]")
435            }
436            UploaderMessage::Bitfield(bitvec) => {
437                write!(f, "Bitfield[len={}]", bitvec.len())
438            }
439            UploaderMessage::Block(info, _) => {
440                write!(f, "Block[{info}]")
441            }
442        }
443    }
444}
445
446// ------
447
448impl From<DownloaderMessage> for PeerMessage {
449    fn from(msg: DownloaderMessage) -> Self {
450        match msg {
451            DownloaderMessage::Interested => PeerMessage::Interested,
452            DownloaderMessage::NotInterested => PeerMessage::NotInterested,
453            DownloaderMessage::Request(info) => PeerMessage::Request {
454                index: info.piece_index as u32,
455                begin: info.in_piece_offset as u32,
456                length: info.block_length as u32,
457            },
458            DownloaderMessage::Cancel(info) => PeerMessage::Cancel {
459                index: info.piece_index as u32,
460                begin: info.in_piece_offset as u32,
461                length: info.block_length as u32,
462            },
463        }
464    }
465}
466
467impl TryFrom<PeerMessage> for DownloaderMessage {
468    type Error = PeerMessage;
469
470    fn try_from(msg: PeerMessage) -> Result<Self, Self::Error> {
471        match msg {
472            PeerMessage::Interested => Ok(DownloaderMessage::Interested),
473            PeerMessage::NotInterested => Ok(DownloaderMessage::NotInterested),
474            PeerMessage::Request {
475                index,
476                begin,
477                length,
478            } => Ok(DownloaderMessage::Request(BlockInfo {
479                piece_index: index as usize,
480                in_piece_offset: begin as usize,
481                block_length: length as usize,
482            })),
483            PeerMessage::Cancel {
484                index,
485                begin,
486                length,
487            } => Ok(DownloaderMessage::Cancel(BlockInfo {
488                piece_index: index as usize,
489                in_piece_offset: begin as usize,
490                block_length: length as usize,
491            })),
492            _ => Err(msg),
493        }
494    }
495}
496
497impl fmt::Display for DownloaderMessage {
498    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
499        match self {
500            DownloaderMessage::Interested => {
501                write!(f, "Interested")
502            }
503            DownloaderMessage::NotInterested => {
504                write!(f, "NotInterested")
505            }
506            DownloaderMessage::Request(info) => {
507                write!(f, "Request[{info}]")
508            }
509            DownloaderMessage::Cancel(info) => {
510                write!(f, "Cancel[{info}]")
511            }
512        }
513    }
514}
515
516// ------
517
518impl Extension {
519    const NAME_METADATA: &'static str = "ut_metadata";
520    const NAME_PEX: &'static str = "ut_pex";
521
522    const ID_HANDSHAKE: u8 = 0;
523    const ID_METADATA: u8 = 1;
524    const ID_PEX: u8 = 2;
525
526    fn from_name(name: &str) -> Option<Self> {
527        match name {
528            Self::NAME_METADATA => Some(Self::Metadata),
529            Self::NAME_PEX => Some(Self::PeerExchange),
530            _ => None,
531        }
532    }
533    fn name(&self) -> &'static str {
534        match self {
535            Extension::Metadata => Self::NAME_METADATA,
536            Extension::PeerExchange => Self::NAME_PEX,
537        }
538    }
539    /// ID of the extension used locally, i.e. for parsing of incoming messages.
540    pub const fn local_id(&self) -> u8 {
541        match self {
542            Extension::Metadata => Self::ID_METADATA,
543            Extension::PeerExchange => Self::ID_PEX,
544        }
545    }
546}
547
548impl ExtendedHandshake {
549    const KEY_M: &'static str = "m";
550    const KEY_P: &'static str = "p";
551    const KEY_V: &'static str = "v";
552    const KEY_YOURIP: &'static str = "yourip";
553    const KEY_IPV6: &'static str = "ipv6";
554    const KEY_IPV4: &'static str = "ipv4";
555    const KEY_REQQ: &'static str = "reqq";
556    const KEY_METADATA_SIZE: &'static str = "metadata_size";
557
558    fn decode(payload: &[u8]) -> Option<Self> {
559        use mtorrent_utils::benc::Element::{self, *};
560        if let Dictionary(d) = Element::from_bytes(payload).ok()? {
561            let mut root = benc::convert_dictionary(d);
562            let mut ret = Self::default();
563            if let Some(Integer(port)) = root.remove(Self::KEY_P) {
564                ret.listen_port = port.try_into().ok();
565            }
566            if let Some(ByteString(v)) = root.remove(Self::KEY_V) {
567                ret.client_type = String::from_utf8(v).ok();
568            }
569            if let Some(ByteString(ip)) = root.remove(Self::KEY_YOURIP) {
570                if let Some(ipv6_bytes) = ip.get(0..16) {
571                    let mut octets = [0u8; 16];
572                    octets.copy_from_slice(ipv6_bytes);
573                    ret.yourip = Some(IpAddr::V6(Ipv6Addr::from(octets)));
574                } else if let Some(ipv4_bytes) = ip.get(0..4) {
575                    let mut octets = [0u8; 4];
576                    octets.copy_from_slice(ipv4_bytes);
577                    ret.yourip = Some(IpAddr::V4(Ipv4Addr::from(octets)));
578                }
579            }
580            if let Some(ByteString(ipv4)) = root.remove(Self::KEY_IPV4) {
581                ret.ipv4 = ipv4.get(0..4).map(|bytes| {
582                    let mut octets = [0u8; 4];
583                    octets.copy_from_slice(bytes);
584                    Ipv4Addr::from(octets)
585                });
586            }
587            if let Some(ByteString(ipv6)) = root.remove(Self::KEY_IPV6) {
588                ret.ipv6 = ipv6.get(0..16).map(|bytes| {
589                    let mut octets = [0u8; 16];
590                    octets.copy_from_slice(bytes);
591                    Ipv6Addr::from(octets)
592                });
593            }
594            if let Some(Integer(max_requests)) = root.remove(Self::KEY_REQQ) {
595                ret.request_limit = max_requests.try_into().ok();
596            }
597            if let Some(Integer(metasize)) = root.remove(Self::KEY_METADATA_SIZE) {
598                ret.metadata_size = metasize.try_into().ok();
599            }
600            if let Some(Dictionary(d)) = root.remove(Self::KEY_M) {
601                ret.extensions = d
602                    .into_iter()
603                    .filter_map(|(key, value)| match (key, value) {
604                        (ByteString(key), Integer(value)) => {
605                            let extension_name = String::from_utf8(key).ok()?;
606                            let extension = Extension::from_name(&extension_name)?;
607                            let id = u8::try_from(value).ok()?;
608                            Some((extension, id))
609                        }
610                        _ => None,
611                    })
612                    .collect();
613            }
614            Some(ret)
615        } else {
616            None
617        }
618    }
619
620    fn encode(&self) -> Vec<u8> {
621        use mtorrent_utils::benc::Element::{self, *};
622        let root = {
623            let mut tmp = BTreeMap::new();
624            let mut insert = |key: &str, val| {
625                tmp.insert(Element::from(key), val);
626            };
627            if let Some(p) = self.listen_port {
628                insert(Self::KEY_P, Integer(p as i64));
629            }
630            if let Some(v) = self.client_type.as_ref() {
631                insert(Self::KEY_V, Element::from(v.as_str()));
632            }
633            if let Some(ip) = self.yourip {
634                let value = match ip {
635                    IpAddr::V4(ip) => ByteString(ip.octets().into()),
636                    IpAddr::V6(ip) => ByteString(ip.octets().into()),
637                };
638                insert(Self::KEY_YOURIP, value);
639            }
640            if let Some(v4) = self.ipv4 {
641                insert(Self::KEY_IPV4, ByteString(v4.octets().into()));
642            }
643            if let Some(v6) = self.ipv6 {
644                insert(Self::KEY_IPV6, ByteString(v6.octets().into()));
645            }
646            if let Some(reqq) = self.request_limit {
647                insert(Self::KEY_REQQ, Element::from(reqq as i64));
648            }
649            if let Some(metasize) = self.metadata_size {
650                insert(Self::KEY_METADATA_SIZE, Integer(metasize as i64));
651            }
652            if !self.extensions.is_empty() {
653                let m = self
654                    .extensions
655                    .iter()
656                    .map(|(&extension, &id)| {
657                        (Element::from(extension.name()), Element::from(i64::from(id)))
658                    })
659                    .collect::<BTreeMap<Element, Element>>();
660                insert(Self::KEY_M, Dictionary(m));
661            }
662            Dictionary(tmp)
663        };
664        root.encode()
665    }
666}
667
668impl PeerExchangeData {
669    const KEY_ADDED_V4: &'static str = "added";
670    const KEY_ADDED_V6: &'static str = "added6";
671    const KEY_DROPPED_V4: &'static str = "dropped";
672    const KEY_DROPPED_V6: &'static str = "dropped6";
673    const MAX_PEERS: usize = 50;
674
675    fn decode(payload: &[u8]) -> Option<Self> {
676        use benc::Element::{self, *};
677
678        let Dictionary(data) = Element::from_bytes(payload).ok()? else {
679            return None;
680        };
681        let mut root = benc::convert_dictionary(data);
682
683        let mut added = HashSet::new();
684        let mut dropped = HashSet::new();
685
686        if let Some(ByteString(s)) = root.remove(Self::KEY_ADDED_V4) {
687            added.extend(net::SocketAddrV4BytesIter(&s).map(SocketAddr::V4));
688        }
689        if let Some(ByteString(s)) = root.remove(Self::KEY_ADDED_V6) {
690            added.extend(net::SocketAddrV6BytesIter(&s).map(SocketAddr::V6));
691        }
692        if let Some(ByteString(s)) = root.remove(Self::KEY_DROPPED_V4) {
693            dropped.extend(net::SocketAddrV4BytesIter(&s).map(SocketAddr::V4));
694        }
695        if let Some(ByteString(s)) = root.remove(Self::KEY_DROPPED_V6) {
696            dropped.extend(net::SocketAddrV6BytesIter(&s).map(SocketAddr::V6));
697        }
698        Some(PeerExchangeData { added, dropped })
699    }
700
701    fn encode(&self) -> Vec<u8> {
702        use mtorrent_utils::benc::Element::{self, *};
703        let mut root = BTreeMap::<Element, Element>::new();
704
705        let mut added_ipv4 = Vec::new();
706        let mut added_ipv6 = Vec::new();
707        for addr in self.added.iter().take(Self::MAX_PEERS) {
708            match addr {
709                SocketAddr::V4(addr) => {
710                    added_ipv4.extend_from_slice(&addr.ip().octets());
711                    added_ipv4.extend_from_slice(&u16::to_be_bytes(addr.port()));
712                }
713                SocketAddr::V6(addr) => {
714                    added_ipv6.extend_from_slice(&addr.ip().octets());
715                    added_ipv6.extend_from_slice(&u16::to_be_bytes(addr.port()));
716                }
717            }
718        }
719        root.insert(Self::KEY_ADDED_V4.into(), ByteString(added_ipv4));
720        root.insert(Self::KEY_ADDED_V6.into(), ByteString(added_ipv6));
721
722        let mut dropped_ipv4 = Vec::new();
723        let mut dropped_ipv6 = Vec::new();
724        for addr in self.dropped.iter().take(Self::MAX_PEERS) {
725            match addr {
726                SocketAddr::V4(addr) => {
727                    dropped_ipv4.extend_from_slice(&addr.ip().octets());
728                    dropped_ipv4.extend_from_slice(&u16::to_be_bytes(addr.port()));
729                }
730                SocketAddr::V6(addr) => {
731                    dropped_ipv6.extend_from_slice(&addr.ip().octets());
732                    dropped_ipv6.extend_from_slice(&u16::to_be_bytes(addr.port()));
733                }
734            }
735        }
736        root.insert(Self::KEY_DROPPED_V4.into(), ByteString(dropped_ipv4));
737        root.insert(Self::KEY_DROPPED_V6.into(), ByteString(dropped_ipv6));
738
739        Dictionary(root).encode()
740    }
741}
742
743enum MetadataMsg {
744    Request {
745        piece: usize,
746    },
747    Block {
748        piece: usize,
749        total_size: usize,
750        data: Vec<u8>,
751    },
752    Reject {
753        piece: usize,
754    },
755}
756
757impl MetadataMsg {
758    const TYPE_REQUEST: u8 = 0;
759    const TYPE_BLOCK: u8 = 1;
760    const TYPE_REJECT: u8 = 2;
761
762    const KEY_TYPE: &'static str = "msg_type";
763    const KEY_PIECE: &'static str = "piece";
764    const KEY_TOTAL_SIZE: &'static str = "total_size";
765
766    fn decode(mut payload: Vec<u8>) -> Result<Self, Vec<u8>> {
767        use benc::Element::{self, *};
768        let (bencode, bencode_len) = match Element::from_bytes_with_len(&payload) {
769            Ok(b) => b,
770            Err(_) => return Err(payload),
771        };
772        let (msg_type, piece, total_size) = match &bencode {
773            Dictionary(d) => {
774                let msg_type = d.get(&Element::from(Self::KEY_TYPE)).and_then(|b| match b {
775                    Integer(msg_type) => u8::try_from(*msg_type).ok(),
776                    _ => None,
777                });
778                let piece = d.get(&Element::from(Self::KEY_PIECE)).and_then(|b| match b {
779                    Integer(piece) => usize::try_from(*piece).ok(),
780                    _ => None,
781                });
782                let total_size =
783                    d.get(&Element::from(Self::KEY_TOTAL_SIZE)).and_then(|e| match e {
784                        Element::Integer(total_size) => usize::try_from(*total_size).ok(),
785                        _ => None,
786                    });
787                match (msg_type, piece) {
788                    (Some(msg_type), Some(piece)) => (msg_type, piece, total_size),
789                    _ => return Err(payload),
790                }
791            }
792            _ => return Err(payload),
793        };
794        match (msg_type, total_size) {
795            (Self::TYPE_REQUEST, _) => Ok(Self::Request { piece }),
796            (Self::TYPE_REJECT, _) => Ok(Self::Reject { piece }),
797            (Self::TYPE_BLOCK, Some(total_size)) => {
798                let header_len = bencode_len;
799                let total_len = payload.len();
800                // remove bencode from the front and retain only the data that follows
801                payload.copy_within(header_len..total_len, 0);
802                payload.truncate(total_len - header_len);
803                Ok(Self::Block {
804                    piece,
805                    total_size,
806                    data: payload,
807                })
808            }
809            _ => Err(payload),
810        }
811    }
812
813    fn encode(self) -> Vec<u8> {
814        use benc::Element::{self, *};
815        let mut root = BTreeMap::<Element, Element>::new();
816        match self {
817            MetadataMsg::Request { piece } => {
818                root.insert(Self::KEY_TYPE.into(), Integer(Self::TYPE_REQUEST.into()));
819                root.insert(Self::KEY_PIECE.into(), Integer(piece as i64));
820                Dictionary(root).encode()
821            }
822            MetadataMsg::Reject { piece } => {
823                root.insert(Self::KEY_TYPE.into(), Integer(Self::TYPE_REJECT.into()));
824                root.insert(Self::KEY_PIECE.into(), Integer(piece as i64));
825                Dictionary(root).encode()
826            }
827            MetadataMsg::Block {
828                piece,
829                total_size,
830                mut data,
831            } => {
832                root.insert(Self::KEY_TYPE.into(), Integer(Self::TYPE_BLOCK.into()));
833                root.insert(Self::KEY_PIECE.into(), Integer(piece as i64));
834                root.insert(Self::KEY_TOTAL_SIZE.into(), Integer(total_size as i64));
835                let header = Dictionary(root).encode();
836                let header_len = header.len();
837                let data_len = data.len();
838                // add header in front of data
839                data.resize(data_len + header_len, 0u8);
840                data.copy_within(0..data_len, header_len);
841                data[0..header_len].copy_from_slice(&header);
842                data
843            }
844        }
845    }
846}
847
848impl From<(ExtendedMessage, u8)> for PeerMessage {
849    fn from((extmsg, id): (ExtendedMessage, u8)) -> Self {
850        match extmsg {
851            ExtendedMessage::Handshake(hs) => PeerMessage::Extended {
852                id: Extension::ID_HANDSHAKE,
853                data: hs.encode(),
854            },
855            ExtendedMessage::PeerExchange(pex) => PeerMessage::Extended {
856                id,
857                data: pex.encode(),
858            },
859            ExtendedMessage::MetadataRequest { piece } => {
860                let msg = MetadataMsg::Request { piece };
861                PeerMessage::Extended {
862                    id,
863                    data: msg.encode(),
864                }
865            }
866            ExtendedMessage::MetadataReject { piece } => {
867                let msg = MetadataMsg::Reject { piece };
868                PeerMessage::Extended {
869                    id,
870                    data: msg.encode(),
871                }
872            }
873            ExtendedMessage::MetadataBlock {
874                piece,
875                total_size,
876                data,
877            } => {
878                let msg = MetadataMsg::Block {
879                    piece,
880                    total_size,
881                    data,
882                };
883                PeerMessage::Extended {
884                    id,
885                    data: msg.encode(),
886                }
887            }
888        }
889    }
890}
891
892impl TryFrom<PeerMessage> for ExtendedMessage {
893    type Error = PeerMessage;
894
895    fn try_from(msg: PeerMessage) -> Result<Self, Self::Error> {
896        match msg {
897            PeerMessage::Extended {
898                id: Extension::ID_HANDSHAKE,
899                ref data,
900            } => {
901                let handshake = ExtendedHandshake::decode(data).ok_or(msg)?;
902                Ok(ExtendedMessage::Handshake(Box::new(handshake)))
903            }
904            PeerMessage::Extended {
905                id: Extension::ID_PEX,
906                ref data,
907            } => {
908                let peer_data = PeerExchangeData::decode(data).ok_or(msg)?;
909                Ok(ExtendedMessage::PeerExchange(Box::new(peer_data)))
910            }
911            PeerMessage::Extended {
912                id: Extension::ID_METADATA,
913                data,
914            } => match MetadataMsg::decode(data) {
915                Ok(MetadataMsg::Request { piece }) => {
916                    Ok(ExtendedMessage::MetadataRequest { piece })
917                }
918                Ok(MetadataMsg::Reject { piece }) => Ok(ExtendedMessage::MetadataReject { piece }),
919                Ok(MetadataMsg::Block {
920                    piece,
921                    total_size,
922                    data,
923                }) => Ok(ExtendedMessage::MetadataBlock {
924                    piece,
925                    total_size,
926                    data,
927                }),
928                Err(data) => Err(PeerMessage::Extended {
929                    id: Extension::ID_METADATA,
930                    data,
931                }),
932            },
933            _ => Err(msg),
934        }
935    }
936}
937
938impl fmt::Display for ExtendedMessage {
939    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
940        match self {
941            ExtendedMessage::Handshake(hs) => {
942                write!(f, "ExtendedHandshake[{hs}]")
943            }
944            ExtendedMessage::MetadataRequest { piece } => {
945                write!(f, "MetadataRequest[piece={piece}]")
946            }
947            ExtendedMessage::MetadataBlock { piece, data, .. } => {
948                write!(f, "MetadataBlock[piece={} len={}]", piece, data.len())
949            }
950            ExtendedMessage::MetadataReject { piece } => {
951                write!(f, "MetadataReject[piece={piece}]")
952            }
953            ExtendedMessage::PeerExchange(pex) => {
954                write!(f, "PEX[{pex}]")
955            }
956        }
957    }
958}
959
960#[cfg(test)]
961mod tests {
962    use super::*;
963
964    #[test]
965    fn test_bitfield_is_parsed_correctly() {
966        let bits = b"\x82";
967        let msg = UploaderMessage::Bitfield(BitVec::from_slice(bits));
968
969        let bitfield = if let UploaderMessage::Bitfield(bitfield) = msg {
970            bitfield
971        } else {
972            panic!()
973        };
974
975        assert!(bitfield[0]);
976        assert!(!bitfield[1]);
977        assert!(!bitfield[2]);
978        assert!(!bitfield[3]);
979
980        assert!(!bitfield[4]);
981        assert!(!bitfield[5]);
982        assert!(bitfield[6]);
983        assert!(!bitfield[7]);
984    }
985
986    #[test]
987    fn test_handshake_payload_is_parsed_correctly() {
988        let payload =
989            Vec::from(b"d1:md11:ut_metadatai1e6:ut_pexi2ee1:pi6881e1:v13:\xc2\xb5Torrent 1.2e");
990
991        let parsed = ExtendedHandshake::decode(&payload).unwrap();
992        assert_eq!(
993            HashMap::from([(Extension::Metadata, 1), (Extension::PeerExchange, 2)]),
994            parsed.extensions
995        );
996        assert_eq!(6881, parsed.listen_port.unwrap());
997        assert_eq!("µTorrent 1.2", parsed.client_type.as_deref().unwrap());
998        assert!(parsed.ipv4.is_none());
999        assert!(parsed.ipv6.is_none());
1000        assert!(parsed.yourip.is_none());
1001        assert!(parsed.request_limit.is_none());
1002        assert!(parsed.metadata_size.is_none());
1003
1004        let data = parsed.encode();
1005        assert_eq!(payload, data);
1006    }
1007
1008    #[test]
1009    fn test_handshake_payload_is_serialized_correctly() {
1010        let hs = ExtendedHandshake {
1011            extensions: HashMap::from([(Extension::Metadata, 1), (Extension::PeerExchange, 2)]),
1012            listen_port: Some(6881),
1013            client_type: Some("µTorrent 1.2".to_owned()),
1014            yourip: None,
1015            ipv4: None,
1016            ipv6: None,
1017            request_limit: None,
1018            metadata_size: None,
1019        };
1020        assert_eq!(
1021            Vec::from(b"d1:md11:ut_metadatai1e6:ut_pexi2ee1:pi6881e1:v13:\xc2\xb5Torrent 1.2e"),
1022            hs.encode()
1023        );
1024    }
1025
1026    #[test]
1027    fn test_metadata_block_payload_is_parsed_correctly() {
1028        let payload = Vec::from("d8:msg_typei1e5:piecei0e10:total_sizei34256eexxxxxxxx");
1029        let msg = PeerMessage::Extended {
1030            id: Extension::ID_METADATA,
1031            data: payload,
1032        };
1033        let parsed = ExtendedMessage::try_from(msg).unwrap();
1034        assert!(
1035            parsed
1036                == ExtendedMessage::MetadataBlock {
1037                    piece: 0,
1038                    total_size: 34256,
1039                    data: Vec::from(b"xxxxxxxx"),
1040                }
1041        );
1042    }
1043
1044    #[test]
1045    fn test_metadata_block_payload_is_serialized_correctly() {
1046        let msg = MetadataMsg::Block {
1047            piece: 0,
1048            total_size: 34256,
1049            data: Vec::from(b"xxxxxxxx"),
1050        };
1051        assert_eq!(
1052            Vec::from("d8:msg_typei1e5:piecei0e10:total_sizei34256eexxxxxxxx"),
1053            msg.encode()
1054        );
1055    }
1056
1057    #[test]
1058    fn test_pex_is_parsed_correctly() {
1059        let payload = vec![
1060            100, 53, 58, 97, 100, 100, 101, 100, 49, 50, 58, 1, 2, 3, 4, 48, 57, 5, 6, 7, 8, 212,
1061            49, 54, 58, 97, 100, 100, 101, 100, 54, 51, 54, 58, 0, 8, 0, 7, 0, 6, 0, 5, 0, 4, 0, 3,
1062            0, 2, 0, 1, 168, 202, 0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0, 7, 0, 8, 168, 202, 55, 58,
1063            100, 114, 111, 112, 112, 101, 100, 54, 58, 1, 2, 3, 4, 26, 225, 56, 58, 100, 114, 111,
1064            112, 112, 101, 100, 54, 49, 56, 58, 0, 1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0, 7, 0, 8,
1065            168, 202, 101,
1066        ];
1067        let pex = PeerExchangeData::decode(&payload).unwrap();
1068        assert_eq!(
1069            pex.added,
1070            HashSet::from([
1071                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)), 12345),
1072                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(5, 6, 7, 8)), 54321),
1073                SocketAddr::new(IpAddr::V6(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8)), 43210),
1074                SocketAddr::new(IpAddr::V6(Ipv6Addr::new(8, 7, 6, 5, 4, 3, 2, 1)), 43210)
1075            ])
1076        );
1077        assert_eq!(
1078            pex.dropped,
1079            HashSet::from([
1080                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)), 6881),
1081                SocketAddr::new(IpAddr::V6(Ipv6Addr::new(1, 2, 3, 4, 5, 6, 7, 8)), 43210),
1082            ])
1083        );
1084    }
1085}
mtorrent_core/pwp/message.rs

mtorrent_core/pwp/
message.rs
