stun_coder/message/encode.rs
1use std::io::{Cursor, Write};
2
3pub use super::errors::{IntegrityKeyGenerationError, MessageDecodeError, MessageEncodeError};
4use super::message::StunMessage;
5
6use crate::attribute::StunAttribute;
7
8impl StunMessage {
9 /// Encodes the STUN message into a binary representation
10 ///
11 /// Arguments:
12 ///
13 /// * `integrity_password`: Optionally set key that will be used for message integrity generation. Required if a MessageIntegrity attribute is present.
14 pub fn encode(&self, integrity_password: Option<&str>) -> Result<Vec<u8>, MessageEncodeError> {
15 let attr_count = self.attributes.len();
16 let mut cursor = Cursor::new(Vec::new());
17
18 // Encode and write the header
19 let encoded_header = &self.header.encode()?;
20 cursor.write_all(encoded_header)?;
21
22 // Mark that a message integrity attribute is present
23 let mut msg_integrity_present = false;
24
25 // Track for username/realm occurrences
26 let mut username = None;
27 let mut realm = None;
28
29 // Process each attribute and encode it
30 for (idx, attr) in self.attributes.clone().iter().enumerate() {
31 let processed_attr = match attr {
32 // Track the username attribute being set
33 StunAttribute::Username { value } => {
34 username = Some(value.clone());
35
36 attr.clone()
37 }
38 // Track the realm attribute being set
39 StunAttribute::Realm { value } => {
40 realm = Some(value.clone());
41
42 attr.clone()
43 }
44 // Process the Fingerprint attribute
45 StunAttribute::Fingerprint { value } => {
46 // Make sure that the Fingerprint attribute is the last one
47 if attr_count - 1 != idx {
48 return Err(MessageEncodeError::IncorrectFingerprintAttributePosition {
49 attr_count,
50 fingerprint_attr_idx: idx,
51 });
52 }
53 // Check if it contains a placeholder value and replace it with the computed fingerprint
54 if *value == 0 {
55 // Update the encoded message length so the correct fingerprint can be calculated
56 self.set_message_length(cursor.get_mut(), 8);
57
58 // Update the fingerprint value
59 let fingerprint = Self::calculate_fingerprint(cursor.get_ref());
60
61 StunAttribute::Fingerprint { value: fingerprint }
62 } else {
63 attr.clone()
64 }
65 }
66 // Process the MessageIntegrity attribute
67 StunAttribute::MessageIntegrity { key } => {
68 // Mark it as present
69 msg_integrity_present = true;
70
71 // In case of placeholder data, replace it with the calculated HMAC value
72 if key.is_empty() {
73 if let Some(integrity_password) = integrity_password {
74 // Calculate the integrity key
75 let integrity_key = Self::calculate_integrity_key(
76 integrity_password,
77 realm.clone(),
78 username.clone(),
79 )?;
80
81 let hmac =
82 Self::calculate_integrity_hash(&integrity_key, cursor.get_ref());
83
84 StunAttribute::MessageIntegrity { key: hmac }
85 } else {
86 // Return an error if no `integrity_password` is submitted with placeholder data
87 return Err(MessageEncodeError::MissingIntegrityPassword());
88 }
89 } else {
90 // Return pre-submitted data
91 attr.clone()
92 }
93 }
94 _ => {
95 // If a MessageIntegrity attribute has been already addded,
96 // no other attributes (except Fingerprint) can be added after it
97 if msg_integrity_present {
98 return Err(MessageEncodeError::AttributeAfterIntegrity());
99 }
100
101 attr.clone()
102 }
103 };
104
105 // Encode and write the attribute
106 let encoded_attr = processed_attr.encode(self.header.transaction_id)?;
107 cursor.write_all(&encoded_attr)?;
108 }
109
110 // Update the encoded message length
111 self.set_message_length(cursor.get_mut(), 0);
112
113 // Return the encoded data
114 Ok(cursor.get_ref().to_vec())
115 }
116}