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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct MessageType {
34 method: MessageMethod,
35 class: MessageClass,
36}
37
38impl MessageType {
39 pub fn new(method: MessageMethod, class: MessageClass) -> Self {
44 Self { method, class }
45 }
46
47 pub fn class(&self) -> MessageClass {
49 self.class
50 }
51
52 pub fn method(&self) -> MessageMethod {
54 self.method
55 }
56
57 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 let val = value & 0x3FFF;
71 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
131pub struct MessageMethod(pub(crate) u16);
132
133impl MessageMethod {
134 pub fn as_u16(&self) -> u16 {
136 self.0
137 }
138
139 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
167pub enum MessageClass {
168 Request,
170 Indication,
172 SuccessResponse,
174 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#[derive(Debug)]
251pub struct StunMessageBuilder(StunMessageParameters);
252
253impl StunMessageBuilder {
254 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 pub fn with_transaction_id(mut self, transaction_id: TransactionId) -> Self {
270 self.0.transaction_id = Some(transaction_id);
271 self
272 }
273
274 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 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#[derive(Debug)]
303pub struct StunMessage {
304 method: MessageMethod,
305 class: MessageClass,
306 transaction_id: TransactionId,
307 attributes: Vec<StunAttribute>,
308}
309
310impl StunMessage {
311 pub fn method(&self) -> MessageMethod {
313 self.method
314 }
315
316 pub fn class(&self) -> MessageClass {
318 self.class
319 }
320
321 pub fn transaction_id(&self) -> &TransactionId {
323 &self.transaction_id
324 }
325
326 pub fn attributes(&self) -> &[StunAttribute] {
328 &self.attributes
329 }
330
331 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}