From 81017b3f90f5f038eae57f2546a981778c130d3f Mon Sep 17 00:00:00 2001
From: Faerbit <faerbit@posteo.net>
Date: Sun, 21 Jul 2024 13:32:10 +0200
Subject: [PATCH] [bencode] Decode bytes

---
 src/bencode.rs | 81 +++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 63 insertions(+), 18 deletions(-)

diff --git a/src/bencode.rs b/src/bencode.rs
index 7656630..89ba5d2 100644
--- a/src/bencode.rs
+++ b/src/bencode.rs
@@ -24,6 +24,16 @@ struct Torrent {
 #[derive(Debug, PartialEq, Eq, Hash)]
 pub struct ByteString(Vec<u8>);
 
+impl ByteString {
+    pub fn from_str(input: &str) -> Self {
+        ByteString(input.as_bytes().to_vec())
+    }
+
+    pub fn from_slice(input: &[u8]) -> Self {
+        ByteString(input.to_vec())
+    }
+}
+
 #[derive(Debug, PartialEq)]
 pub enum Bencode {
     Integer(i64),
@@ -42,21 +52,14 @@ enum BencodeType {
 
 impl Bencode {
     pub fn decode(input: &str) -> Result<Bencode> {
-        Self::decode_slice(input.as_bytes())
+        Self::decode_type(input.as_bytes())
     }
 
-    fn decode_slice(input: &[u8]) -> Result<Bencode>{
+    fn decode_type(input: &[u8]) -> Result<Bencode>{
         if input.len() == 0 {
             return Err(anyhow!("Empty string is not valid bencode"));
         }
         let char = input[0] as char;
-        // let type_  = match char {
-        //     'i'  => Ok(Discriminant(Bencode::Integer)),
-        //     'l'  => Ok(Discriminant(Bencode::List)),
-        //     'd'  => Ok(Discriminant(Bencode::Dict)),
-        //     '0'..='9' => Ok(Discriminant(Bencode::Integer)),
-        //     _ => Err(anyhow!("Unknown bencoding char: {char}")),
-        // }?;
         let type_  = match char {
             'i'  => Ok(BencodeType::Integer),
             'l'  => Ok(BencodeType::List),
@@ -68,37 +71,60 @@ impl Bencode {
             BencodeType::Integer | BencodeType::List | BencodeType::Dict=> 'e',
             BencodeType::Bytes => ':'
         };
-        let mut end_pos = 1;
+        let mut end_pos= 1;
         for &c in &input[1..] {
             if c == end_char as u8 {
                 break
             }
             end_pos += 1;
         }
+        let to_decode = match type_ {
+            BencodeType::Integer | BencodeType::List | BencodeType::Dict=> &input[1..end_pos],
+            BencodeType::Bytes => {
+                let bytes_len = Self::decode_int_only(&input[0..end_pos])? as usize;
+                let bytes_start= end_pos + 1;
+                &input[bytes_start..bytes_start + bytes_len]
+            },
+        };
+
         match type_ {
-            BencodeType::Integer => Self::decode_int(&input[1..end_pos]),
-            BencodeType::List => Self::dummy(),
-            BencodeType::Dict => Self::dummy(),
-            BencodeType::Bytes => Self::dummy(),
+            BencodeType::Integer => Self::decode_int(to_decode),
+            BencodeType::List => Self::dummy(to_decode),
+            BencodeType::Dict => Self::dummy(to_decode),
+            BencodeType::Bytes => Self::decode_bytes(to_decode),
         }
     }
 
-    fn decode_int(input: &[u8]) -> Result<Bencode> {
+    fn decode_int_only(input: &[u8]) -> Result<i64> {
         let int_str = std::str::from_utf8(input)?;
-        let int = int_str.parse::<i64>()?;
+        int_str.parse::<i64>().map_err(anyhow::Error::msg)
+    }
+
+    fn decode_int(input: &[u8]) -> Result<Bencode> {
+        let int = Self::decode_int_only(input)?;
         Ok(Bencode::Integer(int))
     }
 
-    fn dummy() -> Result<Bencode> {
+    fn decode_bytes(input: &[u8]) -> Result<Bencode> {
+        Ok(Bencode::Bytes(ByteString::from_slice(input)))
+    }
+
+    fn dummy(_input: &[u8]) -> Result<Bencode> {
         Ok(Bencode::Integer(42))
     }
 }
 
 #[cfg(test)]
 mod tests {
-    // Note this useful idiom: importing names from outer (for mod tests) scope.
     use super::*;
 
+    #[test]
+    fn test_integer_only() {
+        assert_eq!(Bencode::decode_int_only("42".as_bytes()).unwrap(), 42);
+        assert_eq!(Bencode::decode_int_only("17".as_bytes()).unwrap(), 17);
+        assert_eq!(Bencode::decode_int_only("-17".as_bytes()).unwrap(), -17);
+    }
+
     #[test]
     fn test_integer() {
         assert_eq!(Bencode::decode_int("42".as_bytes()).unwrap(), Bencode::Integer(42));
@@ -113,4 +139,23 @@ mod tests {
         assert_eq!(Bencode::decode("i-17e").unwrap(), Bencode::Integer(-17));
     }
 
+    #[test]
+    fn test_bytes() {
+        assert_eq!(Bencode::decode_bytes("hallo".as_bytes()).unwrap(), Bencode::Bytes(ByteString::from_str("hallo")));
+        assert_eq!(Bencode::decode_bytes("tschüss".as_bytes()).unwrap(), Bencode::Bytes(ByteString::from_str("tschüss")));
+        assert_eq!(Bencode::decode_bytes("💩".as_bytes()).unwrap(), Bencode::Bytes(ByteString::from_str("💩")));
+    }
+
+    #[test]
+    fn test_bytes_str() {
+        let str = "hallo";
+        assert_eq!(Bencode::decode(format!("{}:{str}", str.len()).as_str()).unwrap(), Bencode::Bytes(ByteString::from_str("hallo")));
+        let str = "tschüss";
+        assert_eq!(Bencode::decode(format!("{}:{str}", str.len()).as_str()).unwrap(), Bencode::Bytes(ByteString::from_str("tschüss")));
+        let str = "💩";
+        assert_eq!(Bencode::decode(format!("{}:{str}", str.len()).as_str()).unwrap(), Bencode::Bytes(ByteString::from_str("💩")));
+        let str = "hallo 💩, this is a long text";
+        assert_eq!(Bencode::decode(format!("{}:{str}", str.len()).as_str()).unwrap(), Bencode::Bytes(ByteString::from_str("hallo 💩, this is a long text")));
+    }
+
 }