stun_codec/
attribute.rs

1use crate::message::Message;
2use bytecodec::bytes::{BytesDecoder, BytesEncoder, RemainingBytesDecoder};
3use bytecodec::combinator::{Length, Peekable};
4use bytecodec::fixnum::{U16beDecoder, U16beEncoder};
5use bytecodec::{ByteCount, Decode, Encode, Eos, ErrorKind, Result, SizedEncode, TryTaggedDecode};
6use std::fmt;
7
8use crate::{rfc5389, rfc5766};
9
10#[derive(Debug)]
11pub enum MyAttribute {
12    Rfc5389(rfc5389::Attribute),
13    Rfc5766(rfc5766::Attribute),
14}
15
16#[derive(Debug, Default)]
17pub struct MyAttributeDecoder {
18    rfc5389: rfc5389::AttributeDecoder,
19    rfc5766: rfc5766::AttributeDecoder,
20    index: usize,
21}
22impl Decode for MyAttributeDecoder {
23    type Item = MyAttribute;
24
25    fn decode(&mut self, buf: &[u8], eos: Eos) -> Result<usize> {
26        match self.index {
27            1 => track!(self.rfc5389.decode(buf, eos)),
28            2 => track!(self.rfc5766.decode(buf, eos)),
29            _ => track_panic!(ErrorKind::InconsistentState),
30        }
31    }
32
33    fn finish_decoding(&mut self) -> Result<Self::Item> {
34        let item = match self.index {
35            1 => track!(self.rfc5389.finish_decoding()).map(MyAttribute::Rfc5389)?,
36            2 => track!(self.rfc5766.finish_decoding()).map(MyAttribute::Rfc5766)?,
37            _ => track_panic!(ErrorKind::InconsistentState),
38        };
39        self.index = 0;
40        Ok(item)
41    }
42
43    fn requiring_bytes(&self) -> ByteCount {
44        match self.index {
45            1 => self.rfc5389.requiring_bytes(),
46            2 => self.rfc5766.requiring_bytes(),
47            _ => ByteCount::Finite(0),
48        }
49    }
50
51    fn is_idle(&self) -> bool {
52        match self.index {
53            1 => self.rfc5389.is_idle(),
54            2 => self.rfc5766.is_idle(),
55            _ => true,
56        }
57    }
58}
59impl TryTaggedDecode for MyAttributeDecoder {
60    type Tag = AttributeType;
61
62    fn try_start_decoding(&mut self, tag: Self::Tag) -> Result<bool> {
63        track_assert_eq!(self.index, 0, ErrorKind::InconsistentState);
64        if track!(self.rfc5389.try_start_decoding(tag))? {
65            self.index = 1;
66            Ok(true)
67        } else if track!(self.rfc5766.try_start_decoding(tag))? {
68            self.index = 2;
69            Ok(true)
70        } else {
71            Ok(false)
72        }
73    }
74}
75
76#[derive(Debug, Default)]
77pub struct MyAttributeEncoder {
78    rfc5389: rfc5389::AttributeEncoder,
79    rfc5766: rfc5766::AttributeEncoder,
80}
81impl Encode for MyAttributeEncoder {
82    type Item = MyAttribute;
83
84    fn encode(&mut self, buf: &mut [u8], eos: Eos) -> Result<usize> {
85        let mut offset = 0;
86        bytecodec_try_encode!(self.rfc5389, offset, buf, eos);
87        bytecodec_try_encode!(self.rfc5766, offset, buf, eos);
88        Ok(offset)
89    }
90
91    fn start_encoding(&mut self, item: Self::Item) -> Result<()> {
92        track_assert!(self.is_idle(), ErrorKind::EncoderFull);
93        match item {
94            MyAttribute::Rfc5389(item) => track!(self.rfc5389.start_encoding(item)),
95            MyAttribute::Rfc5766(item) => track!(self.rfc5766.start_encoding(item)),
96        }
97    }
98
99    fn requiring_bytes(&self) -> ByteCount {
100        self.rfc5389
101            .requiring_bytes()
102            .add_for_encoding(self.rfc5766.requiring_bytes())
103    }
104
105    fn is_idle(&self) -> bool {
106        self.rfc5389.is_idle() && self.rfc5766.is_idle()
107    }
108}
109impl SizedEncode for MyAttributeEncoder {
110    fn exact_requiring_bytes(&self) -> u64 {
111        self.rfc5389.exact_requiring_bytes() + self.rfc5766.exact_requiring_bytes()
112    }
113}
114
115/// STUN attribute.
116///
117/// > **Attribute**:  The STUN term for a Type-Length-Value (TLV) object that
118/// > can be added to a STUN message. Attributes are divided into two
119/// > types: comprehension-required and comprehension-optional. STUN
120/// > agents can safely ignore comprehension-optional attributes they
121/// > don't understand, but cannot successfully process a message if it
122/// > contains comprehension-required attributes that are not
123/// > understood.
124/// >
125/// > [RFC 5389 -- 5. Definitions]
126///
127/// [RFC 5389 -- 5. Definitions]: https://tools.ietf.org/html/rfc5389#section-5
128pub trait Attribute: Sized + Clone {
129    /// The decoder of the value part of the attribute.
130    type Decoder: Default + TryTaggedDecode<Tag = AttributeType, Item = Self>;
131
132    /// The encoder of the value part of the attribute.
133    type Encoder: Default + SizedEncode<Item = Self>;
134
135    /// Returns the type of the attribute.
136    fn get_type(&self) -> AttributeType;
137
138    /// This method is called before encoding the attribute.
139    ///
140    /// `message` is the message to which the attribute belongs.
141    /// The message only contains the attributes preceding to `self`.
142    ///
143    /// The default implementation simply returns `Ok(())`.
144    #[allow(unused_variables)]
145    fn before_encode<A: Attribute>(&mut self, message: &Message<A>) -> Result<()> {
146        Ok(())
147    }
148
149    /// This method is called after decoding the attribute and before being appended to the given message.
150    ///
151    /// The default implementation simply returns `Ok(())`.
152    #[allow(unused_variables)]
153    fn after_decode<A: Attribute>(&mut self, message: &Message<A>) -> Result<()> {
154        Ok(())
155    }
156}
157
158/// Attribute type.
159///
160/// > Attributes are divided into two
161/// > types: comprehension-required and comprehension-optional. STUN
162/// > agents can safely ignore comprehension-optional attributes they
163/// > don't understand, but cannot successfully process a message if it
164/// > contains comprehension-required attributes that are not
165/// > understood.
166/// >
167/// > [RFC 5389 -- 5. Definitions]
168/// >
169/// > ---
170/// >
171/// > A STUN Attribute type is a hex number in the range 0x0000 - 0xFFFF.
172/// > STUN attribute types in the range 0x0000 - 0x7FFF are considered
173/// > comprehension-required; STUN attribute types in the range 0x8000 -
174/// > 0xFFFF are considered comprehension-optional.
175/// >
176/// > [RFC 5389 -- 18.2. STUN Attribute Registry]
177///
178/// [RFC 5389 -- 5. Definitions]: https://tools.ietf.org/html/rfc5389#section-5
179/// [RFC 5389 -- 18.2. STUN Attribute Registry]: https://tools.ietf.org/html/rfc5389#section-18.2
180#[derive(Debug, Clone, Copy, PartialOrd, Ord, PartialEq, Eq, Hash)]
181pub struct AttributeType(u16);
182impl AttributeType {
183    /// Makes a new `Type` instance which corresponding to `codepoint`.
184    pub fn new(codepoint: u16) -> Self {
185        AttributeType(codepoint)
186    }
187
188    /// Returns the attribute codepoint corresponding this instance.
189    pub fn as_u16(self) -> u16 {
190        self.0
191    }
192
193    /// Returns `true` if this is a comprehension-required type.
194    pub fn is_comprehension_required(self) -> bool {
195        self.0 < 0x8000
196    }
197
198    /// Returns `true` if this is a comprehension-optional type.
199    pub fn is_comprehension_optional(self) -> bool {
200        !self.is_comprehension_required()
201    }
202}
203impl From<u16> for AttributeType {
204    fn from(f: u16) -> Self {
205        Self::new(f)
206    }
207}
208
209/// An [`Attribute`] implementation that has raw value bytes.
210#[derive(Debug, Clone)]
211pub struct RawAttribute {
212    attr_type: AttributeType,
213    value: Vec<u8>,
214}
215impl RawAttribute {
216    /// Makes a new `RawAttribute` instance.
217    pub fn new(attr_type: AttributeType, value: Vec<u8>) -> Self {
218        RawAttribute { attr_type, value }
219    }
220
221    /// Returns a reference to the value bytes of the attribute.
222    pub fn value(&self) -> &[u8] {
223        &self.value
224    }
225
226    /// Takes ownership of this instance, and returns the value bytes.
227    pub fn into_value(self) -> Vec<u8> {
228        self.value
229    }
230}
231impl Attribute for RawAttribute {
232    type Decoder = RawAttributeDecoder;
233    type Encoder = RawAttributeEncoder;
234
235    fn get_type(&self) -> AttributeType {
236        self.attr_type
237    }
238}
239
240/// [`RawAttribute`] decoder.
241#[derive(Debug, Default)]
242pub struct RawAttributeDecoder {
243    attr_type: Option<AttributeType>,
244    value: RemainingBytesDecoder,
245}
246impl RawAttributeDecoder {
247    /// Makes a new `RawAttributeDecoder` instance.
248    pub fn new() -> Self {
249        Self::default()
250    }
251}
252impl Decode for RawAttributeDecoder {
253    type Item = RawAttribute;
254
255    fn decode(&mut self, buf: &[u8], eos: Eos) -> Result<usize> {
256        track!(self.value.decode(buf, eos))
257    }
258
259    fn finish_decoding(&mut self) -> Result<Self::Item> {
260        let attr_type = track_assert_some!(self.attr_type.take(), ErrorKind::InconsistentState);
261        let value = track!(self.value.finish_decoding())?;
262        Ok(RawAttribute { attr_type, value })
263    }
264
265    fn requiring_bytes(&self) -> ByteCount {
266        self.value.requiring_bytes()
267    }
268
269    fn is_idle(&self) -> bool {
270        self.value.is_idle()
271    }
272}
273impl TryTaggedDecode for RawAttributeDecoder {
274    type Tag = AttributeType;
275
276    fn try_start_decoding(&mut self, attr_type: Self::Tag) -> Result<bool> {
277        self.attr_type = Some(attr_type);
278        Ok(true)
279    }
280}
281
282/// [`RawAttribute`] encoder.
283#[derive(Debug, Default)]
284pub struct RawAttributeEncoder {
285    value: BytesEncoder,
286}
287impl RawAttributeEncoder {
288    /// Makes a new `RawAttributeEncoder` instance.
289    pub fn new() -> Self {
290        Self::default()
291    }
292}
293impl Encode for RawAttributeEncoder {
294    type Item = RawAttribute;
295
296    fn encode(&mut self, buf: &mut [u8], eos: Eos) -> Result<usize> {
297        track!(self.value.encode(buf, eos))
298    }
299
300    fn start_encoding(&mut self, item: Self::Item) -> Result<()> {
301        track!(self.value.start_encoding(item.into_value()))
302    }
303
304    fn requiring_bytes(&self) -> ByteCount {
305        ByteCount::Finite(self.exact_requiring_bytes())
306    }
307
308    fn is_idle(&self) -> bool {
309        self.value.is_idle()
310    }
311}
312impl SizedEncode for RawAttributeEncoder {
313    fn exact_requiring_bytes(&self) -> u64 {
314        self.value.exact_requiring_bytes()
315    }
316}
317
318#[derive(Debug, Clone)]
319pub enum LosslessAttribute<T> {
320    Known {
321        inner: T,
322        padding: Option<Padding>,
323    },
324    Unknown {
325        inner: RawAttribute,
326        padding: Option<Padding>,
327    },
328}
329impl<T: Attribute> LosslessAttribute<T> {
330    pub fn new(inner: T) -> Self {
331        LosslessAttribute::Known {
332            inner,
333            padding: None,
334        }
335    }
336
337    pub fn as_known(&self) -> Option<&T> {
338        match self {
339            LosslessAttribute::Known { inner, .. } => Some(inner),
340            LosslessAttribute::Unknown { .. } => None,
341        }
342    }
343
344    pub fn as_unknown(&self) -> Option<&RawAttribute> {
345        match self {
346            LosslessAttribute::Known { .. } => None,
347            LosslessAttribute::Unknown { inner, .. } => Some(inner),
348        }
349    }
350
351    pub fn get_type(&self) -> AttributeType {
352        match self {
353            LosslessAttribute::Known { inner, .. } => inner.get_type(),
354            LosslessAttribute::Unknown { inner, .. } => inner.get_type(),
355        }
356    }
357
358    pub fn before_encode<A: Attribute>(&mut self, message: &Message<A>) -> Result<()> {
359        match self {
360            LosslessAttribute::Known { inner, .. } => inner.before_encode(message),
361            LosslessAttribute::Unknown { inner, .. } => inner.before_encode(message),
362        }
363    }
364
365    pub fn after_decode<A: Attribute>(&mut self, message: &Message<A>) -> Result<()> {
366        match self {
367            LosslessAttribute::Known { inner, .. } => inner.after_decode(message),
368            LosslessAttribute::Unknown { inner, .. } => inner.after_decode(message),
369        }
370    }
371}
372
373pub struct LosslessAttributeDecoder<T: Attribute> {
374    get_type: U16beDecoder,
375    value_len: Peekable<U16beDecoder>,
376    is_known: bool,
377    known_value: Length<T::Decoder>,
378    unknown_value: Length<RawAttributeDecoder>,
379    padding: BytesDecoder<Padding>,
380}
381impl<T: Attribute> fmt::Debug for LosslessAttributeDecoder<T> {
382    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
383        write!(f, "LosslessAttributeDecoder {{ .. }}")
384    }
385}
386impl<T: Attribute> Default for LosslessAttributeDecoder<T> {
387    fn default() -> Self {
388        LosslessAttributeDecoder {
389            get_type: Default::default(),
390            value_len: Default::default(),
391            is_known: false,
392            known_value: Default::default(),
393            unknown_value: Default::default(),
394            padding: Default::default(),
395        }
396    }
397}
398impl<T: Attribute> Decode for LosslessAttributeDecoder<T> {
399    type Item = LosslessAttribute<T>;
400
401    fn decode(&mut self, buf: &[u8], eos: Eos) -> Result<usize> {
402        let mut offset = 0;
403        if !self.value_len.is_idle() {
404            bytecodec_try_decode!(self.get_type, offset, buf, eos);
405            bytecodec_try_decode!(self.value_len, offset, buf, eos);
406
407            let attr_type = AttributeType(track!(self.get_type.finish_decoding())?);
408            let value_len = *self.value_len.peek().expect("never fails");
409
410            self.is_known = track!(self.known_value.inner_mut().try_start_decoding(attr_type))?;
411            if self.is_known {
412                track!(self.known_value.set_expected_bytes(u64::from(value_len)))?;
413            } else {
414                track!(self.unknown_value.inner_mut().try_start_decoding(attr_type))?; // must be `true`
415                track!(self.unknown_value.set_expected_bytes(u64::from(value_len)))?;
416            }
417            self.padding.set_bytes(Padding::new(value_len as usize));
418        }
419        if self.is_known {
420            bytecodec_try_decode!(self.known_value, offset, buf, eos);
421        } else {
422            bytecodec_try_decode!(self.unknown_value, offset, buf, eos);
423        }
424        bytecodec_try_decode!(self.padding, offset, buf, eos);
425        Ok(offset)
426    }
427
428    fn finish_decoding(&mut self) -> Result<Self::Item> {
429        let _ = track!(self.value_len.finish_decoding())?;
430        let padding = track!(self.padding.finish_decoding())?;
431        if self.is_known {
432            let value = track!(self.known_value.finish_decoding())?;
433            Ok(LosslessAttribute::Known {
434                inner: value,
435                padding: Some(padding),
436            })
437        } else {
438            let value = track!(self.unknown_value.finish_decoding())?;
439            Ok(LosslessAttribute::Unknown {
440                inner: value,
441                padding: Some(padding),
442            })
443        }
444    }
445
446    fn requiring_bytes(&self) -> ByteCount {
447        if self.value_len.is_idle() {
448            if self.is_known {
449                self.known_value
450                    .requiring_bytes()
451                    .add_for_decoding(self.padding.requiring_bytes())
452            } else {
453                self.unknown_value
454                    .requiring_bytes()
455                    .add_for_decoding(self.padding.requiring_bytes())
456            }
457        } else {
458            self.get_type
459                .requiring_bytes()
460                .add_for_decoding(self.value_len.requiring_bytes())
461        }
462    }
463
464    fn is_idle(&self) -> bool {
465        self.value_len.is_idle()
466            && if self.is_known {
467                self.known_value.is_idle()
468            } else {
469                self.unknown_value.is_idle()
470            }
471            && self.padding.is_idle()
472    }
473}
474
475pub struct LosslessAttributeEncoder<T: Attribute> {
476    get_type: U16beEncoder,
477    value_len: U16beEncoder,
478    known_value: T::Encoder,
479    unknown_value: RawAttributeEncoder,
480    padding: BytesEncoder<Padding>,
481}
482impl<T: Attribute> fmt::Debug for LosslessAttributeEncoder<T> {
483    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
484        write!(f, "LosslessAttributeEncoder {{ .. }}")
485    }
486}
487impl<T: Attribute> Default for LosslessAttributeEncoder<T> {
488    fn default() -> Self {
489        LosslessAttributeEncoder {
490            get_type: Default::default(),
491            value_len: Default::default(),
492            known_value: Default::default(),
493            unknown_value: Default::default(),
494            padding: Default::default(),
495        }
496    }
497}
498impl<T: Attribute> Encode for LosslessAttributeEncoder<T> {
499    type Item = LosslessAttribute<T>;
500
501    fn encode(&mut self, buf: &mut [u8], eos: Eos) -> Result<usize> {
502        let mut offset = 0;
503        bytecodec_try_encode!(self.get_type, offset, buf, eos);
504        bytecodec_try_encode!(self.value_len, offset, buf, eos);
505        bytecodec_try_encode!(self.known_value, offset, buf, eos);
506        bytecodec_try_encode!(self.unknown_value, offset, buf, eos);
507        bytecodec_try_encode!(self.padding, offset, buf, eos);
508        Ok(offset)
509    }
510
511    fn start_encoding(&mut self, item: Self::Item) -> Result<()> {
512        track!(self.get_type.start_encoding(item.get_type().as_u16()))?;
513        let padding = match item {
514            LosslessAttribute::Known { inner, padding } => {
515                track!(self.known_value.start_encoding(inner))?;
516                padding
517            }
518            LosslessAttribute::Unknown { inner, padding } => {
519                track!(self.unknown_value.start_encoding(inner))?;
520                padding
521            }
522        };
523
524        let value_len =
525            self.known_value.exact_requiring_bytes() + self.unknown_value.exact_requiring_bytes();
526        track_assert!(value_len < 0x10000, ErrorKind::InvalidInput; value_len);
527
528        let padding = padding.unwrap_or_else(|| Padding::new(value_len as usize));
529        track!(self.value_len.start_encoding(value_len as u16))?;
530        track!(self.padding.start_encoding(padding))?;
531        Ok(())
532    }
533
534    fn requiring_bytes(&self) -> ByteCount {
535        ByteCount::Finite(self.exact_requiring_bytes())
536    }
537
538    fn is_idle(&self) -> bool {
539        self.value_len.is_idle()
540            && self.known_value.is_idle()
541            && self.unknown_value.is_idle()
542            && self.padding.is_idle()
543    }
544}
545impl<T: Attribute> SizedEncode for LosslessAttributeEncoder<T> {
546    fn exact_requiring_bytes(&self) -> u64 {
547        self.get_type.exact_requiring_bytes()
548            + self.value_len.exact_requiring_bytes()
549            + self.known_value.exact_requiring_bytes()
550            + self.unknown_value.exact_requiring_bytes()
551            + self.padding.exact_requiring_bytes()
552    }
553}
554
555#[derive(Default, Clone)]
556pub struct Padding {
557    buf: [u8; 3],
558    len: usize,
559}
560impl Padding {
561    fn new(value_len: usize) -> Self {
562        let len = (4 - value_len % 4) % 4;
563        Padding { buf: [0; 3], len }
564    }
565}
566impl fmt::Debug for Padding {
567    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
568        write!(f, "Padding({:?})", self.as_ref())
569    }
570}
571impl AsRef<[u8]> for Padding {
572    fn as_ref(&self) -> &[u8] {
573        &self.buf[..self.len]
574    }
575}
576impl AsMut<[u8]> for Padding {
577    fn as_mut(&mut self) -> &mut [u8] {
578        &mut self.buf[..self.len]
579    }
580}
stun_codec/attribute.rs

stun_codec/
attribute.rs
