#![allow(
clippy::cast_possible_truncation,
clippy::cast_possible_wrap,
clippy::cast_sign_loss,
reason = "M175: BEP 52 file tree — piece counts bounded by torrent size"
)]
use std::collections::BTreeMap;
use irontide_bencode::BencodeValue;
use crate::error::Error;
use crate::hash::Id32;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct V2FileAttr {
pub length: u64,
pub pieces_root: Option<Id32>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FileTreeNode {
File(V2FileAttr),
Directory(BTreeMap<String, Self>),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct V2FileInfo {
pub path: Vec<String>,
pub attr: V2FileAttr,
}
impl FileTreeNode {
pub fn from_bencode(value: &BencodeValue) -> Result<Self, Error> {
let dict = value
.as_dict()
.ok_or_else(|| Error::InvalidTorrent("file tree node must be a dict".into()))?;
if let Some(attr_value) = dict.get(b"".as_ref()) {
let attr = parse_file_attr(attr_value)?;
return Ok(Self::File(attr));
}
let mut children = BTreeMap::new();
for (key, child_value) in dict {
let name = String::from_utf8(key.clone())
.map_err(|_| Error::InvalidTorrent("file tree key is not valid UTF-8".into()))?;
let child = Self::from_bencode(child_value)?;
children.insert(name, child);
}
Ok(Self::Directory(children))
}
#[must_use]
pub fn flatten(&self) -> Vec<V2FileInfo> {
let mut result = Vec::new();
self.flatten_into(&mut result, &mut Vec::new());
result
}
fn flatten_into(&self, result: &mut Vec<V2FileInfo>, path: &mut Vec<String>) {
match self {
Self::File(attr) => {
result.push(V2FileInfo {
path: path.clone(),
attr: attr.clone(),
});
}
Self::Directory(children) => {
for (name, child) in children {
path.push(name.clone());
child.flatten_into(result, path);
path.pop();
}
}
}
}
#[must_use]
pub fn to_bencode(&self) -> BencodeValue {
match self {
Self::File(attr) => {
let mut file_dict = BTreeMap::new();
file_dict.insert(
b"length".to_vec(),
BencodeValue::Integer(attr.length as i64),
);
if let Some(root) = &attr.pieces_root {
file_dict.insert(
b"pieces root".to_vec(),
BencodeValue::Bytes(root.as_bytes().to_vec()),
);
}
let mut node = BTreeMap::new();
node.insert(b"".to_vec(), BencodeValue::Dict(file_dict));
BencodeValue::Dict(node)
}
Self::Directory(children) => {
let mut dict = BTreeMap::new();
for (name, child) in children {
dict.insert(name.as_bytes().to_vec(), child.to_bencode());
}
BencodeValue::Dict(dict)
}
}
}
}
fn parse_file_attr(value: &BencodeValue) -> Result<V2FileAttr, Error> {
let dict = value
.as_dict()
.ok_or_else(|| Error::InvalidTorrent("file attr must be a dict".into()))?;
let length = dict
.get(b"length".as_ref())
.and_then(irontide_bencode::BencodeValue::as_int)
.ok_or_else(|| Error::InvalidTorrent("file attr missing 'length'".into()))?;
if length < 0 {
return Err(Error::InvalidTorrent(format!(
"file attr has negative length: {length}"
)));
}
let pieces_root = if let Some(root_val) = dict.get(b"pieces root".as_ref()) {
let bytes = root_val
.as_bytes_raw()
.ok_or_else(|| Error::InvalidTorrent("pieces root must be bytes".into()))?;
Some(Id32::from_bytes(bytes)?)
} else {
None
};
Ok(V2FileAttr {
length: length as u64,
pieces_root,
})
}
#[cfg(test)]
mod tests {
use super::*;
fn bdict(pairs: Vec<(&[u8], BencodeValue)>) -> BencodeValue {
let mut map = BTreeMap::new();
for (k, v) in pairs {
map.insert(k.to_vec(), v);
}
BencodeValue::Dict(map)
}
fn bint(v: i64) -> BencodeValue {
BencodeValue::Integer(v)
}
fn bbytes(v: &[u8]) -> BencodeValue {
BencodeValue::Bytes(v.to_vec())
}
#[test]
fn single_file() {
let root_hash = [0xABu8; 32];
let tree = bdict(vec![(
b"test.txt",
bdict(vec![(
b"",
bdict(vec![
(b"length", bint(1024)),
(b"pieces root", bbytes(&root_hash)),
]),
)]),
)]);
let node = FileTreeNode::from_bencode(&tree).unwrap();
let files = node.flatten();
assert_eq!(files.len(), 1);
assert_eq!(files[0].path, vec!["test.txt"]);
assert_eq!(files[0].attr.length, 1024);
assert_eq!(files[0].attr.pieces_root, Some(Id32(root_hash)));
}
#[test]
fn nested_directory() {
let tree = bdict(vec![(
b"dir",
bdict(vec![(
b"subfile.dat",
bdict(vec![(b"", bdict(vec![(b"length", bint(512))]))]),
)]),
)]);
let node = FileTreeNode::from_bencode(&tree).unwrap();
let files = node.flatten();
assert_eq!(files.len(), 1);
assert_eq!(files[0].path, vec!["dir", "subfile.dat"]);
assert_eq!(files[0].attr.length, 512);
assert_eq!(files[0].attr.pieces_root, None);
}
#[test]
fn multiple_files_btreemap_ordering() {
let tree = bdict(vec![
(
b"beta.txt",
bdict(vec![(b"", bdict(vec![(b"length", bint(200))]))]),
),
(
b"alpha.txt",
bdict(vec![(b"", bdict(vec![(b"length", bint(100))]))]),
),
]);
let node = FileTreeNode::from_bencode(&tree).unwrap();
let files = node.flatten();
assert_eq!(files.len(), 2);
assert_eq!(files[0].path, vec!["alpha.txt"]);
assert_eq!(files[1].path, vec!["beta.txt"]);
}
#[test]
fn reject_missing_length() {
let tree = bdict(vec![(
b"bad.txt",
bdict(vec![(
b"",
bdict(vec![(b"pieces root", bbytes(&[0u8; 32]))]),
)]),
)]);
assert!(FileTreeNode::from_bencode(&tree).is_err());
}
#[test]
fn reject_non_dict() {
let value = BencodeValue::Integer(42);
assert!(FileTreeNode::from_bencode(&value).is_err());
}
#[test]
fn empty_file_no_pieces_root() {
let tree = bdict(vec![(
b"empty.txt",
bdict(vec![(b"", bdict(vec![(b"length", bint(0))]))]),
)]);
let node = FileTreeNode::from_bencode(&tree).unwrap();
let files = node.flatten();
assert_eq!(files.len(), 1);
assert_eq!(files[0].attr.length, 0);
assert_eq!(files[0].attr.pieces_root, None);
}
}