1use 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#[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 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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 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 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}