stun_rs/
message.rs

1use crate::attributes::{StunAttribute, StunAttributeType};
2use crate::common::check_buffer_boundaries;
3use crate::error::{StunError, StunErrorType};
4use crate::{Encode, TransactionId};
5use byteorder::{BigEndian, ByteOrder};
6use std::convert::{TryFrom, TryInto};
7
8/// The message type defines the message class (request, success
9/// response, error response, or indication) and the message method (the
10/// primary function) of the STUN message.  Although there are four
11/// message classes, there are only two types of transactions in STUN:
12/// request/response transactions (which consist of a request message and
13/// a response message) and indication transactions (which consist of a
14/// single indication message).  Response classes are split into error
15/// and success responses to aid in quickly processing the STUN message.
16/// # Examples
17///```rust
18/// # use stun_rs::{MessageClass, MessageMethod, MessageType};
19/// # use stun_rs::methods::BINDING;
20/// # use std::convert::TryFrom;
21/// # use std::error::Error;
22/// #
23/// # fn main() -> Result<(), Box<dyn Error>> {
24/// let msg_type = MessageType::new(BINDING, MessageClass::SuccessResponse);
25/// assert_eq!(msg_type.as_u16(), 0x0101);
26/// assert_eq!(msg_type.method(), BINDING);
27/// assert_eq!(msg_type.class(), MessageClass::SuccessResponse);
28/// #
29/// #  Ok(())
30/// # }
31///```
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct MessageType {
34    method: MessageMethod,
35    class: MessageClass,
36}
37
38impl MessageType {
39    /// Creates a new message type.
40    /// # Arguments:
41    /// - `method`- the message method.
42    /// - `class` - The message class.
43    pub fn new(method: MessageMethod, class: MessageClass) -> Self {
44        Self { method, class }
45    }
46
47    /// Returns the message class.
48    pub fn class(&self) -> MessageClass {
49        self.class
50    }
51
52    /// Returns the message method
53    pub fn method(&self) -> MessageMethod {
54        self.method
55    }
56
57    /// Returns the [`u16`] representation of this [`MessageType`]
58    pub fn as_u16(&self) -> u16 {
59        ((self.method.0 & 0x1F80) << 2)
60            | ((self.method.as_u16() & 0x0070) << 1)
61            | (self.method.as_u16() & 0x000F)
62            | ((self.class.as_u16() & 0x0002) << 7)
63            | ((self.class.as_u16() & 0x0001) << 4)
64    }
65}
66
67impl From<u16> for MessageType {
68    fn from(value: u16) -> Self {
69        // Discard two most significant bits
70        let val = value & 0x3FFF;
71        // There is not way this can fail. Value gotten will always fit into a u8
72        // and it will be less or equal to 0x0003
73        let class_u8: u8 = (((val & 0x0100) >> 7) | ((val & 0x0010) >> 4))
74            .try_into()
75            .unwrap();
76        let class = MessageClass::try_from(class_u8).unwrap();
77        // There is not way that method number falls out of the range defined 0x000-0xFFF
78        let method_u16: u16 = ((val & 0x3E00) >> 2) | ((val & 0x00E0) >> 1) | (val & 0x000F);
79        let method = MessageMethod::try_from(method_u16).unwrap();
80
81        MessageType::new(method, class)
82    }
83}
84
85impl From<&[u8; 2]> for MessageType {
86    fn from(value: &[u8; 2]) -> Self {
87        MessageType::from(BigEndian::read_u16(value))
88    }
89}
90
91impl Encode for MessageType {
92    fn encode(&self, buffer: &mut [u8]) -> Result<usize, StunError> {
93        check_buffer_boundaries(buffer, 2)?;
94        BigEndian::write_u16(buffer, self.as_u16());
95        Ok(2)
96    }
97}
98
99/// The STUN method is a 12 bits hex number in the range 0x000-0xFFF but
100/// valid values are defined in the range 0x00-0xFF.
101/// STUN methods in the range 0x000-0x07F are assigned by `IETF` Review
102/// [`RFC8126`](https://datatracker.ietf.org/doc/html/rfc8126). STUN
103/// methods in the range 0x080-0x0FF are assigned by Expert Review.
104///
105/// # Examples
106///```rust
107/// # use stun_rs::{MessageMethod, StunErrorType};
108/// # use std::convert::TryFrom;
109/// # use std::error::Error;
110/// #
111/// # fn main() -> Result<(), Box<dyn Error>> {
112/// // Create a binding method
113/// let binding = MessageMethod::try_from(0x001)?;
114/// assert_eq!(binding.as_u16(), 0x001);
115/// // Binding request is within the range of valid values 0x00-0xFF
116/// assert!(binding.is_valid());
117///
118/// // Create a custom method
119/// let method = MessageMethod::try_from(0x100)?;
120/// // This method is out of the range of valid values 0x00-0xFF
121/// assert!(!method.is_valid());
122///
123/// // Creating a message method out of 12 bits range 0x000-0xFFF
124/// // will result in an error
125/// assert_eq!(MessageMethod::try_from(0x1000).expect_err("Error expected"), StunErrorType::InvalidParam);
126/// #
127/// #   Ok(())
128/// # }
129///```
130#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
131pub struct MessageMethod(pub(crate) u16);
132
133impl MessageMethod {
134    /// Returns the [`u16`] representation of this message method.
135    pub fn as_u16(&self) -> u16 {
136        self.0
137    }
138
139    /// Returns true if the method is within the valid range 0x00-0xFF
140    pub fn is_valid(&self) -> bool {
141        (0x00..=0xff).contains(&self.0)
142    }
143}
144
145impl TryFrom<u16> for MessageMethod {
146    type Error = StunError;
147
148    fn try_from(value: u16) -> Result<Self, Self::Error> {
149        (value & 0xF000 == 0)
150            .then_some(MessageMethod(value))
151            .ok_or_else(|| {
152                StunError::new(
153                    StunErrorType::InvalidParam,
154                    format!("Value '{:#02x}' is not a valid a MessageMethod", value),
155                )
156            })
157    }
158}
159
160/// The STUN message class. Although there are four
161/// message classes, there are only two types of transactions in STUN:
162/// request/response transactions (which consist of a request message and
163/// a response message) and indication transactions (which consist of a
164/// single indication message).  Response classes are split into error
165/// and success responses to aid in quickly processing the STUN message.
166#[derive(Debug, Clone, Copy, PartialEq, Eq)]
167pub enum MessageClass {
168    /// request
169    Request,
170    /// indication
171    Indication,
172    /// success response
173    SuccessResponse,
174    /// error response
175    ErrorResponse,
176}
177
178impl MessageClass {
179    fn as_u16(&self) -> u16 {
180        match self {
181            MessageClass::Request => 0b00,
182            MessageClass::Indication => 0b01,
183            MessageClass::SuccessResponse => 0b10,
184            MessageClass::ErrorResponse => 0b11,
185        }
186    }
187}
188
189impl TryFrom<u8> for MessageClass {
190    type Error = StunError;
191
192    fn try_from(value: u8) -> Result<Self, Self::Error> {
193        match value {
194            0b00 => Ok(MessageClass::Request),
195            0b01 => Ok(MessageClass::Indication),
196            0b10 => Ok(MessageClass::SuccessResponse),
197            0b11 => Ok(MessageClass::ErrorResponse),
198            _ => Err(StunError::new(
199                StunErrorType::InvalidParam,
200                format!("Value '{:#02x}' is not a valid a MessageClass", value),
201            )),
202        }
203    }
204}
205
206#[derive(Debug)]
207struct StunMessageParameters {
208    method: MessageMethod,
209    class: MessageClass,
210    transaction_id: Option<TransactionId>,
211    attributes: Vec<StunAttribute>,
212}
213
214/// The [`StunMessageBuilder`] ease the creation of a [`StunMessage`]
215///
216/// # Examples
217///```rust
218/// # use stun_rs::{MessageClass, MessageMethod, StunAttribute, StunMessage, StunMessageBuilder};
219/// # use stun_rs::attributes::stun::{Software, UserName, Nonce};
220/// # use stun_rs::methods::BINDING;
221/// # use std::convert::TryFrom;
222/// # use std::error::Error;
223/// #
224/// # fn main() -> Result<(), Box<dyn Error>> {
225/// // Create a SUN request message with a random transaction ID.
226/// let message = StunMessageBuilder::new(
227///     BINDING,
228///     MessageClass::Request,
229/// )
230/// .with_attribute(UserName::try_from("test-username")?)
231/// .with_attribute(Software::new("test-software")?)
232/// .build();
233///
234/// let username = message.get::<UserName>()
235///   .ok_or("UserName attriute not found")?
236///   .as_user_name()?;
237/// assert_eq!(username, "test-username");
238///
239/// let software = message.get::<Software>()
240///   .ok_or("Software attriute not found")?
241///   .as_software()?;
242/// assert_eq!(software, "test-software");
243///
244/// // Nonce attribute must return None
245/// assert!(message.get::<Nonce>().is_none());
246/// #
247/// #   Ok(())
248/// # }
249///```
250#[derive(Debug)]
251pub struct StunMessageBuilder(StunMessageParameters);
252
253impl StunMessageBuilder {
254    /// Creates a new builder.
255    /// # Arguments:
256    /// - `method` - Message method.
257    /// - `class` - Message class.
258    pub fn new(method: MessageMethod, class: MessageClass) -> StunMessageBuilder {
259        Self(StunMessageParameters {
260            method,
261            class,
262            transaction_id: None,
263            attributes: Vec::new(),
264        })
265    }
266
267    /// Creates a STUN message using an specific transaction ID. If no
268    /// [`TransactionId`] is specified, a random one will be used
269    pub fn with_transaction_id(mut self, transaction_id: TransactionId) -> Self {
270        self.0.transaction_id = Some(transaction_id);
271        self
272    }
273
274    /// Adds an attribute to the message.
275    pub fn with_attribute<T>(mut self, attribute: T) -> Self
276    where
277        T: Into<StunAttribute>,
278    {
279        self.0.attributes.push(attribute.into());
280        self
281    }
282
283    /// Creates the STUN message.
284    pub fn build(self) -> StunMessage {
285        StunMessage {
286            method: self.0.method,
287            class: self.0.class,
288            transaction_id: self.0.transaction_id.unwrap_or_default(),
289            attributes: self.0.attributes,
290        }
291    }
292}
293
294/// The stun message is the basic unit of information interchanged between
295/// two agents implementing the STUN protocol.
296///
297/// All STUN messages comprise a 20-byte header followed by zero or more
298/// attributes. The STUN header contains a STUN message type, message
299/// length, magic cookie, and transaction ID.
300///
301/// STUN messages can be created using the [`StunMessageBuilder`].
302#[derive(Debug)]
303pub struct StunMessage {
304    method: MessageMethod,
305    class: MessageClass,
306    transaction_id: TransactionId,
307    attributes: Vec<StunAttribute>,
308}
309
310impl StunMessage {
311    /// Returns the message method.
312    pub fn method(&self) -> MessageMethod {
313        self.method
314    }
315
316    /// Returns the message class
317    pub fn class(&self) -> MessageClass {
318        self.class
319    }
320
321    /// Returns the transaction-id
322    pub fn transaction_id(&self) -> &TransactionId {
323        &self.transaction_id
324    }
325
326    /// Returns the attributes contained in this STUN message.
327    pub fn attributes(&self) -> &[StunAttribute] {
328        &self.attributes
329    }
330
331    /// Returns the attribute if the message contains the attribute type
332    /// or None if there is no such attribute.
333    /// If there are more than one attributes of this type, this function
334    /// will return the first one.
335    pub fn get<A>(&self) -> Option<&StunAttribute>
336    where
337        A: StunAttributeType,
338    {
339        self.attributes
340            .iter()
341            .find(|&attr| attr.attribute_type() == A::get_type())
342    }
343}
344
345#[cfg(test)]
346mod tests {
347    use crate::{message::*, methods::BINDING};
348
349    #[test]
350    fn message_class() {
351        let cls = MessageClass::try_from(0).expect("Can not create MessageClass");
352        assert_eq!(cls.as_u16(), 0);
353
354        let cls = MessageClass::try_from(1).expect("Can not create MessageClass");
355        assert_eq!(cls.as_u16(), 1);
356
357        let cls = MessageClass::try_from(2).expect("Can not create MessageClass");
358        assert_eq!(cls.as_u16(), 2);
359
360        let cls = MessageClass::try_from(3).expect("Can not create MessageClass");
361        assert_eq!(cls.as_u16(), 3);
362
363        MessageClass::try_from(4).expect_err("MessageClass should not be created");
364    }
365
366    #[test]
367    fn message_method() {
368        let m = MessageMethod::try_from(0x0000).expect("Can not create MessageMethod");
369        assert_eq!(m.as_u16(), 0x0000);
370
371        let m = MessageMethod::try_from(0x0001).expect("Can not create MessageMethod");
372        assert_eq!(m.as_u16(), 0x0001);
373
374        let m = MessageMethod::try_from(0x0FFF).expect("Can not create MessageMethod");
375        assert_eq!(m.as_u16(), 0x0FFF);
376
377        MessageMethod::try_from(0x1000).expect_err("MessageMethod should not be created");
378    }
379
380    #[test]
381    fn message_type() {
382        let cls = MessageClass::Request;
383        let method = MessageMethod::try_from(0x0001).expect("Can not create MessageMethod");
384        let msg_type = MessageType::new(method, cls);
385
386        assert_eq!(msg_type.class(), cls);
387        assert_eq!(msg_type.method(), method);
388
389        let mut buffer: [u8; 2] = [0; 2];
390        assert_eq!(msg_type.encode(&mut buffer), Ok(2));
391        assert_eq!(buffer, [0x00, 0x01]);
392    }
393
394    #[test]
395    fn encode_message_type() {
396        let cls = MessageClass::Request;
397        let method = MessageMethod::try_from(0x08D8).expect("Can not create MessageMethod");
398        let msg_type = MessageType::new(method, cls);
399
400        let mut buffer: [u8; 2] = [0; 2];
401        assert_eq!(msg_type.encode(&mut buffer), Ok(2));
402        assert_eq!(buffer, [0x22, 0xA8]);
403
404        let cls = MessageClass::Indication;
405        let msg_type = MessageType::new(method, cls);
406        let mut buffer: [u8; 2] = [0; 2];
407        assert_eq!(msg_type.encode(&mut buffer), Ok(2));
408        assert_eq!(buffer, [0x22, 0xB8]);
409
410        let cls = MessageClass::SuccessResponse;
411        let msg_type = MessageType::new(method, cls);
412        let mut buffer: [u8; 2] = [0; 2];
413        assert_eq!(msg_type.encode(&mut buffer), Ok(2));
414        assert_eq!(buffer, [0x23, 0xA8]);
415
416        let cls = MessageClass::ErrorResponse;
417        let msg_type = MessageType::new(method, cls);
418        let mut buffer: [u8; 2] = [0; 2];
419        assert_eq!(msg_type.encode(&mut buffer), Ok(2));
420        assert_eq!(buffer, [0x23, 0xB8]);
421
422        let cls = MessageClass::ErrorResponse;
423        let msg_type = MessageType::new(method, cls);
424        let mut buffer: [u8; 1] = [0; 1];
425        assert_eq!(
426            msg_type.encode(&mut buffer).expect_err("Error expected"),
427            StunErrorType::SmallBuffer
428        );
429    }
430
431    #[test]
432    fn message_type_from() {
433        let method = MessageMethod::try_from(0x08D8).expect("Can not create MessageMethod");
434
435        let buffer = [0x22, 0xA8];
436        let msg_type = MessageType::from(&buffer);
437        assert_eq!(msg_type.class(), MessageClass::Request);
438        assert_eq!(msg_type.method(), method);
439
440        let buffer = [0x22, 0xB8];
441        let msg_type = MessageType::from(&buffer);
442        assert_eq!(msg_type.class(), MessageClass::Indication);
443        assert_eq!(msg_type.method(), method);
444
445        let buffer = [0x23, 0xA8];
446        let msg_type = MessageType::from(&buffer);
447        assert_eq!(msg_type.class(), MessageClass::SuccessResponse);
448        assert_eq!(msg_type.method(), method);
449
450        let buffer = [0x23, 0xB8];
451        let msg_type = MessageType::from(&buffer);
452        assert_eq!(msg_type.class(), MessageClass::ErrorResponse);
453        assert_eq!(msg_type.method(), method);
454
455        let buffer = [0x23, 0xB8];
456        let msg_type = MessageType::from(&buffer);
457        assert_eq!(msg_type.class(), MessageClass::ErrorResponse);
458        assert_eq!(msg_type.method(), method);
459    }
460
461    #[test]
462    fn fmt() {
463        let cls = MessageClass::Request;
464        let method = MessageMethod::try_from(0x0001).expect("Can not create MessageMethod");
465        let msg_type = MessageType::new(method, cls);
466        let _val = format!("{:?}", msg_type);
467
468        let builder = StunMessageBuilder::new(BINDING, MessageClass::Request);
469        let _val = format!("{:?}", builder);
470
471        let msg = builder.build();
472        let _val = format!("{:?}", msg);
473    }
474}