bencode_decode/
lib.rs

1use std::{collections::BTreeMap, io::Read, result::Result};
2// Bencoding spec
3// https://wiki.theory.org/index.php/BitTorrentSpecification#Bencoding
4
5#[derive(PartialEq, Ord, PartialOrd, Eq, Debug, Clone)]
6pub enum Value {
7    ByteString(Vec<u8>),
8    Integer(i64),
9    List(Vec<Value>),
10    Dictionary(BTreeMap<Vec<u8>, Value>),
11}
12
13#[derive(PartialEq, Debug)]
14pub enum ParseResult {
15    ValueType(Value),
16    ListStart,
17    DictStart,
18    End,
19    EOF,
20}
21
22/// Constructs a `Parser` for bencoded data from a reader implementing
23/// `std::io::read`. The only exposed interface is an iterator, which
24/// will emit parsed tokens `ParseResult` up until (but not including)
25/// EOF.
26///
27/// ```
28/// use bencode_decode::Parser;
29/// let input = std::io::Cursor::new(
30///            "d9:publisher3:bob17:publisher-webpage15:www.example.com18:publisher.location4:homee"
31///                .to_string()
32///                .into_bytes(),
33///        );
34/// let mut parser = Parser::new(input);
35/// for item in parser {
36///     println!("{:?}", item);
37/// }
38/// ```
39pub struct Parser<R: Read> {
40    reader: R,
41}
42impl<R: Read> Parser<R> {
43    pub fn new(reader: R) -> Self {
44        Self { reader }
45    }
46}
47
48impl<R: Read> Iterator for Parser<R> {
49    type Item = ParseResult;
50    fn next(&mut self) -> Option<Self::Item> {
51        let res = parse(&mut self.reader).ok();
52        if res == Some(ParseResult::EOF) {
53            None
54        } else {
55            res
56        }
57    }
58}
59
60use ParseResult::*;
61use Value::*;
62/// Given a token parser `parser`, will try to decode `ParseResult` into
63/// `Value`s. This function does obviously not attempt to drain the passed
64/// reader instance, but rather expects one top-level value to parse form.
65///
66/// ```
67/// use bencode_decode::{Parser, decode};
68/// use std::fs::File;
69///
70/// let f = File::open("./test/ubuntu-18.04.4-live-server-amd64.iso.torrent").unwrap();
71/// let mut parser = Parser::new(f);
72/// let res = decode(&mut parser, None).unwrap();
73/// ```
74pub fn decode<R: Read>(parser: &mut Parser<R>, current: Option<ParseResult>) -> Option<Value> {
75    match current.or_else(|| parser.next()) {
76        Some(ValueType(val)) => Some(val),
77        Some(t @ DictStart) | Some(t @ ListStart) => {
78            let mut data = vec![];
79            let mut next = parser.next().expect("Unexpected EOF");
80            while next != End {
81                data.push(decode(parser, Some(next)).unwrap());
82                next = parser.next().expect("Unexpected EOF");
83            }
84            if t == ListStart {
85                Some(Value::List(data))
86            } else {
87                let mut map = BTreeMap::new();
88                let mut input = data.into_iter();
89                while let (Some(ByteString(key)), Some(value)) = (input.next(), input.next()) {
90                    map.insert(key, value);
91                }
92                Some(Dictionary(map))
93            }
94        }
95        Some(End) => unreachable!(),
96        Some(EOF) => unreachable!(),
97        None => None,
98    }
99}
100
101fn parse<R: Read>(reader: &mut R) -> Result<ParseResult, Box<dyn std::error::Error>> {
102    let mut buf = [0; 1];
103    let mut vec = vec![];
104    loop {
105        let read_bytes = reader.read(&mut buf)?;
106        if read_bytes == 0 {
107            return Ok(EOF);
108        }
109        match buf[0] {
110            n @ b'0'..=b'9' => vec.push(n),
111            b':' => {
112                let size = String::from_utf8(vec)?.parse()?;
113                let mut str = vec![0; size];
114                reader.read_exact(&mut str)?;
115                return Ok(ValueType(ByteString(str)));
116            }
117            b'i' => {
118                let mut b = [0; 1];
119                reader.read_exact(&mut b)?;
120                let mut vec = vec![];
121                while b[0] != b'e' {
122                    vec.push(b[0]);
123                    reader.read_exact(&mut b)?;
124                }
125                let int: i64 = String::from_utf8(vec)?.parse()?;
126                return Ok(ValueType(Integer(int)));
127            }
128            b'e' => return Ok(End),
129            b'l' => return Ok(ListStart),
130            b'd' => return Ok(DictStart),
131            _ => unreachable!("unexpected token"),
132        }
133    }
134}
135
136#[cfg(test)]
137mod test {
138    use super::*;
139    use std::fs::File;
140    #[test]
141    fn torrent() {
142        let f = File::open("./test/ubuntu-18.04.4-live-server-amd64.iso.torrent").unwrap();
143        let mut parser = Parser::new(f);
144        let res = decode(&mut parser, None).unwrap();
145        if let Value::Dictionary(x) = res {
146            if let Value::Dictionary(y) = x.get(&b"info".to_vec()).unwrap() {
147                let path = y.get(&b"name".to_vec()).unwrap();
148                let length = y.get(&b"length".to_vec()).unwrap();
149                if let (Value::ByteString(path), Value::Integer(length)) = (path, length) {
150                    let path = String::from_utf8_lossy(path);
151                    println!("{} -> {} bytes", path, length);
152                    assert_eq!(path, "ubuntu-18.04.4-live-server-amd64.iso");
153                    assert_eq!(*length, 912_261_120);
154                }
155            }
156        }
157    }
158
159    #[test]
160    fn spec() {
161        let input = std::io::Cursor::new(
162            "d9:publisher3:bob17:publisher-webpage15:www.example.com18:publisher.location4:homee"
163                .to_string()
164                .into_bytes(),
165        );
166        let mut parser = Parser::new(input);
167        let res = decode(&mut parser, None).unwrap();
168        let mut map = BTreeMap::new();
169        vec![
170            ("publisher", "bob"),
171            ("publisher-webpage", "www.example.com"),
172            ("publisher.location", "home"),
173        ]
174        .into_iter()
175        .for_each(|(k, v)| {
176            map.insert(
177                k.as_bytes().to_vec(),
178                Value::ByteString(v.as_bytes().to_vec()),
179            );
180        });
181
182        assert_eq!(res, Value::Dictionary(map));
183    }
184}