moq_transfork/message/
extensions.rs

1use std::collections::HashMap;
2use std::io::Cursor;
3
4use crate::coding::*;
5
6pub trait Extension: Encode + Decode {
7	fn id() -> u64;
8}
9
10#[derive(Default, Debug, Clone)]
11pub struct Extensions(HashMap<u64, Vec<u8>>);
12
13impl Decode for Extensions {
14	fn decode<R: bytes::Buf>(mut r: &mut R) -> Result<Self, DecodeError> {
15		let mut map = HashMap::new();
16
17		// I hate this encoding so much; let me encode my role and get on with my life.
18		let count = u64::decode(r)?;
19		for _ in 0..count {
20			let kind = u64::decode(r)?;
21			if map.contains_key(&kind) {
22				return Err(DecodeError::DupliateParameter);
23			}
24
25			let data = Vec::<u8>::decode(&mut r)?;
26			map.insert(kind, data);
27		}
28
29		Ok(Extensions(map))
30	}
31}
32
33impl Encode for Extensions {
34	fn encode<W: bytes::BufMut>(&self, w: &mut W) {
35		self.0.len().encode(w);
36
37		for (kind, value) in self.0.iter() {
38			kind.encode(w);
39			value.encode(w);
40		}
41	}
42}
43
44impl Extensions {
45	pub fn get<E: Extension>(&self) -> Result<Option<E>, DecodeError> {
46		Ok(match self.0.get(&E::id()) {
47			Some(payload) => {
48				let mut cursor = Cursor::new(payload);
49				Some(E::decode(&mut cursor)?)
50			}
51			None => None,
52		})
53	}
54
55	pub fn set<E: Extension>(&mut self, e: E) {
56		let mut value = Vec::new();
57		e.encode(&mut value);
58		self.0.insert(E::id(), value);
59	}
60}