irontide_wire/
message.rs

1#![allow(
2    clippy::cast_possible_truncation,
3    clippy::cast_possible_wrap,
4    clippy::cast_sign_loss,
5    reason = "M175: BitTorrent wire format — message field widths fixed by BEP 3/6/10 spec (u32 piece indices, u32 chunk lengths)"
6)]
7
8use bytes::{BufMut, Bytes, BytesMut};
9
10use crate::error::{Error, Result};
11
12/// Standard `BitTorrent` peer wire messages (BEP 3).
13///
14/// Generic over buffer type `B` to support both owned (`Bytes`) and borrowed
15/// (`&[u8]`) payloads.  Data-carrying variants (`Piece`, `Bitfield`,
16/// `Extended`) use `B`; fixed-field variants are buffer-agnostic.
17///
18/// The `Piece` variant carries two buffer fields (`data_0`, `data_1`) to
19/// support ring-buffer wrap-around: when a block spans the ring boundary the
20/// two non-contiguous slices are stored separately. In the common (no-wrap)
21/// case `data_1` is empty.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub enum Message<B = Bytes> {
24    /// Keep connection alive (no payload, length=0).
25    KeepAlive,
26    /// Peer is choking us.
27    Choke,
28    /// Peer is unchoking us.
29    Unchoke,
30    /// We're interested in the peer's data.
31    Interested,
32    /// We're not interested.
33    NotInterested,
34    /// Peer has piece `index`.
35    Have {
36        /// Piece index.
37        index: u32,
38    },
39    /// Peer's complete bitfield.
40    Bitfield(B),
41    /// Request a block: piece index, byte offset within piece, length.
42    Request {
43        /// Piece index.
44        index: u32,
45        /// Byte offset within the piece.
46        begin: u32,
47        /// Requested block length in bytes.
48        length: u32,
49    },
50    /// A data block: piece index, byte offset, data.
51    ///
52    /// Two buffer fields support ring-buffer wrap-around.  When the block does
53    /// not straddle the ring boundary, `data_1` is empty.
54    Piece {
55        /// Piece index.
56        index: u32,
57        /// Byte offset within the piece.
58        begin: u32,
59        /// First (or only) contiguous slice of block payload.
60        data_0: B,
61        /// Second contiguous slice when the ring buffer wraps; empty otherwise.
62        data_1: B,
63    },
64    /// Cancel a previously sent request.
65    Cancel {
66        /// Piece index.
67        index: u32,
68        /// Byte offset within the piece.
69        begin: u32,
70        /// Block length in bytes.
71        length: u32,
72    },
73    /// DHT port (BEP 5).
74    Port(u16),
75    /// Extension message (BEP 10). `ext_id=0` is handshake.
76    Extended {
77        /// Extension message ID (0 = handshake).
78        ext_id: u8,
79        /// Bencoded extension payload.
80        payload: B,
81    },
82    /// BEP 6: Suggest a piece for the peer to download.
83    SuggestPiece(u32),
84    /// BEP 6: We have all pieces.
85    HaveAll,
86    /// BEP 6: We have no pieces.
87    HaveNone,
88    /// BEP 6: Reject a request from the peer.
89    RejectRequest {
90        /// Piece index.
91        index: u32,
92        /// Byte offset within the piece.
93        begin: u32,
94        /// Block length in bytes.
95        length: u32,
96    },
97    /// BEP 6: Piece index the peer is allowed to request while choked.
98    AllowedFast(u32),
99    /// BEP 52: Request hashes from a file's Merkle tree.
100    HashRequest {
101        /// File root hash identifying the Merkle tree.
102        pieces_root: irontide_core::Id32,
103        /// Tree layer (0 = leaf/block layer).
104        base: u32,
105        /// Starting node index within the layer.
106        index: u32,
107        /// Number of consecutive hashes requested.
108        count: u32,
109        /// Number of uncle proof layers to include.
110        proof_layers: u32,
111    },
112    /// BEP 52: Response with hashes and uncle proof.
113    Hashes {
114        /// File root hash identifying the Merkle tree.
115        pieces_root: irontide_core::Id32,
116        /// Tree layer (0 = leaf/block layer).
117        base: u32,
118        /// Starting node index within the layer.
119        index: u32,
120        /// Number of consecutive hashes in the response.
121        count: u32,
122        /// Number of uncle proof layers included.
123        proof_layers: u32,
124        /// Hash values followed by uncle proof hashes.
125        hashes: Vec<irontide_core::Id32>,
126    },
127    /// BEP 52: Reject a hash request.
128    HashReject {
129        /// File root hash identifying the Merkle tree.
130        pieces_root: irontide_core::Id32,
131        /// Tree layer that was requested.
132        base: u32,
133        /// Starting node index that was requested.
134        index: u32,
135        /// Number of hashes that was requested.
136        count: u32,
137        /// Number of proof layers that was requested.
138        proof_layers: u32,
139    },
140}
141
142// Message IDs per BEP 3
143const ID_CHOKE: u8 = 0;
144const ID_UNCHOKE: u8 = 1;
145const ID_INTERESTED: u8 = 2;
146const ID_NOT_INTERESTED: u8 = 3;
147const ID_HAVE: u8 = 4;
148const ID_BITFIELD: u8 = 5;
149const ID_REQUEST: u8 = 6;
150const ID_PIECE: u8 = 7;
151const ID_CANCEL: u8 = 8;
152const ID_PORT: u8 = 9;
153const ID_EXTENDED: u8 = 20;
154
155// BEP 6 Fast Extension
156const ID_SUGGEST_PIECE: u8 = 0x0D;
157const ID_HAVE_ALL: u8 = 0x0E;
158const ID_HAVE_NONE: u8 = 0x0F;
159const ID_REJECT_REQUEST: u8 = 0x10;
160const ID_ALLOWED_FAST: u8 = 0x11;
161
162// BEP 52 Hash Messages
163const ID_HASH_REQUEST: u8 = 21;
164const ID_HASHES: u8 = 22;
165const ID_HASH_REJECT: u8 = 23;
166
167impl<B: AsRef<[u8]>> Message<B> {
168    /// Serialize a message to bytes (length-prefix + id + payload).
169    ///
170    /// The returned bytes include the 4-byte length prefix.
171    pub fn to_bytes(&self) -> Bytes {
172        match self {
173            Self::KeepAlive => {
174                let mut buf = BytesMut::with_capacity(4);
175                buf.put_u32(0);
176                buf.freeze()
177            }
178            Self::Choke => fixed_msg(ID_CHOKE),
179            Self::Unchoke => fixed_msg(ID_UNCHOKE),
180            Self::Interested => fixed_msg(ID_INTERESTED),
181            Self::NotInterested => fixed_msg(ID_NOT_INTERESTED),
182            Self::Have { index } => {
183                let mut buf = BytesMut::with_capacity(9);
184                buf.put_u32(5);
185                buf.put_u8(ID_HAVE);
186                buf.put_u32(*index);
187                buf.freeze()
188            }
189            Self::Bitfield(bits) => {
190                let bits = bits.as_ref();
191                let mut buf = BytesMut::with_capacity(5 + bits.len());
192                buf.put_u32(1 + bits.len() as u32);
193                buf.put_u8(ID_BITFIELD);
194                buf.put_slice(bits);
195                buf.freeze()
196            }
197            Self::Request {
198                index,
199                begin,
200                length,
201            } => triple_msg(ID_REQUEST, *index, *begin, *length),
202            Self::Piece {
203                index,
204                begin,
205                data_0,
206                data_1,
207            } => {
208                let d0 = data_0.as_ref();
209                let d1 = data_1.as_ref();
210                let data_len = d0.len() + d1.len();
211                let mut buf = BytesMut::with_capacity(13 + data_len);
212                buf.put_u32(9 + data_len as u32);
213                buf.put_u8(ID_PIECE);
214                buf.put_u32(*index);
215                buf.put_u32(*begin);
216                buf.put_slice(d0);
217                buf.put_slice(d1);
218                buf.freeze()
219            }
220            Self::Cancel {
221                index,
222                begin,
223                length,
224            } => triple_msg(ID_CANCEL, *index, *begin, *length),
225            Self::Port(port) => {
226                let mut buf = BytesMut::with_capacity(7);
227                buf.put_u32(3);
228                buf.put_u8(ID_PORT);
229                buf.put_u16(*port);
230                buf.freeze()
231            }
232            Self::Extended { ext_id, payload } => {
233                let payload = payload.as_ref();
234                let mut buf = BytesMut::with_capacity(6 + payload.len());
235                buf.put_u32(2 + payload.len() as u32);
236                buf.put_u8(ID_EXTENDED);
237                buf.put_u8(*ext_id);
238                buf.put_slice(payload);
239                buf.freeze()
240            }
241            Self::SuggestPiece(index) => {
242                let mut buf = BytesMut::with_capacity(9);
243                buf.put_u32(5);
244                buf.put_u8(ID_SUGGEST_PIECE);
245                buf.put_u32(*index);
246                buf.freeze()
247            }
248            Self::HaveAll => fixed_msg(ID_HAVE_ALL),
249            Self::HaveNone => fixed_msg(ID_HAVE_NONE),
250            Self::RejectRequest {
251                index,
252                begin,
253                length,
254            } => triple_msg(ID_REJECT_REQUEST, *index, *begin, *length),
255            Self::AllowedFast(index) => {
256                let mut buf = BytesMut::with_capacity(9);
257                buf.put_u32(5);
258                buf.put_u8(ID_ALLOWED_FAST);
259                buf.put_u32(*index);
260                buf.freeze()
261            }
262            Self::HashRequest {
263                pieces_root,
264                base,
265                index,
266                count,
267                proof_layers,
268            }
269            | Self::HashReject {
270                pieces_root,
271                base,
272                index,
273                count,
274                proof_layers,
275            } => {
276                let id = match self {
277                    Self::HashRequest { .. } => ID_HASH_REQUEST,
278                    _ => ID_HASH_REJECT,
279                };
280                let mut buf = BytesMut::with_capacity(53);
281                buf.put_u32(49); // 1 + 32 + 4*4
282                buf.put_u8(id);
283                buf.put_slice(&pieces_root.0);
284                buf.put_u32(*base);
285                buf.put_u32(*index);
286                buf.put_u32(*count);
287                buf.put_u32(*proof_layers);
288                buf.freeze()
289            }
290            Self::Hashes {
291                pieces_root,
292                base,
293                index,
294                count,
295                proof_layers,
296                hashes,
297            } => {
298                let hash_bytes = hashes.len() * 32;
299                let payload_len = 1 + 32 + 16 + hash_bytes;
300                let mut buf = BytesMut::with_capacity(4 + payload_len);
301                buf.put_u32(payload_len as u32);
302                buf.put_u8(ID_HASHES);
303                buf.put_slice(&pieces_root.0);
304                buf.put_u32(*base);
305                buf.put_u32(*index);
306                buf.put_u32(*count);
307                buf.put_u32(*proof_layers);
308                for h in hashes {
309                    buf.put_slice(&h.0);
310                }
311                buf.freeze()
312            }
313        }
314    }
315
316    /// Encode this message (with length prefix) directly into a buffer.
317    ///
318    /// Unlike [`to_bytes`](Self::to_bytes), this writes directly into `dst`
319    /// without allocating an intermediate `Bytes`, avoiding a double-copy
320    /// when used with `tokio_util::codec::Encoder`.
321    pub fn encode_into(&self, dst: &mut BytesMut) {
322        match self {
323            Self::KeepAlive => {
324                dst.put_u32(0);
325            }
326            Self::Choke => encode_fixed_into(dst, ID_CHOKE),
327            Self::Unchoke => encode_fixed_into(dst, ID_UNCHOKE),
328            Self::Interested => encode_fixed_into(dst, ID_INTERESTED),
329            Self::NotInterested => encode_fixed_into(dst, ID_NOT_INTERESTED),
330            Self::Have { index } => {
331                dst.put_u32(5);
332                dst.put_u8(ID_HAVE);
333                dst.put_u32(*index);
334            }
335            Self::Bitfield(bits) => {
336                let bits = bits.as_ref();
337                dst.reserve(5 + bits.len());
338                dst.put_u32(1 + bits.len() as u32);
339                dst.put_u8(ID_BITFIELD);
340                dst.put_slice(bits);
341            }
342            Self::Request {
343                index,
344                begin,
345                length,
346            } => encode_triple_into(dst, ID_REQUEST, *index, *begin, *length),
347            Self::Piece {
348                index,
349                begin,
350                data_0,
351                data_1,
352            } => {
353                let d0 = data_0.as_ref();
354                let d1 = data_1.as_ref();
355                let data_len = d0.len() + d1.len();
356                dst.reserve(13 + data_len);
357                dst.put_u32(9 + data_len as u32);
358                dst.put_u8(ID_PIECE);
359                dst.put_u32(*index);
360                dst.put_u32(*begin);
361                dst.put_slice(d0);
362                dst.put_slice(d1);
363            }
364            Self::Cancel {
365                index,
366                begin,
367                length,
368            } => encode_triple_into(dst, ID_CANCEL, *index, *begin, *length),
369            Self::Port(port) => {
370                dst.put_u32(3);
371                dst.put_u8(ID_PORT);
372                dst.put_u16(*port);
373            }
374            Self::Extended { ext_id, payload } => {
375                let payload = payload.as_ref();
376                dst.reserve(6 + payload.len());
377                dst.put_u32(2 + payload.len() as u32);
378                dst.put_u8(ID_EXTENDED);
379                dst.put_u8(*ext_id);
380                dst.put_slice(payload);
381            }
382            Self::SuggestPiece(index) => {
383                dst.put_u32(5);
384                dst.put_u8(ID_SUGGEST_PIECE);
385                dst.put_u32(*index);
386            }
387            Self::HaveAll => encode_fixed_into(dst, ID_HAVE_ALL),
388            Self::HaveNone => encode_fixed_into(dst, ID_HAVE_NONE),
389            Self::RejectRequest {
390                index,
391                begin,
392                length,
393            } => encode_triple_into(dst, ID_REJECT_REQUEST, *index, *begin, *length),
394            Self::AllowedFast(index) => {
395                dst.put_u32(5);
396                dst.put_u8(ID_ALLOWED_FAST);
397                dst.put_u32(*index);
398            }
399            Self::HashRequest {
400                pieces_root,
401                base,
402                index,
403                count,
404                proof_layers,
405            }
406            | Self::HashReject {
407                pieces_root,
408                base,
409                index,
410                count,
411                proof_layers,
412            } => {
413                let id = match self {
414                    Self::HashRequest { .. } => ID_HASH_REQUEST,
415                    _ => ID_HASH_REJECT,
416                };
417                dst.put_u32(49); // 1 + 32 + 4*4
418                dst.put_u8(id);
419                dst.put_slice(&pieces_root.0);
420                dst.put_u32(*base);
421                dst.put_u32(*index);
422                dst.put_u32(*count);
423                dst.put_u32(*proof_layers);
424            }
425            Self::Hashes {
426                pieces_root,
427                base,
428                index,
429                count,
430                proof_layers,
431                hashes,
432            } => {
433                let hash_bytes = hashes.len() * 32;
434                let payload_len = 1 + 32 + 16 + hash_bytes;
435                dst.reserve(4 + payload_len);
436                dst.put_u32(payload_len as u32);
437                dst.put_u8(ID_HASHES);
438                dst.put_slice(&pieces_root.0);
439                dst.put_u32(*base);
440                dst.put_u32(*index);
441                dst.put_u32(*count);
442                dst.put_u32(*proof_layers);
443                for h in hashes {
444                    dst.put_slice(&h.0);
445                }
446            }
447        }
448    }
449
450    /// Return the exact encoded wire length in bytes, including the 4-byte
451    /// length prefix.
452    ///
453    /// This is a pure computation — no allocation, no encoding.  Use it to
454    /// check whether a message fits in a fixed-size buffer before calling
455    /// [`encode_to_slice`](Self::encode_to_slice).
456    #[must_use]
457    pub fn wire_len(&self) -> usize {
458        match self {
459            Self::KeepAlive => 4,
460            Self::Choke
461            | Self::Unchoke
462            | Self::Interested
463            | Self::NotInterested
464            | Self::HaveAll
465            | Self::HaveNone => 5,
466            Self::Have { .. } | Self::SuggestPiece(_) | Self::AllowedFast(_) => 9,
467            Self::Port(_) => 7,
468            Self::Request { .. } | Self::Cancel { .. } | Self::RejectRequest { .. } => 17,
469            Self::Bitfield(bits) => 5 + bits.as_ref().len(),
470            Self::Piece { data_0, data_1, .. } => {
471                13 + data_0.as_ref().len() + data_1.as_ref().len()
472            }
473            Self::Extended { payload, .. } => 6 + payload.as_ref().len(),
474            Self::HashRequest { .. } | Self::HashReject { .. } => 53,
475            Self::Hashes { hashes, .. } => 53 + hashes.len() * 32,
476        }
477    }
478
479    /// Encode this message (with length prefix) into a raw byte slice.
480    ///
481    /// Returns the number of bytes written. The caller must ensure `dst` is
482    /// large enough to hold the encoded message. For peer wire messages,
483    /// `MAX_MSG_LEN` (16397) is always sufficient.
484    ///
485    /// Uses `std::io::Cursor<&mut [u8]>` + `std::io::Write` instead of
486    /// `BufMut`, producing identical bytes to [`encode_into`](Self::encode_into).
487    #[must_use]
488    pub fn encode_to_slice(&self, dst: &mut [u8]) -> usize {
489        use std::io::{Cursor, Write};
490
491        let mut cursor = Cursor::new(dst);
492
493        match self {
494            Self::KeepAlive => {
495                cursor.write_all(&0u32.to_be_bytes()).unwrap();
496            }
497            Self::Choke => {
498                cursor.write_all(&1u32.to_be_bytes()).unwrap();
499                cursor.write_all(&[ID_CHOKE]).unwrap();
500            }
501            Self::Unchoke => {
502                cursor.write_all(&1u32.to_be_bytes()).unwrap();
503                cursor.write_all(&[ID_UNCHOKE]).unwrap();
504            }
505            Self::Interested => {
506                cursor.write_all(&1u32.to_be_bytes()).unwrap();
507                cursor.write_all(&[ID_INTERESTED]).unwrap();
508            }
509            Self::NotInterested => {
510                cursor.write_all(&1u32.to_be_bytes()).unwrap();
511                cursor.write_all(&[ID_NOT_INTERESTED]).unwrap();
512            }
513            Self::Have { index } => {
514                cursor.write_all(&5u32.to_be_bytes()).unwrap();
515                cursor.write_all(&[ID_HAVE]).unwrap();
516                cursor.write_all(&index.to_be_bytes()).unwrap();
517            }
518            Self::Bitfield(bits) => {
519                let bits = bits.as_ref();
520                cursor
521                    .write_all(&(1 + bits.len() as u32).to_be_bytes())
522                    .unwrap();
523                cursor.write_all(&[ID_BITFIELD]).unwrap();
524                cursor.write_all(bits).unwrap();
525            }
526            Self::Request {
527                index,
528                begin,
529                length,
530            } => {
531                cursor.write_all(&13u32.to_be_bytes()).unwrap();
532                cursor.write_all(&[ID_REQUEST]).unwrap();
533                cursor.write_all(&index.to_be_bytes()).unwrap();
534                cursor.write_all(&begin.to_be_bytes()).unwrap();
535                cursor.write_all(&length.to_be_bytes()).unwrap();
536            }
537            Self::Piece {
538                index,
539                begin,
540                data_0,
541                data_1,
542            } => {
543                let d0 = data_0.as_ref();
544                let d1 = data_1.as_ref();
545                let data_len = d0.len() + d1.len();
546                cursor
547                    .write_all(&(9 + data_len as u32).to_be_bytes())
548                    .unwrap();
549                cursor.write_all(&[ID_PIECE]).unwrap();
550                cursor.write_all(&index.to_be_bytes()).unwrap();
551                cursor.write_all(&begin.to_be_bytes()).unwrap();
552                cursor.write_all(d0).unwrap();
553                cursor.write_all(d1).unwrap();
554            }
555            Self::Cancel {
556                index,
557                begin,
558                length,
559            } => {
560                cursor.write_all(&13u32.to_be_bytes()).unwrap();
561                cursor.write_all(&[ID_CANCEL]).unwrap();
562                cursor.write_all(&index.to_be_bytes()).unwrap();
563                cursor.write_all(&begin.to_be_bytes()).unwrap();
564                cursor.write_all(&length.to_be_bytes()).unwrap();
565            }
566            Self::Port(port) => {
567                cursor.write_all(&3u32.to_be_bytes()).unwrap();
568                cursor.write_all(&[ID_PORT]).unwrap();
569                cursor.write_all(&port.to_be_bytes()).unwrap();
570            }
571            Self::Extended { ext_id, payload } => {
572                let payload = payload.as_ref();
573                cursor
574                    .write_all(&(2 + payload.len() as u32).to_be_bytes())
575                    .unwrap();
576                cursor.write_all(&[ID_EXTENDED]).unwrap();
577                cursor.write_all(&[*ext_id]).unwrap();
578                cursor.write_all(payload).unwrap();
579            }
580            Self::SuggestPiece(index) => {
581                cursor.write_all(&5u32.to_be_bytes()).unwrap();
582                cursor.write_all(&[ID_SUGGEST_PIECE]).unwrap();
583                cursor.write_all(&index.to_be_bytes()).unwrap();
584            }
585            Self::HaveAll => {
586                cursor.write_all(&1u32.to_be_bytes()).unwrap();
587                cursor.write_all(&[ID_HAVE_ALL]).unwrap();
588            }
589            Self::HaveNone => {
590                cursor.write_all(&1u32.to_be_bytes()).unwrap();
591                cursor.write_all(&[ID_HAVE_NONE]).unwrap();
592            }
593            Self::RejectRequest {
594                index,
595                begin,
596                length,
597            } => {
598                cursor.write_all(&13u32.to_be_bytes()).unwrap();
599                cursor.write_all(&[ID_REJECT_REQUEST]).unwrap();
600                cursor.write_all(&index.to_be_bytes()).unwrap();
601                cursor.write_all(&begin.to_be_bytes()).unwrap();
602                cursor.write_all(&length.to_be_bytes()).unwrap();
603            }
604            Self::AllowedFast(index) => {
605                cursor.write_all(&5u32.to_be_bytes()).unwrap();
606                cursor.write_all(&[ID_ALLOWED_FAST]).unwrap();
607                cursor.write_all(&index.to_be_bytes()).unwrap();
608            }
609            Self::HashRequest {
610                pieces_root,
611                base,
612                index,
613                count,
614                proof_layers,
615            }
616            | Self::HashReject {
617                pieces_root,
618                base,
619                index,
620                count,
621                proof_layers,
622            } => {
623                let id = match self {
624                    Self::HashRequest { .. } => ID_HASH_REQUEST,
625                    _ => ID_HASH_REJECT,
626                };
627                cursor.write_all(&49u32.to_be_bytes()).unwrap();
628                cursor.write_all(&[id]).unwrap();
629                cursor.write_all(&pieces_root.0).unwrap();
630                cursor.write_all(&base.to_be_bytes()).unwrap();
631                cursor.write_all(&index.to_be_bytes()).unwrap();
632                cursor.write_all(&count.to_be_bytes()).unwrap();
633                cursor.write_all(&proof_layers.to_be_bytes()).unwrap();
634            }
635            Self::Hashes {
636                pieces_root,
637                base,
638                index,
639                count,
640                proof_layers,
641                hashes,
642            } => {
643                let hash_bytes = hashes.len() * 32;
644                let payload_len = 1 + 32 + 16 + hash_bytes;
645                cursor
646                    .write_all(&(payload_len as u32).to_be_bytes())
647                    .unwrap();
648                cursor.write_all(&[ID_HASHES]).unwrap();
649                cursor.write_all(&pieces_root.0).unwrap();
650                cursor.write_all(&base.to_be_bytes()).unwrap();
651                cursor.write_all(&index.to_be_bytes()).unwrap();
652                cursor.write_all(&count.to_be_bytes()).unwrap();
653                cursor.write_all(&proof_layers.to_be_bytes()).unwrap();
654                for h in hashes {
655                    cursor.write_all(&h.0).unwrap();
656                }
657            }
658        }
659
660        cursor.position() as usize
661    }
662}
663
664impl Message<&[u8]> {
665    /// Convert a borrowed message to an owned `Message<Bytes>`.
666    ///
667    /// Fixed-field variants are zero-cost (no data to copy).
668    /// Data-carrying variants (`Piece`, `Bitfield`, `Extended`) copy their
669    /// slices into fresh `Bytes` allocations.
670    #[must_use]
671    pub fn to_owned_bytes(&self) -> Message<Bytes> {
672        match *self {
673            Message::KeepAlive => Message::KeepAlive,
674            Message::Choke => Message::Choke,
675            Message::Unchoke => Message::Unchoke,
676            Message::Interested => Message::Interested,
677            Message::NotInterested => Message::NotInterested,
678            Message::Have { index } => Message::Have { index },
679            Message::Bitfield(data) => Message::Bitfield(Bytes::copy_from_slice(data)),
680            Message::Request {
681                index,
682                begin,
683                length,
684            } => Message::Request {
685                index,
686                begin,
687                length,
688            },
689            Message::Piece {
690                index,
691                begin,
692                data_0,
693                data_1,
694            } => Message::Piece {
695                index,
696                begin,
697                data_0: Bytes::copy_from_slice(data_0),
698                data_1: Bytes::copy_from_slice(data_1),
699            },
700            Message::Cancel {
701                index,
702                begin,
703                length,
704            } => Message::Cancel {
705                index,
706                begin,
707                length,
708            },
709            Message::Port(port) => Message::Port(port),
710            Message::Extended { ext_id, payload } => Message::Extended {
711                ext_id,
712                payload: Bytes::copy_from_slice(payload),
713            },
714            Message::SuggestPiece(index) => Message::SuggestPiece(index),
715            Message::HaveAll => Message::HaveAll,
716            Message::HaveNone => Message::HaveNone,
717            Message::RejectRequest {
718                index,
719                begin,
720                length,
721            } => Message::RejectRequest {
722                index,
723                begin,
724                length,
725            },
726            Message::AllowedFast(index) => Message::AllowedFast(index),
727            Message::HashRequest {
728                pieces_root,
729                base,
730                index,
731                count,
732                proof_layers,
733            } => Message::HashRequest {
734                pieces_root,
735                base,
736                index,
737                count,
738                proof_layers,
739            },
740            Message::Hashes {
741                ref pieces_root,
742                base,
743                index,
744                count,
745                proof_layers,
746                ref hashes,
747            } => Message::Hashes {
748                pieces_root: *pieces_root,
749                base,
750                index,
751                count,
752                proof_layers,
753                hashes: hashes.clone(),
754            },
755            Message::HashReject {
756                pieces_root,
757                base,
758                index,
759                count,
760                proof_layers,
761            } => Message::HashReject {
762                pieces_root,
763                base,
764                index,
765                count,
766                proof_layers,
767            },
768        }
769    }
770}
771
772impl Message<Bytes> {
773    /// Parse a message from its payload (after the 4-byte length prefix has
774    /// been consumed). `payload` is everything after the length prefix.
775    ///
776    /// # Errors
777    ///
778    /// Returns an error if the message ID is unknown or the payload is malformed.
779    #[allow(clippy::needless_pass_by_value, reason = "pub API stability")]
780    pub fn from_payload(payload: Bytes) -> Result<Self> {
781        if payload.is_empty() {
782            return Ok(Self::KeepAlive);
783        }
784
785        let id = payload[0];
786        let body = &payload[1..];
787
788        match id {
789            ID_CHOKE => Ok(Self::Choke),
790            ID_UNCHOKE => Ok(Self::Unchoke),
791            ID_INTERESTED => Ok(Self::Interested),
792            ID_NOT_INTERESTED => Ok(Self::NotInterested),
793            ID_HAVE => {
794                ensure_len(body, 4, "Have")?;
795                Ok(Self::Have {
796                    index: read_u32(body),
797                })
798            }
799            ID_BITFIELD => Ok(Self::Bitfield(payload.slice(1..))),
800            ID_REQUEST => {
801                ensure_len(body, 12, "Request")?;
802                Ok(Self::Request {
803                    index: read_u32(body),
804                    begin: read_u32(&body[4..]),
805                    length: read_u32(&body[8..]),
806                })
807            }
808            ID_PIECE => {
809                ensure_len(body, 8, "Piece")?;
810                let index = read_u32(body);
811                let begin = read_u32(&body[4..]);
812                Ok(Self::Piece {
813                    index,
814                    begin,
815                    data_0: payload.slice(9..),
816                    data_1: Bytes::new(),
817                })
818            }
819            ID_CANCEL => {
820                ensure_len(body, 12, "Cancel")?;
821                Ok(Self::Cancel {
822                    index: read_u32(body),
823                    begin: read_u32(&body[4..]),
824                    length: read_u32(&body[8..]),
825                })
826            }
827            ID_PORT => {
828                ensure_len(body, 2, "Port")?;
829                Ok(Self::Port(u16::from_be_bytes([body[0], body[1]])))
830            }
831            ID_EXTENDED => {
832                ensure_len(body, 1, "Extended")?;
833                let ext_id = body[0];
834                Ok(Self::Extended {
835                    ext_id,
836                    payload: payload.slice(2..),
837                })
838            }
839            ID_SUGGEST_PIECE => {
840                ensure_len(body, 4, "SuggestPiece")?;
841                Ok(Self::SuggestPiece(read_u32(body)))
842            }
843            ID_HAVE_ALL => Ok(Self::HaveAll),
844            ID_HAVE_NONE => Ok(Self::HaveNone),
845            ID_REJECT_REQUEST => {
846                ensure_len(body, 12, "RejectRequest")?;
847                Ok(Self::RejectRequest {
848                    index: read_u32(body),
849                    begin: read_u32(&body[4..]),
850                    length: read_u32(&body[8..]),
851                })
852            }
853            ID_ALLOWED_FAST => {
854                ensure_len(body, 4, "AllowedFast")?;
855                Ok(Self::AllowedFast(read_u32(body)))
856            }
857            ID_HASH_REQUEST | ID_HASH_REJECT => {
858                ensure_len(body, 48, "HashRequest/Reject")?;
859                let mut root = [0u8; 32];
860                root.copy_from_slice(&body[..32]);
861                let pieces_root = irontide_core::Id32(root);
862                let base = read_u32(&body[32..]);
863                let index = read_u32(&body[36..]);
864                let count = read_u32(&body[40..]);
865                let proof_layers = read_u32(&body[44..]);
866                if id == ID_HASH_REQUEST {
867                    Ok(Self::HashRequest {
868                        pieces_root,
869                        base,
870                        index,
871                        count,
872                        proof_layers,
873                    })
874                } else {
875                    Ok(Self::HashReject {
876                        pieces_root,
877                        base,
878                        index,
879                        count,
880                        proof_layers,
881                    })
882                }
883            }
884            ID_HASHES => {
885                ensure_len(body, 48, "Hashes")?;
886                let mut root = [0u8; 32];
887                root.copy_from_slice(&body[..32]);
888                let pieces_root = irontide_core::Id32(root);
889                let base = read_u32(&body[32..]);
890                let index = read_u32(&body[36..]);
891                let count = read_u32(&body[40..]);
892                let proof_layers = read_u32(&body[44..]);
893                let hash_data = &body[48..];
894                if !hash_data.len().is_multiple_of(32) {
895                    return Err(Error::MessageTooShort {
896                        expected: 48 + 32,
897                        got: body.len(),
898                    });
899                }
900                let hashes = hash_data
901                    .chunks_exact(32)
902                    .map(|chunk| {
903                        let mut h = [0u8; 32];
904                        h.copy_from_slice(chunk);
905                        irontide_core::Id32(h)
906                    })
907                    .collect();
908                Ok(Self::Hashes {
909                    pieces_root,
910                    base,
911                    index,
912                    count,
913                    proof_layers,
914                    hashes,
915                })
916            }
917            _ => Err(Error::InvalidMessageId(id)),
918        }
919    }
920}
921
922fn encode_fixed_into(dst: &mut BytesMut, id: u8) {
923    dst.put_u32(1);
924    dst.put_u8(id);
925}
926
927fn encode_triple_into(dst: &mut BytesMut, id: u8, a: u32, b: u32, c: u32) {
928    dst.put_u32(13);
929    dst.put_u8(id);
930    dst.put_u32(a);
931    dst.put_u32(b);
932    dst.put_u32(c);
933}
934
935fn fixed_msg(id: u8) -> Bytes {
936    let mut buf = BytesMut::with_capacity(5);
937    buf.put_u32(1);
938    buf.put_u8(id);
939    buf.freeze()
940}
941
942fn triple_msg(id: u8, a: u32, b: u32, c: u32) -> Bytes {
943    let mut buf = BytesMut::with_capacity(17);
944    buf.put_u32(13);
945    buf.put_u8(id);
946    buf.put_u32(a);
947    buf.put_u32(b);
948    buf.put_u32(c);
949    buf.freeze()
950}
951
952fn read_u32(buf: &[u8]) -> u32 {
953    let mut b = [0u8; 4];
954    b.copy_from_slice(&buf[..4]);
955    u32::from_be_bytes(b)
956}
957
958fn ensure_len(body: &[u8], min: usize, _name: &str) -> Result<()> {
959    if body.len() < min {
960        Err(Error::MessageTooShort {
961            expected: min,
962            got: body.len(),
963        })
964    } else {
965        Ok(())
966    }
967}
968
969/// BEP 6 Allowed-Fast set generation.
970///
971/// Generates a deterministic set of piece indices that a peer is allowed
972/// to request even while choked. Uses IP masking + `info_hash` + SHA1.
973///
974/// For IPv4: masks to /24 (matching BEP 6 spec).
975/// For IPv6: masks to /48 (matching libtorrent convention).
976#[must_use]
977pub fn allowed_fast_set(
978    info_hash: &irontide_core::Id20,
979    peer_ip: std::net::Ipv4Addr,
980    num_pieces: u32,
981    count: usize,
982) -> Vec<u32> {
983    allowed_fast_set_for_ip(info_hash, std::net::IpAddr::V4(peer_ip), num_pieces, count)
984}
985
986/// BEP 6 Allowed-Fast set generation for any IP address family.
987///
988/// IPv4: /24 prefix mask. IPv6: /48 prefix mask (libtorrent convention).
989#[must_use]
990pub fn allowed_fast_set_for_ip(
991    info_hash: &irontide_core::Id20,
992    peer_ip: std::net::IpAddr,
993    num_pieces: u32,
994    count: usize,
995) -> Vec<u32> {
996    use irontide_core::sha1;
997
998    if num_pieces == 0 {
999        return Vec::new();
1000    }
1001
1002    let count = count.min(num_pieces as usize);
1003    let mut result = Vec::with_capacity(count);
1004
1005    // Build masked IP bytes based on address family
1006    let masked: Vec<u8> = match peer_ip {
1007        std::net::IpAddr::V4(ipv4) => {
1008            // Mask to /24
1009            let o = ipv4.octets();
1010            vec![o[0], o[1], o[2], 0]
1011        }
1012        std::net::IpAddr::V6(ipv6) => {
1013            // Mask to /48: keep first 6 bytes, zero the rest
1014            let o = ipv6.octets();
1015            let mut masked = [0u8; 16];
1016            masked[..6].copy_from_slice(&o[..6]);
1017            masked.to_vec()
1018        }
1019    };
1020
1021    // Initial hash: SHA1(masked_ip + info_hash)
1022    let mut input = Vec::with_capacity(masked.len() + 20);
1023    input.extend_from_slice(&masked);
1024    input.extend_from_slice(info_hash.as_bytes());
1025    let mut hash = sha1(&input);
1026
1027    while result.len() < count {
1028        let hash_bytes = hash.as_bytes();
1029        // Each 20-byte hash gives us 5 candidate indices (4 bytes each)
1030        for i in (0..20).step_by(4) {
1031            if result.len() >= count {
1032                break;
1033            }
1034            let index = u32::from_be_bytes([
1035                hash_bytes[i],
1036                hash_bytes[i + 1],
1037                hash_bytes[i + 2],
1038                hash_bytes[i + 3],
1039            ]) % num_pieces;
1040            if !result.contains(&index) {
1041                result.push(index);
1042            }
1043        }
1044        // Re-hash for more candidates
1045        hash = sha1(hash.as_bytes());
1046    }
1047
1048    result
1049}
1050
1051#[cfg(test)]
1052mod tests {
1053    use super::*;
1054
1055    #[allow(clippy::needless_pass_by_value, reason = "test helper convenience")]
1056    fn round_trip(msg: Message) {
1057        let bytes = msg.to_bytes();
1058        // Skip the 4-byte length prefix for parsing
1059        let parsed = Message::from_payload(Bytes::copy_from_slice(&bytes[4..])).unwrap();
1060        assert_eq!(msg, parsed);
1061    }
1062
1063    #[test]
1064    fn keepalive() {
1065        round_trip(Message::KeepAlive);
1066    }
1067
1068    #[test]
1069    fn choke_unchoke() {
1070        round_trip(Message::Choke);
1071        round_trip(Message::Unchoke);
1072    }
1073
1074    #[test]
1075    fn interested() {
1076        round_trip(Message::Interested);
1077        round_trip(Message::NotInterested);
1078    }
1079
1080    #[test]
1081    fn have() {
1082        round_trip(Message::Have { index: 42 });
1083    }
1084
1085    #[test]
1086    fn bitfield() {
1087        round_trip(Message::Bitfield(Bytes::from_static(&[0xFF, 0x80])));
1088    }
1089
1090    #[test]
1091    fn request() {
1092        round_trip(Message::Request {
1093            index: 1,
1094            begin: 0,
1095            length: 16384,
1096        });
1097    }
1098
1099    #[test]
1100    fn piece() {
1101        round_trip(Message::Piece {
1102            index: 1,
1103            begin: 0,
1104            data_0: Bytes::from_static(b"hello world"),
1105            data_1: Bytes::new(),
1106        });
1107    }
1108
1109    #[test]
1110    fn cancel() {
1111        round_trip(Message::Cancel {
1112            index: 1,
1113            begin: 0,
1114            length: 16384,
1115        });
1116    }
1117
1118    #[test]
1119    fn port() {
1120        round_trip(Message::Port(6881));
1121    }
1122
1123    #[test]
1124    fn extended() {
1125        round_trip(Message::Extended {
1126            ext_id: 1,
1127            payload: Bytes::from_static(b"test payload"),
1128        });
1129    }
1130
1131    #[test]
1132    fn invalid_message_id() {
1133        assert!(Message::from_payload(Bytes::from_static(&[99u8])).is_err());
1134    }
1135
1136    #[test]
1137    fn suggest_piece() {
1138        round_trip(Message::SuggestPiece(42));
1139    }
1140
1141    #[test]
1142    fn have_all() {
1143        round_trip(Message::HaveAll);
1144    }
1145
1146    #[test]
1147    fn have_none() {
1148        round_trip(Message::HaveNone);
1149    }
1150
1151    #[test]
1152    fn reject_request() {
1153        round_trip(Message::RejectRequest {
1154            index: 1,
1155            begin: 0,
1156            length: 16384,
1157        });
1158    }
1159
1160    #[test]
1161    fn allowed_fast() {
1162        round_trip(Message::AllowedFast(7));
1163    }
1164
1165    #[test]
1166    fn allowed_fast_set_deterministic() {
1167        use irontide_core::Id20;
1168        let ih = Id20::from_hex("aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d").unwrap();
1169        let ip: std::net::Ipv4Addr = "192.168.1.100".parse().unwrap();
1170        let set1 = allowed_fast_set(&ih, ip, 1000, 10);
1171        let set2 = allowed_fast_set(&ih, ip, 1000, 10);
1172        assert_eq!(set1, set2);
1173        assert_eq!(set1.len(), 10);
1174        // All indices in range
1175        assert!(set1.iter().all(|&i| i < 1000));
1176    }
1177
1178    #[test]
1179    fn allowed_fast_set_unique() {
1180        use irontide_core::Id20;
1181        let ih = Id20::from_hex("aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d").unwrap();
1182        let ip: std::net::Ipv4Addr = "10.0.0.1".parse().unwrap();
1183        let set = allowed_fast_set(&ih, ip, 50, 10);
1184        let unique: std::collections::HashSet<u32> = set.iter().copied().collect();
1185        assert_eq!(set.len(), unique.len(), "all indices should be unique");
1186    }
1187
1188    #[test]
1189    fn allowed_fast_set_empty_torrent() {
1190        use irontide_core::Id20;
1191        let ih = Id20::ZERO;
1192        let ip: std::net::Ipv4Addr = "127.0.0.1".parse().unwrap();
1193        assert!(allowed_fast_set(&ih, ip, 0, 10).is_empty());
1194    }
1195
1196    #[test]
1197    fn allowed_fast_set_ipv6() {
1198        use irontide_core::Id20;
1199        let ih = Id20::from_hex("aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d").unwrap();
1200        let ip: std::net::IpAddr = "2001:db8::1".parse().unwrap();
1201        let set = allowed_fast_set_for_ip(&ih, ip, 1000, 10);
1202        assert_eq!(set.len(), 10);
1203        assert!(set.iter().all(|&i| i < 1000));
1204
1205        // Same /48 prefix should produce same set
1206        let ip2: std::net::IpAddr = "2001:db8::ffff".parse().unwrap();
1207        let set2 = allowed_fast_set_for_ip(&ih, ip2, 1000, 10);
1208        assert_eq!(set, set2);
1209
1210        // Different /48 prefix should produce different set
1211        let ip3: std::net::IpAddr = "2001:db9::1".parse().unwrap();
1212        let set3 = allowed_fast_set_for_ip(&ih, ip3, 1000, 10);
1213        assert_ne!(set, set3);
1214    }
1215
1216    #[test]
1217    fn hash_request_round_trip() {
1218        let msg = Message::HashRequest {
1219            pieces_root: irontide_core::Id32::ZERO,
1220            base: 7,
1221            index: 0,
1222            count: 512,
1223            proof_layers: 3,
1224        };
1225        round_trip(msg);
1226    }
1227
1228    #[test]
1229    fn hash_reject_round_trip() {
1230        let msg = Message::HashReject {
1231            pieces_root: irontide_core::Id32::ZERO,
1232            base: 7,
1233            index: 0,
1234            count: 512,
1235            proof_layers: 3,
1236        };
1237        round_trip(msg);
1238    }
1239
1240    #[test]
1241    fn hashes_round_trip() {
1242        let h1 = irontide_core::sha256(b"block1");
1243        let h2 = irontide_core::sha256(b"block2");
1244        let uncle = irontide_core::sha256(b"uncle");
1245        let msg = Message::Hashes {
1246            pieces_root: irontide_core::Id32::ZERO,
1247            base: 0,
1248            index: 0,
1249            count: 2,
1250            proof_layers: 1,
1251            hashes: vec![h1, h2, uncle],
1252        };
1253        round_trip(msg);
1254    }
1255
1256    #[test]
1257    fn hash_request_exact_wire_size() {
1258        let msg: Message = Message::HashRequest {
1259            pieces_root: irontide_core::Id32::ZERO,
1260            base: 0,
1261            index: 0,
1262            count: 1,
1263            proof_layers: 0,
1264        };
1265        let bytes = msg.to_bytes();
1266        // 4 (length prefix) + 1 (msg id) + 32 (root) + 4*4 (fields) = 53
1267        assert_eq!(bytes.len(), 53);
1268    }
1269
1270    #[test]
1271    fn hashes_variable_length() {
1272        let h = irontide_core::sha256(b"test");
1273        let msg: Message = Message::Hashes {
1274            pieces_root: irontide_core::Id32::ZERO,
1275            base: 0,
1276            index: 0,
1277            count: 1,
1278            proof_layers: 0,
1279            hashes: vec![h],
1280        };
1281        let bytes = msg.to_bytes();
1282        // 4 + 1 + 32 + 4*4 + 1*32 = 85
1283        assert_eq!(bytes.len(), 85);
1284    }
1285
1286    #[test]
1287    fn hash_request_too_short() {
1288        // msg id 21, but only 10 bytes of body (need 48)
1289        let mut payload = vec![21u8];
1290        payload.extend_from_slice(&[0u8; 10]);
1291        assert!(Message::from_payload(Bytes::from(payload)).is_err());
1292    }
1293
1294    #[test]
1295    fn encode_into_matches_to_bytes() {
1296        let messages = vec![
1297            Message::KeepAlive,
1298            Message::Choke,
1299            Message::Unchoke,
1300            Message::Interested,
1301            Message::NotInterested,
1302            Message::Have { index: 42 },
1303            Message::Bitfield(Bytes::from_static(b"\xff\x00")),
1304            Message::Request {
1305                index: 1,
1306                begin: 0,
1307                length: 16384,
1308            },
1309            Message::Piece {
1310                index: 0,
1311                begin: 0,
1312                data_0: Bytes::from_static(b"block data here"),
1313                data_1: Bytes::new(),
1314            },
1315            Message::Cancel {
1316                index: 1,
1317                begin: 0,
1318                length: 16384,
1319            },
1320            Message::Port(6881),
1321            Message::Extended {
1322                ext_id: 0,
1323                payload: Bytes::from_static(b"ext payload"),
1324            },
1325            Message::SuggestPiece(7),
1326            Message::HaveAll,
1327            Message::HaveNone,
1328            Message::RejectRequest {
1329                index: 1,
1330                begin: 0,
1331                length: 16384,
1332            },
1333            Message::AllowedFast(5),
1334            Message::HashRequest {
1335                pieces_root: irontide_core::Id32::ZERO,
1336                base: 7,
1337                index: 0,
1338                count: 512,
1339                proof_layers: 3,
1340            },
1341            Message::HashReject {
1342                pieces_root: irontide_core::Id32::ZERO,
1343                base: 7,
1344                index: 0,
1345                count: 512,
1346                proof_layers: 3,
1347            },
1348            Message::Hashes {
1349                pieces_root: irontide_core::Id32::ZERO,
1350                base: 0,
1351                index: 0,
1352                count: 2,
1353                proof_layers: 1,
1354                hashes: vec![
1355                    irontide_core::sha256(b"block1"),
1356                    irontide_core::sha256(b"block2"),
1357                    irontide_core::sha256(b"uncle"),
1358                ],
1359            },
1360        ];
1361        for msg in messages {
1362            let expected = msg.to_bytes();
1363            let mut buf = BytesMut::new();
1364            msg.encode_into(&mut buf);
1365            assert_eq!(&expected[..], &buf[..], "mismatch for {msg:?}");
1366        }
1367    }
1368
1369    #[test]
1370    fn allowed_fast_set_ipv4_compat() {
1371        // allowed_fast_set (IPv4-only) and allowed_fast_set_for_ip with V4 should match
1372        use irontide_core::Id20;
1373        let ih = Id20::from_hex("aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d").unwrap();
1374        let ipv4: std::net::Ipv4Addr = "192.168.1.100".parse().unwrap();
1375        let set_v4 = allowed_fast_set(&ih, ipv4, 1000, 10);
1376        let set_ip = allowed_fast_set_for_ip(&ih, std::net::IpAddr::V4(ipv4), 1000, 10);
1377        assert_eq!(set_v4, set_ip);
1378    }
1379
1380    // ── M158: BEP 6 AllowedFast spec vector conformance ──
1381
1382    #[test]
1383    fn allowed_fast_bep6_spec_vector_k7() {
1384        // BEP 6 exact test vector: IP 80.4.4.200, InfoHash 0xaaaa..aa, 1313 pieces, k=7
1385        use irontide_core::Id20;
1386        let ih = Id20([0xaa; 20]);
1387        let ip: std::net::Ipv4Addr = "80.4.4.200".parse().unwrap();
1388        let set = allowed_fast_set(&ih, ip, 1313, 7);
1389        assert_eq!(set, vec![1059, 431, 808, 1217, 287, 376, 1188]);
1390    }
1391
1392    #[test]
1393    fn allowed_fast_bep6_spec_vector_k9() {
1394        // BEP 6 exact test vector: same inputs, k=9 — extends k=7 with [353, 508]
1395        use irontide_core::Id20;
1396        let ih = Id20([0xaa; 20]);
1397        let ip: std::net::Ipv4Addr = "80.4.4.200".parse().unwrap();
1398        let set = allowed_fast_set(&ih, ip, 1313, 9);
1399        assert_eq!(set, vec![1059, 431, 808, 1217, 287, 376, 1188, 353, 508]);
1400    }
1401
1402    #[test]
1403    fn allowed_fast_ip_masking() {
1404        // BEP 6 spec: peer IP is masked to /24 before hashing.
1405        // 80.4.4.200 should mask to 80.4.4.0 = [0x50, 0x04, 0x04, 0x00].
1406        // Verify that any IP in the same /24 produces the same set.
1407        use irontide_core::Id20;
1408        let ih = Id20([0xaa; 20]);
1409        let ip_a: std::net::Ipv4Addr = "80.4.4.200".parse().unwrap();
1410        let ip_b: std::net::Ipv4Addr = "80.4.4.0".parse().unwrap();
1411        let ip_c: std::net::Ipv4Addr = "80.4.4.255".parse().unwrap();
1412        let set_a = allowed_fast_set(&ih, ip_a, 1313, 7);
1413        let set_b = allowed_fast_set(&ih, ip_b, 1313, 7);
1414        let set_c = allowed_fast_set(&ih, ip_c, 1313, 7);
1415        assert_eq!(
1416            set_a, set_b,
1417            "80.4.4.200 and 80.4.4.0 must produce same set (same /24)"
1418        );
1419        assert_eq!(
1420            set_a, set_c,
1421            "80.4.4.200 and 80.4.4.255 must produce same set (same /24)"
1422        );
1423
1424        // Verify the masked bytes are correct: 80=0x50, 4=0x04, 4=0x04, 0=0x00
1425        let octets = ip_a.octets();
1426        let masked = [octets[0], octets[1], octets[2], 0u8];
1427        assert_eq!(masked, [0x50, 0x04, 0x04, 0x00]);
1428    }
1429
1430    // ── M110: Generic Message<B> tests ──
1431
1432    #[test]
1433    fn message_piece_two_fields_round_trip() {
1434        // Encode a Piece with data_0 only (data_1 empty) — common case.
1435        let msg = Message::Piece {
1436            index: 5,
1437            begin: 16384,
1438            data_0: Bytes::from_static(b"block payload here"),
1439            data_1: Bytes::new(),
1440        };
1441        let bytes = msg.to_bytes();
1442        let parsed = Message::from_payload(Bytes::copy_from_slice(&bytes[4..])).unwrap();
1443        assert_eq!(msg, parsed);
1444    }
1445
1446    #[test]
1447    fn message_piece_split_data_round_trip() {
1448        // Encode a Piece with both data_0 and data_1 populated (ring-wrap case).
1449        // The wire format concatenates them, so decoding puts everything into data_0.
1450        let msg = Message::Piece {
1451            index: 3,
1452            begin: 0,
1453            data_0: Bytes::from_static(b"first half"),
1454            data_1: Bytes::from_static(b" second half"),
1455        };
1456        let bytes = msg.to_bytes();
1457        let parsed = Message::from_payload(Bytes::copy_from_slice(&bytes[4..])).unwrap();
1458        // After round-trip through the wire, the data is concatenated into data_0
1459        // with data_1 empty (the receiver doesn't know about ring-wrap).
1460        assert_eq!(
1461            parsed,
1462            Message::Piece {
1463                index: 3,
1464                begin: 0,
1465                data_0: Bytes::from_static(b"first half second half"),
1466                data_1: Bytes::new(),
1467            }
1468        );
1469    }
1470
1471    #[test]
1472    fn message_generic_encode_borrowed() {
1473        // Verify that Message<&[u8]> can encode_into just like Message<Bytes>.
1474        let borrowed: Message<&[u8]> = Message::Piece {
1475            index: 1,
1476            begin: 0,
1477            data_0: b"borrowed data",
1478            data_1: b"",
1479        };
1480        let owned: Message<Bytes> = Message::Piece {
1481            index: 1,
1482            begin: 0,
1483            data_0: Bytes::from_static(b"borrowed data"),
1484            data_1: Bytes::new(),
1485        };
1486        let mut buf_borrowed = BytesMut::new();
1487        borrowed.encode_into(&mut buf_borrowed);
1488        let mut buf_owned = BytesMut::new();
1489        owned.encode_into(&mut buf_owned);
1490        assert_eq!(
1491            buf_borrowed, buf_owned,
1492            "borrowed and owned encode identically"
1493        );
1494
1495        // Also test to_bytes
1496        assert_eq!(borrowed.to_bytes(), owned.to_bytes());
1497
1498        // Test borrowed Bitfield
1499        let bf_borrowed: Message<&[u8]> = Message::Bitfield(b"\xff\x80");
1500        let bf_owned: Message<Bytes> = Message::Bitfield(Bytes::from_static(b"\xff\x80"));
1501        assert_eq!(bf_borrowed.to_bytes(), bf_owned.to_bytes());
1502
1503        // Test borrowed Extended
1504        let ext_borrowed: Message<&[u8]> = Message::Extended {
1505            ext_id: 1,
1506            payload: b"payload",
1507        };
1508        let ext_owned: Message<Bytes> = Message::Extended {
1509            ext_id: 1,
1510            payload: Bytes::from_static(b"payload"),
1511        };
1512        assert_eq!(ext_borrowed.to_bytes(), ext_owned.to_bytes());
1513    }
1514
1515    // ── M115: encode_to_slice tests ──
1516
1517    /// Build the complete set of message variants used for `encode_to_slice` tests.
1518    fn all_message_variants() -> Vec<Message> {
1519        vec![
1520            Message::KeepAlive,
1521            Message::Choke,
1522            Message::Unchoke,
1523            Message::Interested,
1524            Message::NotInterested,
1525            Message::Have { index: 42 },
1526            Message::Bitfield(Bytes::from_static(b"\xff\x00")),
1527            Message::Request {
1528                index: 1,
1529                begin: 0,
1530                length: 16384,
1531            },
1532            Message::Piece {
1533                index: 0,
1534                begin: 0,
1535                data_0: Bytes::from_static(b"block data here"),
1536                data_1: Bytes::new(),
1537            },
1538            Message::Piece {
1539                index: 3,
1540                begin: 0,
1541                data_0: Bytes::from_static(b"first half"),
1542                data_1: Bytes::from_static(b" second half"),
1543            },
1544            Message::Cancel {
1545                index: 1,
1546                begin: 0,
1547                length: 16384,
1548            },
1549            Message::Port(6881),
1550            Message::Extended {
1551                ext_id: 0,
1552                payload: Bytes::from_static(b"ext payload"),
1553            },
1554            Message::SuggestPiece(7),
1555            Message::HaveAll,
1556            Message::HaveNone,
1557            Message::RejectRequest {
1558                index: 1,
1559                begin: 0,
1560                length: 16384,
1561            },
1562            Message::AllowedFast(5),
1563            Message::HashRequest {
1564                pieces_root: irontide_core::Id32::ZERO,
1565                base: 7,
1566                index: 0,
1567                count: 512,
1568                proof_layers: 3,
1569            },
1570            Message::HashReject {
1571                pieces_root: irontide_core::Id32::ZERO,
1572                base: 7,
1573                index: 0,
1574                count: 512,
1575                proof_layers: 3,
1576            },
1577            Message::Hashes {
1578                pieces_root: irontide_core::Id32::ZERO,
1579                base: 0,
1580                index: 0,
1581                count: 2,
1582                proof_layers: 1,
1583                hashes: vec![
1584                    irontide_core::sha256(b"block1"),
1585                    irontide_core::sha256(b"block2"),
1586                    irontide_core::sha256(b"uncle"),
1587                ],
1588            },
1589        ]
1590    }
1591
1592    #[test]
1593    fn encode_to_slice_roundtrip() {
1594        for msg in all_message_variants() {
1595            let mut buf = [0u8; 4096];
1596            let n = msg.encode_to_slice(&mut buf);
1597            // Skip the 4-byte length prefix for parsing
1598            let parsed = Message::from_payload(Bytes::copy_from_slice(&buf[4..n])).unwrap();
1599            // For Piece with split data, wire format concatenates into data_0
1600            match &msg {
1601                Message::Piece {
1602                    index,
1603                    begin,
1604                    data_0,
1605                    data_1,
1606                } if !data_1.is_empty() => {
1607                    let mut combined = Vec::from(data_0.as_ref());
1608                    combined.extend_from_slice(data_1.as_ref());
1609                    let expected = Message::Piece {
1610                        index: *index,
1611                        begin: *begin,
1612                        data_0: Bytes::from(combined),
1613                        data_1: Bytes::new(),
1614                    };
1615                    assert_eq!(parsed, expected, "roundtrip mismatch for split Piece");
1616                }
1617                _ => {
1618                    assert_eq!(msg, parsed, "roundtrip mismatch for {msg:?}");
1619                }
1620            }
1621        }
1622    }
1623
1624    #[test]
1625    fn encode_to_slice_matches_encode_into() {
1626        for msg in all_message_variants() {
1627            let mut slice_buf = [0u8; 4096];
1628            let n = msg.encode_to_slice(&mut slice_buf);
1629
1630            let mut bytes_buf = BytesMut::new();
1631            msg.encode_into(&mut bytes_buf);
1632
1633            assert_eq!(
1634                &slice_buf[..n],
1635                &bytes_buf[..],
1636                "encode_to_slice vs encode_into mismatch for {msg:?}"
1637            );
1638        }
1639    }
1640
1641    #[test]
1642    fn wire_len_matches_encoded_size() {
1643        for msg in all_message_variants() {
1644            let expected = msg.to_bytes().len();
1645            assert_eq!(msg.wire_len(), expected, "wire_len mismatch for {msg:?}");
1646        }
1647    }
1648
1649    #[test]
1650    fn wire_len_large_bitfield() {
1651        // Bitfield for a torrent with >131K pieces (16,384 bytes of bitfield data).
1652        let bits = vec![0xFFu8; 20_000];
1653        let msg = Message::Bitfield(Bytes::from(bits.clone()));
1654        assert_eq!(msg.wire_len(), 5 + bits.len());
1655        assert_eq!(msg.wire_len(), msg.to_bytes().len());
1656    }
1657}
irontide_wire/message.rs

irontide_wire/
message.rs
