btpc_core/
edit.rs

1//! Validated owned metainfo editing.
2
3use std::collections::BTreeMap;
4
5use crate::bencode::{OwnedValue, Value, ValueKind, parse};
6use crate::{Error, Metainfo, Result};
7
8const RESERVED_TOP_LEVEL: &[&[u8]] = &[
9    b"announce",
10    b"announce-list",
11    b"url-list",
12    b"nodes",
13    b"comment",
14    b"created by",
15    b"creation date",
16    b"info",
17    b"piece layers",
18];
19
20/// Owned metainfo editor preserving untouched known and unknown fields.
21#[derive(Clone, Debug)]
22pub struct MetainfoEditor {
23    top_level: BTreeMap<Vec<u8>, OwnedValue>,
24    info: EditorInfo,
25}
26
27#[derive(Clone, Debug)]
28enum EditorInfo {
29    Raw(Vec<u8>),
30    Owned(BTreeMap<Vec<u8>, OwnedValue>),
31}
32
33impl MetainfoEditor {
34    /// Starts from a parsed validated metainfo object.
35    ///
36    /// # Errors
37    ///
38    /// Returns an error if parsing or owned top-level conversion fails.
39    pub fn from_metainfo(metainfo: &Metainfo) -> Result<Self> {
40        let parsed = parse(metainfo.original_bytes())?;
41        let ValueKind::Dictionary(entries) = parsed.kind() else {
42            unreachable!("validated metainfo roots are dictionaries");
43        };
44        let mut top_level = BTreeMap::new();
45        let mut info = None;
46        for (key, value) in entries {
47            if key.bytes() == b"info" {
48                info = Some(EditorInfo::Raw(
49                    metainfo.original_bytes()[value.span().range()].to_vec(),
50                ));
51            } else {
52                top_level.insert(key.bytes().to_vec(), owned(value)?);
53            }
54        }
55        Ok(Self {
56            top_level,
57            info: info.ok_or_else(|| Error::metainfo_field("info", "missing required field"))?,
58        })
59    }
60
61    /// Starts from an owned validated info dictionary.
62    ///
63    /// # Errors
64    ///
65    /// Returns an error when `info` is not a dictionary.
66    pub fn from_info(info: OwnedValue) -> Result<Self> {
67        let OwnedValue::Dictionary(info) = info else {
68            return Err(Error::metainfo_field("info", "must be a dictionary"));
69        };
70        Ok(Self {
71            top_level: BTreeMap::new(),
72            info: EditorInfo::Owned(info),
73        })
74    }
75
76    /// Replaces tracker tiers.
77    #[must_use]
78    pub fn trackers(mut self, tiers: impl IntoIterator<Item = Vec<Vec<u8>>>) -> Self {
79        let tiers = tiers.into_iter().collect::<Vec<_>>();
80        self.top_level.remove(b"announce".as_slice());
81        self.top_level.remove(b"announce-list".as_slice());
82        if let Some(announce) = tiers.first().and_then(|tier| tier.first()) {
83            self.top_level
84                .insert(b"announce".to_vec(), OwnedValue::bytes(announce.clone()));
85        }
86        if !tiers.is_empty() {
87            self.top_level.insert(
88                b"announce-list".to_vec(),
89                OwnedValue::list(
90                    tiers
91                        .into_iter()
92                        .map(|tier| OwnedValue::list(tier.into_iter().map(OwnedValue::bytes))),
93                ),
94            );
95        }
96        self
97    }
98
99    /// Replaces web seed URLs.
100    #[must_use]
101    pub fn web_seeds(mut self, seeds: impl IntoIterator<Item = Vec<u8>>) -> Self {
102        let seeds = seeds.into_iter().collect::<Vec<_>>();
103        set_optional(
104            &mut self.top_level,
105            b"url-list",
106            (!seeds.is_empty()).then(|| OwnedValue::list(seeds.into_iter().map(OwnedValue::bytes))),
107        );
108        self
109    }
110
111    /// Replaces DHT bootstrap nodes.
112    #[must_use]
113    pub fn nodes(mut self, nodes: impl IntoIterator<Item = (Vec<u8>, u16)>) -> Self {
114        let nodes = nodes.into_iter().collect::<Vec<_>>();
115        set_optional(
116            &mut self.top_level,
117            b"nodes",
118            (!nodes.is_empty()).then(|| {
119                OwnedValue::list(nodes.into_iter().map(|(host, port)| {
120                    OwnedValue::list([
121                        OwnedValue::bytes(host),
122                        OwnedValue::integer(i64::from(port)),
123                    ])
124                }))
125            }),
126        );
127        self
128    }
129
130    /// Sets or removes the top-level comment.
131    #[must_use]
132    pub fn comment(mut self, value: Option<Vec<u8>>) -> Self {
133        set_optional(
134            &mut self.top_level,
135            b"comment",
136            value.map(OwnedValue::bytes),
137        );
138        self
139    }
140
141    /// Sets or removes the top-level creator.
142    #[must_use]
143    pub fn created_by(mut self, value: Option<Vec<u8>>) -> Self {
144        set_optional(
145            &mut self.top_level,
146            b"created by",
147            value.map(OwnedValue::bytes),
148        );
149        self
150    }
151
152    /// Sets or removes the top-level creation timestamp.
153    #[must_use]
154    pub fn creation_date(mut self, value: Option<i64>) -> Self {
155        set_optional(
156            &mut self.top_level,
157            b"creation date",
158            value.map(OwnedValue::integer),
159        );
160        self
161    }
162
163    /// Sets or removes the private flag inside `info`.
164    #[must_use]
165    pub fn private(mut self, value: Option<bool>) -> Self {
166        if let Ok(info) = self.info_mut() {
167            set_optional(
168                info,
169                b"private",
170                value.map(|value| OwnedValue::integer(i64::from(value))),
171            );
172        }
173        self
174    }
175
176    /// Sets or removes the source field inside `info`.
177    #[must_use]
178    pub fn source(mut self, value: Option<Vec<u8>>) -> Self {
179        if let Ok(info) = self.info_mut() {
180            set_optional(info, b"source", value.map(OwnedValue::bytes));
181        }
182        self
183    }
184
185    /// Adds or replaces an unknown top-level extension.
186    ///
187    /// # Errors
188    ///
189    /// Returns an error for reserved keys.
190    pub fn raw_top_level(mut self, key: Vec<u8>, value: OwnedValue) -> Result<Self> {
191        if RESERVED_TOP_LEVEL.contains(&key.as_slice()) {
192            return Err(Error::metainfo_field(
193                "top-level field",
194                "reserved key must use its typed editor method",
195            ));
196        }
197        self.top_level.insert(key, value);
198        Ok(self)
199    }
200
201    /// Sets file attributes for a validated v1 or v2 path.
202    ///
203    /// # Errors
204    ///
205    /// Returns an error when the path is absent or has an unsupported shape.
206    pub fn file_attributes(mut self, path: &[Vec<u8>], attributes: Vec<u8>) -> Result<Self> {
207        let attributes = attributes.into_boxed_slice();
208        update_file_attributes(self.info_mut()?, path, &attributes, false)?;
209        Ok(self)
210    }
211
212    /// Serializes top-level fields canonically and validates the edited result.
213    ///
214    /// Untouched `info` bytes retain their exact source encoding. Any info-level
215    /// edit canonicalizes the updated `info` dictionary. Calling
216    /// [`Metainfo::to_bytes`] on the result explicitly canonicalizes the complete
217    /// metainfo object.
218    ///
219    /// # Errors
220    ///
221    /// Returns serialization or protocol validation errors.
222    pub fn to_metainfo(self) -> Result<Metainfo> {
223        let bytes = match self.info {
224            EditorInfo::Raw(info) => encode_with_raw_info(self.top_level, &info)?,
225            EditorInfo::Owned(info) => {
226                let mut root = self.top_level;
227                root.insert(b"info".to_vec(), OwnedValue::Dictionary(info));
228                OwnedValue::Dictionary(root).to_vec()?
229            }
230        };
231        Metainfo::from_bytes(&bytes)
232    }
233
234    fn info_mut(&mut self) -> Result<&mut BTreeMap<Vec<u8>, OwnedValue>> {
235        if let EditorInfo::Raw(bytes) = &self.info {
236            let parsed = parse(bytes)?;
237            let OwnedValue::Dictionary(info) = owned(&parsed)? else {
238                return Err(Error::metainfo_field("info", "must be a dictionary"));
239            };
240            self.info = EditorInfo::Owned(info);
241        }
242        let EditorInfo::Owned(info) = &mut self.info else {
243            unreachable!("raw info was converted to owned info");
244        };
245        Ok(info)
246    }
247}
248
249fn encode_with_raw_info(top_level: BTreeMap<Vec<u8>, OwnedValue>, info: &[u8]) -> Result<Vec<u8>> {
250    let mut output = vec![b'd'];
251    let mut wrote_info = false;
252    for (key, value) in top_level {
253        if !wrote_info && key.as_slice() > b"info".as_slice() {
254            write_bytes(&mut output, b"info");
255            output.extend_from_slice(info);
256            wrote_info = true;
257        }
258        write_bytes(&mut output, &key);
259        value.write_to(&mut output)?;
260    }
261    if !wrote_info {
262        write_bytes(&mut output, b"info");
263        output.extend_from_slice(info);
264    }
265    output.push(b'e');
266    Ok(output)
267}
268
269fn write_bytes(output: &mut Vec<u8>, bytes: &[u8]) {
270    output.extend_from_slice(bytes.len().to_string().as_bytes());
271    output.push(b':');
272    output.extend_from_slice(bytes);
273}
274
275fn update_v1_attributes(
276    info: &mut BTreeMap<Vec<u8>, OwnedValue>,
277    path: &[Vec<u8>],
278    attributes: &[u8],
279) -> Result<Option<bool>> {
280    if let Some(files) = info.get_mut(b"files".as_slice()) {
281        let OwnedValue::List(files) = files else {
282            return Err(Error::metainfo_field("info.files", "must be a list"));
283        };
284        for file in files {
285            let OwnedValue::Dictionary(fields) = file else {
286                return Err(Error::metainfo_field(
287                    "info.files",
288                    "entry must be a dictionary",
289                ));
290            };
291            if path_matches(fields.get(b"path".as_slice()), path) {
292                let padding = matches!(
293                    fields.get(b"attr".as_slice()),
294                    Some(OwnedValue::Bytes(current)) if current.contains(&b'p')
295                );
296                set_optional(
297                    fields,
298                    b"attr",
299                    (!attributes.is_empty()).then(|| OwnedValue::bytes(attributes.to_vec())),
300                );
301                return Ok(Some(padding));
302            }
303        }
304    } else if path.len() == 1
305        && info
306            .get(b"name".as_slice())
307            .is_some_and(|name| matches!(name, OwnedValue::Bytes(bytes) if bytes == &path[0]))
308        && info.contains_key(b"length".as_slice())
309    {
310        let padding = matches!(
311            info.get(b"attr".as_slice()),
312            Some(OwnedValue::Bytes(current)) if current.contains(&b'p')
313        );
314        set_optional(
315            info,
316            b"attr",
317            (!attributes.is_empty()).then(|| OwnedValue::bytes(attributes.to_vec())),
318        );
319        return Ok(Some(padding));
320    }
321    Ok(None)
322}
323
324fn update_v2_attributes(
325    info: &mut BTreeMap<Vec<u8>, OwnedValue>,
326    path: &[Vec<u8>],
327    attributes: &[u8],
328) -> Result<bool> {
329    let Some(mut node) = info.get_mut(b"file tree".as_slice()) else {
330        return Ok(false);
331    };
332    for component in path {
333        let OwnedValue::Dictionary(entries) = node else {
334            return Err(Error::metainfo_field(
335                "info.file tree",
336                "invalid path shape",
337            ));
338        };
339        let Some(next) = entries.get_mut(component.as_slice()) else {
340            return Ok(false);
341        };
342        node = next;
343    }
344    let OwnedValue::Dictionary(entries) = node else {
345        return Err(Error::metainfo_field("info.file tree", "invalid leaf"));
346    };
347    let Some(properties) = entries.get_mut(b"".as_slice()) else {
348        return Err(Error::metainfo_field(
349            "info.file tree",
350            "leaf marker missing",
351        ));
352    };
353    let OwnedValue::Dictionary(properties) = properties else {
354        return Err(Error::metainfo_field(
355            "info.file tree",
356            "invalid properties",
357        ));
358    };
359    set_optional(
360        properties,
361        b"attr",
362        (!attributes.is_empty()).then(|| OwnedValue::bytes(attributes.to_vec())),
363    );
364    Ok(true)
365}
366
367fn update_file_attributes(
368    info: &mut BTreeMap<Vec<u8>, OwnedValue>,
369    path: &[Vec<u8>],
370    attributes: &[u8],
371    inject_failure_after_v1: bool,
372) -> Result<()> {
373    let mut candidate = info.clone();
374    let has_v1 =
375        candidate.contains_key(b"files".as_slice()) || candidate.contains_key(b"length".as_slice());
376    let has_v2 = candidate.contains_key(b"file tree".as_slice());
377    let v1_padding = has_v1
378        .then(|| update_v1_attributes(&mut candidate, path, attributes))
379        .transpose()?
380        .flatten();
381    if inject_failure_after_v1 && v1_padding.is_some() {
382        return Err(Error::metainfo_field("file attributes", "injected failure"));
383    }
384    if v1_padding == Some(true) {
385        *info = candidate;
386        return Ok(());
387    }
388    let v2_updated = has_v2 && update_v2_attributes(&mut candidate, path, attributes)?;
389    match (has_v1, has_v2, v1_padding, v2_updated) {
390        (true, true, Some(false), true)
391        | (true, false, Some(false), false)
392        | (false, true, None, true) => {
393            *info = candidate;
394            Ok(())
395        }
396        (true, true, _, _) => Err(Error::metainfo_field(
397            "file path",
398            "hybrid v1 and v2 representations are inconsistent",
399        )),
400        _ => Err(Error::metainfo_field("file path", "path not found")),
401    }
402}
403
404fn set_optional(
405    dictionary: &mut BTreeMap<Vec<u8>, OwnedValue>,
406    key: &[u8],
407    value: Option<OwnedValue>,
408) {
409    if let Some(value) = value {
410        dictionary.insert(key.to_vec(), value);
411    } else {
412        dictionary.remove(key);
413    }
414}
415
416fn path_matches(value: Option<&OwnedValue>, expected: &[Vec<u8>]) -> bool {
417    let Some(OwnedValue::List(values)) = value else {
418        return false;
419    };
420    values.len() == expected.len()
421        && values.iter().zip(expected).all(
422            |(value, expected)| matches!(value, OwnedValue::Bytes(actual) if actual == expected),
423        )
424}
425
426fn owned(value: &Value<'_>) -> Result<OwnedValue> {
427    match value.kind() {
428        ValueKind::Integer(value) => OwnedValue::integer_bytes(value.canonical_bytes()),
429        ValueKind::Bytes(bytes) => Ok(OwnedValue::bytes(bytes.to_vec())),
430        ValueKind::List(values) => values
431            .iter()
432            .map(owned)
433            .collect::<Result<Vec<_>>>()
434            .map(OwnedValue::list),
435        ValueKind::Dictionary(entries) => OwnedValue::dictionary(
436            entries
437                .iter()
438                .map(|(key, value)| Ok((key.bytes().to_vec(), owned(value)?)))
439                .collect::<Result<Vec<_>>>()?,
440        ),
441    }
442}
443
444#[cfg(test)]
445mod tests {
446    use super::{OwnedValue, update_file_attributes};
447
448    #[test]
449    fn injected_hybrid_attribute_failure_leaves_info_unchanged() {
450        let leaf = OwnedValue::dictionary([(
451            Vec::new(),
452            OwnedValue::dictionary([(b"length".to_vec(), OwnedValue::integer(1))]).unwrap(),
453        )])
454        .unwrap();
455        let mut info = std::collections::BTreeMap::from([
456            (
457                b"files".to_vec(),
458                OwnedValue::list([OwnedValue::dictionary([
459                    (b"length".to_vec(), OwnedValue::integer(1)),
460                    (
461                        b"path".to_vec(),
462                        OwnedValue::list([OwnedValue::bytes(b"file".to_vec())]),
463                    ),
464                ])
465                .unwrap()]),
466            ),
467            (
468                b"file tree".to_vec(),
469                OwnedValue::dictionary([(b"file".to_vec(), leaf)]).unwrap(),
470            ),
471        ]);
472        let original = info.clone();
473        assert!(update_file_attributes(&mut info, &[b"file".to_vec()], b"x", true).is_err());
474        assert_eq!(info, original);
475    }
476}