just commit stuff to start fresh

src/client/mod.rs

@@ -1,6 +1,6 @@
 use std::{
     collections::HashMap,
-    fmt::{Debug, Display},
+    fmt::Display,
     io::{Read, Write},
     net::TcpStream,
     time::Duration,
@@ -10,6 +10,7 @@ use crate::tracker::{Peer, Tracker};
 
 const PSTR: &[u8; 19] = b"BitTorrent protocol";
 const HANDSHAKE_LEN: usize = 49 + PSTR.len();
+const REQUEST_LEN: u32 = 2 << 14;
 
 pub struct PeerConnectionError {
     pub peer: Peer,
@@ -20,6 +21,36 @@ pub enum HandshakePhase {
     Receive,
 }
 
+pub enum MessageId {
+    Choke = 0,
+    Unchoke = 1,
+    Interested = 2,
+    NotInterested = 3,
+    Have = 4,
+    Bitfield = 5,
+    Request = 6,
+    Piece = 7,
+    Cancel = 8,
+    Port = 9,
+}
+
+impl MessageId {
+    fn value(&self) -> u8 {
+        match self {
+            MessageId::Choke => 0,
+            MessageId::Unchoke => 1,
+            MessageId::Interested => 2,
+            MessageId::NotInterested => 3,
+            MessageId::Have => 4,
+            MessageId::Bitfield => 5,
+            MessageId::Request => 6,
+            MessageId::Piece => 7,
+            MessageId::Cancel => 8,
+            MessageId::Port => 9,
+        }
+    }
+}
+
 impl Display for HandshakePhase {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -36,13 +67,20 @@ pub struct HandshakeError {
     message: String,
 }
 
+pub struct SendMessageError {
+    peer: Peer,
+    message: Message,
+    error: String,
+}
+
 pub enum ClientError {
     ValidateHandshakeError(String),
     GetPeersError(String),
     HandshakeError(HandshakeError),
+    SendMessageError(SendMessageError),
 }
 
-impl Debug for ClientError {
+impl Display for ClientError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             ClientError::ValidateHandshakeError(e) => write!(f, "ValidateHandshakeError: {}", e),
@@ -58,21 +96,99 @@ impl Debug for ClientError {
                     .map(|b| format!("{:02x}", b))
                     .collect::<String>()
             ),
+            ClientError::SendMessageError(e) => write!(
+                f,
+                "SendMessageError: Peer: {}, Message: {:02x?}, Error: {}",
+                e.peer,
+                e.message.serialize(),
+                e.error
+            ),
+        }
+    }
+}
+
+struct Message {
+    len: u32,
+    id: u8,
+    payload: Vec<u8>,
+}
+
+impl Message {
+    fn new(id: MessageId, payload: &Vec<u8>) -> Self {
+        Self {
+            len: payload.len() as u32 + 1,
+            id: id.value(),
+            payload: payload.clone(),
+        }
+    }
+
+    fn serialize(&self) -> Vec<u8> {
+        let mut message = Vec::new();
+        message.extend_from_slice(&self.len.to_be_bytes());
+        message.push(self.id);
+        message.extend_from_slice(&self.payload);
+        message
+    }
+}
+
+impl Clone for Message {
+    fn clone(&self) -> Self {
+        Self {
+            len: self.len,
+            id: self.id,
+            payload: self.payload.clone(),
         }
     }
 }
 
 pub struct Client {
     tracker: Tracker,
     connections: HashMap<Vec<u8>, TcpStream>,
+    bitfield: Vec<u8>,
 }
 
 impl Client {
     pub fn new(tracker: Tracker) -> Self {
+        // divide the number of pieces by 8 to get the number of bytes needed to represent the bitfield
+        let bitfield = vec![0u8; tracker.get_metainfo().get_peices().len().div_ceil(8)];
         Self {
             tracker,
             connections: HashMap::new(),
+            bitfield,
+        }
+    }
+
+    pub async fn download(&mut self) -> Result<(), ClientError> {
+        println!("Starting download...");
+        // self.connect_to_peers(30).await?;
+        self.connect_to_peers(1).await?;
+        self.send_message(Message::new(MessageId::Bitfield, &self.bitfield))?;
+        self.send_message(Message::new(MessageId::Interested, &Vec::new()))?;
+        Ok(())
+    }
+
+    fn send_message(&mut self, message: Message) -> Result<(), ClientError> {
+        let serialized_message = message.serialize();
+        println!("Sending message: {:?}", serialized_message);
+        for (_, stream) in self.connections.iter_mut() {
+            stream.write_all(&serialized_message).map_err(|e| {
+                ClientError::SendMessageError(SendMessageError {
+                    peer: Peer {
+                        peer_id: None,
+                        addr: stream.peer_addr().unwrap(),
+                    },
+                    message: message.clone(),
+                    error: e.to_string(),
+                })
+            })?;
+
+            let mut response = vec![0u8; serialized_message.len()];
+            stream.read_exact(&mut response).map_err(|e| {
+                ClientError::GetPeersError(format!("Failed to read response: {}", e))
+            })?;
+            println!("Response: {:?}", response);
         }
+        Ok(())
     }
 
     fn get_handshake(&self) -> Result<Vec<u8>, ClientError> {
@@ -90,8 +206,9 @@ impl Client {
         handshake.extend_from_slice(PSTR);
         handshake.extend_from_slice(&[0; 8]);
         handshake.extend_from_slice(&info_hash);
-        handshake.extend_from_slice(peer_id.as_bytes());
+        handshake.extend_from_slice(&peer_id);
 
+        #[cfg(debug_assertions)]
         assert_eq!(handshake.len(), HANDSHAKE_LEN);
         Ok(handshake)
     }
@@ -127,7 +244,8 @@ impl Client {
         Ok(peer_id)
     }
 
-    pub async fn connect_to_peers(&mut self, min_connections: usize) -> Result<(), ClientError> {
+    async fn connect_to_peers(&mut self, min_connections: usize) -> Result<(), ClientError> {
+        println!("Connecting to peers...");
         while self.connections.len() < min_connections {
             let mut handles = Vec::new();
             for peer in self
@@ -137,7 +255,7 @@ impl Client {
                 .map_err(|_| ClientError::GetPeersError(String::from("Failed to get peers")))?
             {
                 if self.connections.len() >= min_connections {
-                    break;
+                    return Ok(());
                 }
 
                 let handshake = self.get_handshake()?;
@@ -183,26 +301,27 @@ impl Client {
 
             for handle in handles {
                 if self.connections.len() >= min_connections {
-                    break;
-                }
-
-                match handle
-                    .await
-                    .map_err(|e| ClientError::GetPeersError(String::from(e.to_string())))?
-                {
-                    Ok((peer_id, stream)) => {
-                        println!(
-                            "Connected to peer: {} at {}",
-                            String::from_utf8_lossy(&peer_id),
-                            stream
-                                .peer_addr()
-                                .map(|addr| addr.to_string())
-                                .unwrap_or("Unknown".to_string())
-                        );
-                        self.connections.insert(peer_id, stream);
-                    }
-                    Err(e) => {
-                        eprintln!("{:?}", e);
+                    handle.abort();
+                } else {
+                    match handle
+                        .await
+                        .map_err(|e| ClientError::GetPeersError(String::from(e.to_string())))?
+                    {
+                        Ok((peer_id, stream)) => {
+                            println!(
+                                "Connected to peer: {} at {}",
+                                String::from_utf8_lossy(&peer_id),
+                                stream
+                                    .peer_addr()
+                                    .map(|addr| addr.to_string())
+                                    .unwrap_or("Unknown".to_string())
+                            );
+                            self.connections.insert(peer_id, stream);
+                        }
+                        Err(e) => {
+                            #[cfg(debug_assertions)]
+                            eprintln!("{}", e);
+                        }
                     }
                 }
             }

src/main.rs

@@ -40,8 +40,8 @@ async fn main() {
     let tracker = Tracker::new(bencode_value).expect("Failed to create tracker");
     let mut client = Client::new(tracker);
 
-    client
-        .connect_to_peers(10)
-        .await
-        .expect("Failed to connect to peers");
+    match client.download().await {
+        Ok(()) => println!("Download completed"),
+        Err(e) => eprintln!("Error downloading: {}", e),
+    }
 }

src/metainfo/mod.rs

@@ -8,7 +8,7 @@ use crate::bencode::{BencodeString, BencodeValue};
 #[derive(Debug, PartialEq)]
 pub struct BaseInfo {
     // shared by both single and multi file mode
-    pub pieces: Vec<u8>,
+    pub pieces: Vec<Vec<u8>>,
     pub piece_length: i64,
     pub private: Option<i64>,
 }
@@ -103,9 +103,22 @@ impl Metainfo {
         Ok(result.to_vec())
     }
 
+    pub fn get_peices(&self) -> &Vec<Vec<u8>> {
+        match self.info {
+            Info::SingleFile(ref info) => &info.base_info.pieces,
+            Info::MultiFile(ref info) => &info.base_info.pieces,
+        }
+    }
+
     fn dict_to_base_info(dict: &BTreeMap<String, BencodeValue>) -> Result<BaseInfo, MetaInfoError> {
         let pieces = match dict.get("pieces") {
-            Some(BencodeValue::String(BencodeString::Bytes(b))) => b.clone(),
+            Some(BencodeValue::String(BencodeString::Bytes(b))) => {
+                let mut pieces = Vec::new();
+                for i in (0..b.len()).step_by(20) {
+                    pieces.push(b[i..i + 20].to_vec());
+                }
+                pieces
+            }
             _ => {
                 return Err(MetaInfoError::InvalidAttribute(AttributeError {
                     content: BencodeValue::Dict(dict.clone()),

src/tracker/mod.rs

@@ -52,13 +52,13 @@ impl Debug for TrackerError {
 #[derive(Debug)]
 pub struct Tracker {
     metainfo: Metainfo,
-    peer_id: String,
+    peer_id: Vec<u8>,
 }
 
 #[derive(Debug)]
 pub struct Peer {
     pub addr: SocketAddr,
-    pub peer_id: Option<String>,
+    pub peer_id: Option<Vec<u8>>,
 }
 
 impl Clone for Peer {
@@ -73,7 +73,7 @@ impl Clone for Peer {
 impl Display for Peer {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match &self.peer_id {
-            Some(peer_id) => write!(f, "{}: {}", peer_id, self.addr),
+            Some(peer_id) => write!(f, "{}: {}", String::from_utf8_lossy(peer_id), self.addr),
             None => write!(f, "{}", self.addr),
         }
     }
@@ -116,8 +116,8 @@ impl Tracker {
         &self.metainfo
     }
 
-    pub fn peer_id(&self) -> String {
-        self.peer_id.to_string()
+    pub fn peer_id(&self) -> Vec<u8> {
+        self.peer_id.clone()
     }
 
     pub async fn get_peers(&self) -> Result<Peers, TrackerError> {
@@ -177,6 +177,9 @@ impl Tracker {
                                 .get("peer id")
                                 .map(|peer_id| match peer_id {
                                     BencodeValue::String(BencodeString::String(peer_id)) => {
+                                        Some(peer_id.bytes().collect::<Vec<u8>>())
+                                    }
+                                    BencodeValue::String(BencodeString::Bytes(peer_id)) => {
                                         Some(peer_id.clone())
                                     }
                                     _ => None,
@@ -322,7 +325,13 @@ impl Tracker {
             url::form_urlencoded::byte_serialize(&info_hash).collect::<String>();
 
         url.push_str(format!("?info_hash={}", url_encoded_info_hash).as_str());
-        url.push_str(format!("&peer_id={}", self.peer_id).as_str());
+        url.push_str(
+            format!(
+                "&peer_id={}",
+                String::from_utf8(self.peer_id.clone()).unwrap()
+            )
+            .as_str(),
+        );
         url.push_str("&port=6881");
 
         println!("GET {}", &url);
@@ -351,11 +360,11 @@ impl Tracker {
         Tracker::to_tracker_response(&parsed_bencode)
     }
 
-    fn get_peer_id() -> String {
-        let mut peer_id = String::from("-RT0001-");
+    fn get_peer_id() -> Vec<u8> {
+        let mut peer_id = Vec::from(b"-rT0001-");
         let mut rng = rand::thread_rng();
         for _ in 0..(20 - peer_id.len()) {
-            let random_char = (rng.gen_range(0..26) + 97) as u8 as char;
+            let random_char = (rng.gen_range(0..26) + 97) as u8;
             peer_id.push(random_char);
         }
         peer_id