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}