Skip to main content

moq_transport/coding/
kvp.rs

1use crate::coding::{Decode, DecodeError, Encode, EncodeError};
2use std::fmt;
3
4#[derive(Clone, Eq, PartialEq)]
5pub enum Value {
6    IntValue(u64),
7    BytesValue(Vec<u8>),
8}
9
10impl fmt::Debug for Value {
11    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
12        match self {
13            Value::IntValue(v) => write!(f, "{}", v),
14            Value::BytesValue(bytes) => {
15                // Show up to 16 bytes in hex for readability
16                let preview: Vec<String> = bytes
17                    .iter()
18                    .take(16)
19                    .map(|b| format!("{:02X}", b))
20                    .collect();
21                write!(f, "[{}]", preview.join(" "))
22            }
23        }
24    }
25}
26
27#[derive(Clone, Eq, PartialEq)]
28pub struct KeyValuePair {
29    pub key: u64,
30    pub value: Value,
31}
32
33impl KeyValuePair {
34    pub fn new(key: u64, value: Value) -> Self {
35        Self { key, value }
36    }
37
38    pub fn new_int(key: u64, value: u64) -> Self {
39        Self {
40            key,
41            value: Value::IntValue(value),
42        }
43    }
44
45    pub fn new_bytes(key: u64, value: Vec<u8>) -> Self {
46        Self {
47            key,
48            value: Value::BytesValue(value),
49        }
50    }
51}
52
53impl Decode for KeyValuePair {
54    fn decode<R: bytes::Buf>(r: &mut R) -> Result<Self, DecodeError> {
55        let key = u64::decode(r)?;
56
57        if key % 2 == 0 {
58            // VarInt variant
59            let value = u64::decode(r)?;
60            tracing::trace!("[KVP] Decoded even key={}, value={}", key, value);
61            Ok(KeyValuePair::new_int(key, value))
62        } else {
63            // Bytes variant
64            let length = usize::decode(r)?;
65            tracing::trace!("[KVP] Decoded odd key={}, length={}", key, length);
66            if length > u16::MAX as usize {
67                tracing::error!(
68                    "[KVP] Length exceeded! key={}, length={} (max={})",
69                    key,
70                    length,
71                    u16::MAX
72                );
73                return Err(DecodeError::KeyValuePairLengthExceeded());
74            }
75
76            Self::decode_remaining(r, length)?;
77            let mut buf = vec![0; length];
78            r.copy_to_slice(&mut buf);
79            Ok(KeyValuePair::new_bytes(key, buf))
80        }
81    }
82}
83
84impl Encode for KeyValuePair {
85    fn encode<W: bytes::BufMut>(&self, w: &mut W) -> Result<(), EncodeError> {
86        match &self.value {
87            Value::IntValue(v) => {
88                // key must be even for IntValue
89                if !self.key.is_multiple_of(2) {
90                    return Err(EncodeError::InvalidValue);
91                }
92                self.key.encode(w)?;
93                (*v).encode(w)?;
94                Ok(())
95            }
96            Value::BytesValue(v) => {
97                // key must be odd for BytesValue
98                if self.key.is_multiple_of(2) {
99                    return Err(EncodeError::InvalidValue);
100                }
101                self.key.encode(w)?;
102                v.len().encode(w)?;
103                Self::encode_remaining(w, v.len())?;
104                w.put_slice(v);
105                Ok(())
106            }
107        }
108    }
109}
110
111impl fmt::Debug for KeyValuePair {
112    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
113        write!(f, "{{{}: {:?}}}", self.key, self.value)
114    }
115}
116
117/// A collection of KeyValuePair entries, where the number of key-value-pairs are encoded/decoded first.
118/// This structure is appropriate for Control message parameters.
119/// Since duplicate parameters are allowed for unknown parameters, we don't do duplicate checking here.
120#[derive(Default, Clone, Eq, PartialEq)]
121pub struct KeyValuePairs(pub Vec<KeyValuePair>);
122
123// TODO: These set/get API's all assume no duplicate keys. We can add API's to support duplicates if needed.
124impl KeyValuePairs {
125    pub fn new() -> Self {
126        Self::default()
127    }
128
129    /// Insert or replace a KeyValuePair with the same key.
130    pub fn set(&mut self, kvp: KeyValuePair) {
131        if let Some(existing) = self.0.iter_mut().find(|k| k.key == kvp.key) {
132            *existing = kvp;
133        } else {
134            self.0.push(kvp);
135        }
136    }
137
138    pub fn set_intvalue(&mut self, key: u64, value: u64) {
139        self.set(KeyValuePair::new_int(key, value));
140    }
141
142    pub fn set_bytesvalue(&mut self, key: u64, value: Vec<u8>) {
143        self.set(KeyValuePair::new_bytes(key, value));
144    }
145
146    pub fn has(&self, key: u64) -> bool {
147        self.0.iter().any(|k| k.key == key)
148    }
149
150    pub fn get(&self, key: u64) -> Option<&KeyValuePair> {
151        self.0.iter().find(|k| k.key == key)
152    }
153}
154
155impl Decode for KeyValuePairs {
156    fn decode<R: bytes::Buf>(mut r: &mut R) -> Result<Self, DecodeError> {
157        let mut kvps = Vec::new();
158
159        let count = u64::decode(r)?;
160        for _ in 0..count {
161            let kvp = KeyValuePair::decode(&mut r)?;
162            kvps.push(kvp);
163        }
164
165        Ok(KeyValuePairs(kvps))
166    }
167}
168
169impl Encode for KeyValuePairs {
170    fn encode<W: bytes::BufMut>(&self, w: &mut W) -> Result<(), EncodeError> {
171        self.0.len().encode(w)?;
172
173        for kvp in &self.0 {
174            kvp.encode(w)?;
175        }
176
177        Ok(())
178    }
179}
180
181impl fmt::Debug for KeyValuePairs {
182    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
183        write!(f, "{{ ")?;
184        for (i, kv) in self.0.iter().enumerate() {
185            if i > 0 {
186                write!(f, ", ")?;
187            }
188            write!(f, "{:?}", kv)?;
189        }
190        write!(f, " }}")
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use bytes::Bytes;
198    use bytes::BytesMut;
199
200    #[test]
201    fn encode_decode_keyvaluepair() {
202        let mut buf = BytesMut::new();
203
204        // Type=1, VarInt value=0 - illegal with odd key/type
205        let kvp = KeyValuePair::new(1, Value::IntValue(0));
206        let encoded = kvp.encode(&mut buf);
207        assert!(matches!(encoded.unwrap_err(), EncodeError::InvalidValue));
208
209        // Type=0, VarInt value=0
210        let kvp = KeyValuePair::new(0, Value::IntValue(0));
211        kvp.encode(&mut buf).unwrap();
212        assert_eq!(buf.to_vec(), vec![0x00, 0x00]);
213        let decoded = KeyValuePair::decode(&mut buf).unwrap();
214        assert_eq!(decoded, kvp);
215
216        // Type=100, VarInt value=100
217        let kvp = KeyValuePair::new(100, Value::IntValue(100));
218        kvp.encode(&mut buf).unwrap();
219        assert_eq!(buf.to_vec(), vec![0x40, 0x64, 0x40, 0x64]); // 2 2-byte VarInts with first 2 bits as 01
220        let decoded = KeyValuePair::decode(&mut buf).unwrap();
221        assert_eq!(decoded, kvp);
222
223        // Type=0, Bytes value=[1,2,3,4,5] - illegal with even key/type
224        let kvp = KeyValuePair::new(0, Value::BytesValue(vec![0x01, 0x02, 0x03, 0x04, 0x05]));
225        let decoded = kvp.encode(&mut buf);
226        assert!(matches!(decoded.unwrap_err(), EncodeError::InvalidValue));
227
228        // Type=1, Bytes value=[1,2,3,4,5]
229        let kvp = KeyValuePair::new(1, Value::BytesValue(vec![0x01, 0x02, 0x03, 0x04, 0x05]));
230        kvp.encode(&mut buf).unwrap();
231        assert_eq!(buf.to_vec(), vec![0x01, 0x05, 0x01, 0x02, 0x03, 0x04, 0x05]);
232        let decoded = KeyValuePair::decode(&mut buf).unwrap();
233        assert_eq!(decoded, kvp);
234    }
235
236    #[test]
237    fn decode_badtype() {
238        // Simulate a VarInt value of 5, but with an odd key/type
239        let data: Vec<u8> = vec![0x01, 0x05];
240        let mut buf: Bytes = data.into();
241        let decoded = KeyValuePair::decode(&mut buf);
242        assert!(matches!(decoded.unwrap_err(), DecodeError::More(_))); // Framing will be off now
243    }
244
245    #[test]
246    fn encode_decode_keyvaluepairs() {
247        let mut buf = BytesMut::new();
248
249        let mut kvps = KeyValuePairs::new();
250        kvps.set_bytesvalue(1, vec![0x01, 0x02, 0x03, 0x04, 0x05]);
251        kvps.encode(&mut buf).unwrap();
252        assert_eq!(
253            buf.to_vec(),
254            vec![
255                0x01, // 1 KeyValuePair
256                0x01, 0x05, 0x01, 0x02, 0x03, 0x04, 0x05, // Key=1, Value=[1,2,3,4,5]
257            ]
258        );
259        let decoded = KeyValuePairs::decode(&mut buf).unwrap();
260        assert_eq!(decoded, kvps);
261
262        let mut kvps = KeyValuePairs::new();
263        kvps.set_intvalue(0, 0);
264        kvps.set_intvalue(100, 100);
265        kvps.set_bytesvalue(1, vec![0x01, 0x02, 0x03, 0x04, 0x05]);
266        kvps.encode(&mut buf).unwrap();
267        let buf_vec = buf.to_vec();
268        //  Validate the encoded length and the KeyValuePair count
269        assert_eq!(14, buf_vec.len()); // 14 bytes total
270        assert_eq!(3, buf_vec[0]); // 3 KeyValuePairs
271        let decoded = KeyValuePairs::decode(&mut buf).unwrap();
272        assert_eq!(decoded, kvps);
273    }
274}