Handshake fixes

src/main.rs

@@ -19,18 +19,18 @@ mod tracker;
 mod types;
 
 fn main() {
-    let path = PathBuf::from("data/academic_test.torrent");
+    let path = PathBuf::from("data/knoppix.torrent");
     let bencoded = fs::read(path).unwrap();
     let metainfo_dict = match parse_bencoded(bencoded) {
         (Some(metadata), left) if left.is_empty() => metadata,
         _ => panic!("metadata file parsing error"),
     };
-    println!("{metainfo_dict:?}");
+    println!("metainfo dict: {metainfo_dict:?}");
     let metainfo = match Metainfo::try_from(metainfo_dict.clone()) {
         Ok(info) => info,
         Err(e) => panic!("metadata file structure error: {e}"),
     };
-    println!("{metainfo:?}");
+    println!("metainfo: {metainfo:?}");
     let info_dict_str = match metainfo_dict {
         bencode::BencodeValue::Dict(d) => d.get("info").unwrap().encode(),
         _ => unreachable!(),
@@ -52,11 +52,8 @@ fn main() {
     if let TrackerResponse::Success(resp) = tracker_response {
         for p in resp.peers {
             match handshake(&p, &info_hash, &peer_id) {
-                Ok(_) => {
-                    println!("successfull handshake with peer {:?}", p);
-                    break;
-                }
-                Err(e) => eprintln!("{}", e),
+                Ok(_) => println!("successfull handshake with peer {:?}", p),
+                Err(e) => eprintln!("handshake error: {}", e),
             }
         }
     }

src/peer.rs

@@ -1,27 +1,74 @@
-use std::io::{self, Read, Write};
+use std::io::{self, BufReader, Read, Write};
 use std::net::{IpAddr, SocketAddr, TcpStream};
 use std::str::FromStr;
-use std::thread::sleep;
 use std::time::Duration;
 
 use crate::hex::hex;
 use crate::tracker::TrackerPeer;
 use crate::types::ByteString;
 
+pub struct HandshakePacket {
+    info_hash: Vec<u8>,
+    peer_id: Vec<u8>,
+}
+
+impl From<HandshakePacket> for Vec<u8> {
+    fn from(value: HandshakePacket) -> Self {
+        let pstr = "BitTorrent protocol";
+        let pstrlen = &[pstr.len() as u8];
+        let reserved = &[0u8; 8];
+        [
+            pstrlen,
+            pstr.as_bytes(),
+            reserved,
+            &value.info_hash,
+            &value.peer_id,
+        ]
+        .concat()
+    }
+}
+
+impl TryFrom<Vec<u8>> for HandshakePacket {
+    type Error = String;
+
+    fn try_from(value: Vec<u8>) -> Result<Self, Self::Error> {
+        if value.len() != 68 {
+            return Err(format!("invalid handshake len: {}", value.len()));
+        }
+        let pstrlen = &value.as_slice()[0..1];
+        if pstrlen != [19u8] {
+            return Err(format!("invalid pstrlen: {}", hex(pstrlen)));
+        }
+        let pstr = &value.as_slice()[1..20];
+        if pstr != "BitTorrent protocol".as_bytes() {
+            return Err(format!("invalid pstr: {}", hex(pstr)));
+        }
+        Ok(HandshakePacket {
+            info_hash: value.as_slice()[29..48].to_vec(),
+            peer_id: value.as_slice()[49..68].to_vec(),
+        })
+    }
+}
+
 pub fn handshake(
     peer: &TrackerPeer,
     info_hash: &ByteString,
     peer_id: &ByteString,
-) -> io::Result<()> {
+) -> io::Result<HandshakePacket> {
+    let timeout = Duration::new(4, 0);
     println!("connecting to peer {peer:?}");
-    let timeout = Duration::new(2, 0);
     let mut stream = TcpStream::connect_timeout(
         &SocketAddr::new(IpAddr::from_str(&peer.ip).unwrap(), peer.port as u16),
         timeout,
     )?;
     stream.set_read_timeout(Some(timeout))?;
     stream.set_write_timeout(Some(timeout))?;
-    let handshake = handshake_packet(info_hash, peer_id);
+    let handshake: Vec<u8> = HandshakePacket {
+        info_hash: info_hash.clone(),
+        peer_id: peer_id.clone(),
+    }
+    .into();
+
     println!("writing handshake {}", hex(&handshake.to_vec()));
     match stream.write_all(&handshake) {
         Err(e) => {
@@ -31,34 +78,22 @@ pub fn handshake(
         _ => println!("write ok"),
     };
     stream.flush()?;
-    let mut reader = stream;
-    let mut read_packet = vec![];
-    println!("reading response");
-    let mut retry = 0;
-    loop {
-        if retry > 3 {
-            return Err(io::Error::new(io::ErrorKind::Other, "read timeout"));
-        };
-        match reader.read_to_end(&mut read_packet) {
-            Err(e) => {
-                eprintln!("read error: {}", e);
-                return Err(e);
-            }
-            Ok(n) if n > 0 => {
-                println!("peer response: {}", hex(&read_packet.to_vec()));
-            }
-            _ => {
-                println!("no data");
-                retry += 1;
-                sleep(Duration::new(1, 0));
+
+    let mut reader = BufReader::new(stream);
+    let mut read_packet = [0u8; 68];
+    println!("reading handshake");
+    match reader.read_exact(&mut read_packet) {
+        Ok(_) => {
+            let msg: Vec<u8> = read_packet.to_vec();
+            println!("peer response: {}", hex(&msg));
+            match HandshakePacket::try_from(msg) {
+                Ok(p) => Ok(p),
+                Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
             }
-        };
+        }
+        Err(e) => {
+            eprintln!("read error: {}", e);
+            Err(e)
+        }
     }
 }
-
-pub fn handshake_packet(info_hash: &ByteString, peer_id: &ByteString) -> Vec<u8> {
-    let pstr = "BitTorrent protocol";
-    let pstrlen = &[pstr.len() as u8];
-    let reserved = &[0u8; 8];
-    [pstrlen, pstr.as_bytes(), reserved, &info_hash, &peer_id].concat()
-}

src/tracker.rs

@@ -203,8 +203,7 @@ pub fn tracker_request(
     announce: String,
     request: TrackerRequest,
 ) -> Result<TrackerResponse, String> {
-    let params = request.to_params();
-    let query = format!(
+    let params = format!(
         "?{}",
         request
             .to_params()
@@ -213,14 +212,15 @@ pub fn tracker_request(
             .collect::<Vec<_>>()
             .join("&")
     );
-    println!("params: {params:?}");
+    let url = format!("{announce}{params}");
+    println!("url: {url}");
     let resp = Client::new()
-        .get(format!("{announce}{query}"))
+        .get(url)
         .send()
         .map_err(|e| format!("request error: {}", e))?
         .bytes()
         .map_err(|e| format!("request body error: {}", e))?;
-    println!("{}", String::from_utf8_lossy(&resp));
+    println!("raw response: {}", String::from_utf8_lossy(&resp));
     let resp_dict = parse_bencoded(resp.to_vec())
         .0
         .ok_or("Malformed response")?;