From d258a9afe299027403723a2c35c1981592778129 Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Wed, 6 Dec 2023 23:25:52 +0000
Subject: [PATCH] Processing cancel messages

---
 crates/librqbit/src/torrent_state/live/mod.rs |  5 +++-
 crates/peer_binary_protocol/src/lib.rs        | 25 +++++++++++++++----
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/crates/librqbit/src/torrent_state/live/mod.rs b/crates/librqbit/src/torrent_state/live/mod.rs
index 13bc197..a7c0821 100644
--- a/crates/librqbit/src/torrent_state/live/mod.rs
+++ b/crates/librqbit/src/torrent_state/live/mod.rs
@@ -856,7 +856,10 @@ impl<'a> PeerConnectionHandler for &'a PeerHandler {
             }
             Message::Have(h) => self.on_have(h),
             Message::NotInterested => {
-                debug!("received \"not interested\", but we don't care yet")
+                debug!("received \"not interested\", but we don't process it yet")
+            }
+            Message::Cancel(_) => {
+                debug!("received \"cancel\", but we don't process it yet")
             }
             message => {
                 warn!("received unsupported message {:?}, ignoring", message);
diff --git a/crates/peer_binary_protocol/src/lib.rs b/crates/peer_binary_protocol/src/lib.rs
index b99edc5..4811ab4 100644
--- a/crates/peer_binary_protocol/src/lib.rs
+++ b/crates/peer_binary_protocol/src/lib.rs
@@ -40,6 +40,7 @@ const MSGID_HAVE: u8 = 4;
 const MSGID_BITFIELD: u8 = 5;
 const MSGID_REQUEST: u8 = 6;
 const MSGID_PIECE: u8 = 7;
+const MSGID_CANCEL: u8 = 8;
 const MSGID_EXTENDED: u8 = 20;
 
 pub const MY_EXTENDED_UT_METADATA: u8 = 3;
@@ -169,6 +170,7 @@ impl From<anyhow::Error> for MessageDeserializeError {
 #[derive(Debug)]
 pub enum Message<ByteBuf: std::hash::Hash + Eq> {
     Request(Request),
+    Cancel(Request),
     Bitfield(ByteBuf),
     KeepAlive,
     Have(u32),
@@ -200,6 +202,7 @@ where
     fn clone_to_owned(&self) -> Self::Target {
         match self {
             Message::Request(req) => Message::Request(*req),
+            Message::Cancel(req) => Message::Cancel(*req),
             Message::Bitfield(b) => Message::Bitfield(b.clone_to_owned()),
             Message::Choke => Message::Choke,
             Message::Unchoke => Message::Unchoke,
@@ -240,7 +243,7 @@ where
 {
     pub fn len_prefix_and_msg_id(&self) -> (u32, u8) {
         match self {
-            Message::Request(_) => (LEN_PREFIX_REQUEST, MSGID_REQUEST),
+            Message::Request(_) | Message::Cancel(_) => (LEN_PREFIX_REQUEST, MSGID_REQUEST),
             Message::Bitfield(b) => (1 + b.as_ref().len() as u32, MSGID_BITFIELD),
             Message::Choke => (LEN_PREFIX_CHOKE, MSGID_CHOKE),
             Message::Unchoke => (LEN_PREFIX_UNCHOKE, MSGID_UNCHOKE),
@@ -270,7 +273,7 @@ where
         let ser = bopts();
 
         match self {
-            Message::Request(request) => {
+            Message::Request(request) | Message::Cancel(request) => {
                 const MSG_LEN: usize = PREAMBLE_LEN + 12;
                 out.resize(MSG_LEN, 0);
                 debug_assert_eq!(out[PREAMBLE_LEN..].len(), 12);
@@ -411,16 +414,28 @@ where
                     }
                 }
             }
-            MSGID_REQUEST => {
+            MSGID_REQUEST | MSGID_CANCEL => {
                 let expected_len = 12;
                 match rest.get(..expected_len) {
                     Some(b) => {
                         let request = decoder_config.deserialize::<Request>(b).unwrap();
-                        Ok((Message::Request(request), PREAMBLE_LEN + expected_len))
+                        let req = if msg_id == MSGID_REQUEST {
+                            Message::Request(request)
+                        } else {
+                            Message::Cancel(request)
+                        };
+                        Ok((req, PREAMBLE_LEN + expected_len))
                     }
                     None => {
                         let missing = expected_len - rest.len();
-                        Err(MessageDeserializeError::NotEnoughData(missing, "request"))
+                        Err(MessageDeserializeError::NotEnoughData(
+                            missing,
+                            if msg_id == MSGID_REQUEST {
+                                "request"
+                            } else {
+                                "cancel"
+                            },
+                        ))
                     }
                 }
             }