webtorrent/
bencode_parser.rs

1use crate::error::{Result, WebTorrentError};
2use bytes::Bytes;
3use std::collections::HashMap;
4
5/// Simple bencode value
6#[derive(Debug, Clone)]
7pub enum BencodeValue {
8    Integer(i64),
9    Bytes(Bytes),
10    List(Vec<BencodeValue>),
11    Dict(HashMap<Bytes, BencodeValue>),
12}
13
14impl BencodeValue {
15    pub fn as_integer(&self) -> Option<i64> {
16        match self {
17            BencodeValue::Integer(i) => Some(*i),
18            _ => None,
19        }
20    }
21
22    pub fn as_bytes(&self) -> Option<&Bytes> {
23        match self {
24            BencodeValue::Bytes(b) => Some(b),
25            _ => None,
26        }
27    }
28
29    pub fn as_string(&self) -> Option<String> {
30        self.as_bytes()
31            .and_then(|b| String::from_utf8(b.to_vec()).ok())
32    }
33
34    pub fn as_list(&self) -> Option<&Vec<BencodeValue>> {
35        match self {
36            BencodeValue::List(l) => Some(l),
37            _ => None,
38        }
39    }
40
41    pub fn as_dict(&self) -> Option<&HashMap<Bytes, BencodeValue>> {
42        match self {
43            BencodeValue::Dict(d) => Some(d),
44            _ => None,
45        }
46    }
47
48    pub fn get(&self, key: &[u8]) -> Option<&BencodeValue> {
49        self.as_dict().and_then(|d| {
50            let key_bytes = Bytes::copy_from_slice(key);
51            d.get(&key_bytes)
52        })
53    }
54
55    pub fn encode(&self) -> Bytes {
56        let mut result = Vec::new();
57        self.encode_into(&mut result);
58        Bytes::from(result)
59    }
60
61    fn encode_into(&self, buf: &mut Vec<u8>) {
62        match self {
63            BencodeValue::Integer(i) => {
64                buf.push(b'i');
65                buf.extend_from_slice(i.to_string().as_bytes());
66                buf.push(b'e');
67            }
68            BencodeValue::Bytes(b) => {
69                buf.extend_from_slice(b.len().to_string().as_bytes());
70                buf.push(b':');
71                buf.extend_from_slice(b);
72            }
73            BencodeValue::List(l) => {
74                buf.push(b'l');
75                for item in l {
76                    item.encode_into(buf);
77                }
78                buf.push(b'e');
79            }
80            BencodeValue::Dict(d) => {
81                buf.push(b'd');
82                let mut keys: Vec<_> = d.keys().collect();
83                keys.sort();
84                for key in keys {
85                    BencodeValue::Bytes(key.clone()).encode_into(buf);
86                    d[key].encode_into(buf);
87                }
88                buf.push(b'e');
89            }
90        }
91    }
92}
93
94/// Parse bencode from bytes
95pub fn parse_bencode(data: &[u8]) -> Result<(BencodeValue, usize)> {
96    let mut pos = 0;
97    let value = parse_value(data, &mut pos)?;
98    Ok((value, pos))
99}
100
101fn parse_value(data: &[u8], pos: &mut usize) -> Result<BencodeValue> {
102    if *pos >= data.len() {
103        return Err(WebTorrentError::Bencode(
104            "Unexpected end of data".to_string(),
105        ));
106    }
107
108    match data[*pos] {
109        b'i' => parse_integer(data, pos),
110        b'l' => parse_list(data, pos),
111        b'd' => parse_dict(data, pos),
112        b'0'..=b'9' => parse_bytes(data, pos),
113        _ => Err(WebTorrentError::Bencode(format!(
114            "Unexpected byte: {}",
115            data[*pos]
116        ))),
117    }
118}
119
120fn parse_integer(data: &[u8], pos: &mut usize) -> Result<BencodeValue> {
121    *pos += 1; // skip 'i'
122    let start = *pos;
123    while *pos < data.len() && data[*pos] != b'e' {
124        *pos += 1;
125    }
126    if *pos >= data.len() {
127        return Err(WebTorrentError::Bencode(
128            "Integer not terminated".to_string(),
129        ));
130    }
131    let num_str = String::from_utf8_lossy(&data[start..*pos]);
132    let num = num_str
133        .parse::<i64>()
134        .map_err(|e| WebTorrentError::Bencode(format!("Invalid integer: {}", e)))?;
135    *pos += 1; // skip 'e'
136    Ok(BencodeValue::Integer(num))
137}
138
139fn parse_bytes(data: &[u8], pos: &mut usize) -> Result<BencodeValue> {
140    let start = *pos;
141    while *pos < data.len() && data[*pos] != b':' {
142        *pos += 1;
143    }
144    if *pos >= data.len() {
145        return Err(WebTorrentError::Bencode(
146            "Bytes length not terminated".to_string(),
147        ));
148    }
149    let len_str = String::from_utf8_lossy(&data[start..*pos]);
150    let len = len_str
151        .parse::<usize>()
152        .map_err(|e| WebTorrentError::Bencode(format!("Invalid bytes length: {}", e)))?;
153    *pos += 1; // skip ':'
154    if *pos + len > data.len() {
155        return Err(WebTorrentError::Bencode(
156            "Bytes length exceeds data".to_string(),
157        ));
158    }
159    let bytes = Bytes::copy_from_slice(&data[*pos..*pos + len]);
160    *pos += len;
161    Ok(BencodeValue::Bytes(bytes))
162}
163
164fn parse_list(data: &[u8], pos: &mut usize) -> Result<BencodeValue> {
165    *pos += 1; // skip 'l'
166    let mut list = Vec::new();
167    while *pos < data.len() && data[*pos] != b'e' {
168        list.push(parse_value(data, pos)?);
169    }
170    if *pos >= data.len() {
171        return Err(WebTorrentError::Bencode("List not terminated".to_string()));
172    }
173    *pos += 1; // skip 'e'
174    Ok(BencodeValue::List(list))
175}
176
177fn parse_dict(data: &[u8], pos: &mut usize) -> Result<BencodeValue> {
178    *pos += 1; // skip 'd'
179    let mut dict = HashMap::new();
180    while *pos < data.len() && data[*pos] != b'e' {
181        let key = match parse_value(data, pos)? {
182            BencodeValue::Bytes(b) => b,
183            _ => {
184                return Err(WebTorrentError::Bencode(
185                    "Dictionary key must be bytes".to_string(),
186                ))
187            }
188        };
189        let value = parse_value(data, pos)?;
190        dict.insert(key, value);
191    }
192    if *pos >= data.len() {
193        return Err(WebTorrentError::Bencode(
194            "Dictionary not terminated".to_string(),
195        ));
196    }
197    *pos += 1; // skip 'e'
198    Ok(BencodeValue::Dict(dict))
199}
webtorrent/bencode_parser.rs

webtorrent/
bencode_parser.rs
