1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

use std::io::{Cursor, Write};

pub use super::errors::{IntegrityKeyGenerationError, MessageDecodeError, MessageEncodeError};
use super::message::StunMessage;

use crate::attribute::StunAttribute;

impl StunMessage {
    /// Encodes the STUN message into a binary representation
    ///
    /// Arguments:
    ///
    /// * `integrity_password`: Optionally set key that will be used for message integrity generation. Required if a MessageIntegrity attribute is present.
    pub fn encode(&self, integrity_password: Option<&str>) -> Result<Vec<u8>, MessageEncodeError> {
        let attr_count = self.attributes.len();
        let mut cursor = Cursor::new(Vec::new());

        // Encode and write the header
        let encoded_header = &self.header.encode()?;
        cursor.write_all(encoded_header)?;

        // Mark that a message integrity attribute is present
        let mut msg_integrity_present = false;

        // Track for username/realm occurrences
        let mut username = None;
        let mut realm = None;

        // Process each attribute and encode it
        for (idx, attr) in self.attributes.clone().iter().enumerate() {
            let processed_attr = match attr {
                // Track the username attribute being set
                StunAttribute::Username { value } => {
                    username = Some(value.clone());

                    attr.clone()
                }
                // Track the realm attribute being set
                StunAttribute::Realm { value } => {
                    realm = Some(value.clone());

                    attr.clone()
                }
                // Process the Fingerprint attribute
                StunAttribute::Fingerprint { value } => {
                    // Make sure that the Fingerprint attribute is the last one
                    if attr_count - 1 != idx {
                        return Err(MessageEncodeError::IncorrectFingerprintAttributePosition {
                            attr_count,
                            fingerprint_attr_idx: idx,
                        });
                    }
                    // Check if it contains a placeholder value and replace it with the computed fingerprint
                    if *value == 0 {
                        // Update the encoded message length so the correct fingerprint can be calculated
                        self.set_message_length(cursor.get_mut(), 8);

                        // Update the fingerprint value
                        let fingerprint = Self::calculate_fingerprint(cursor.get_ref());

                        StunAttribute::Fingerprint { value: fingerprint }
                    } else {
                        attr.clone()
                    }
                }
                // Process the MessageIntegrity attribute
                StunAttribute::MessageIntegrity { key } => {
                    // Mark it as present
                    msg_integrity_present = true;

                    // In case of placeholder data, replace it with the calculated HMAC value
                    if key.is_empty() {
                        if let Some(integrity_password) = integrity_password {
                            // Calculate the integrity key
                            let integrity_key = Self::calculate_integrity_key(
                                integrity_password,
                                realm.clone(),
                                username.clone(),
                            )?;

                            let hmac =
                                Self::calculate_integrity_hash(&integrity_key, cursor.get_ref());

                            StunAttribute::MessageIntegrity { key: hmac }
                        } else {
                            // Return an error if no `integrity_password` is submitted with placeholder data
                            return Err(MessageEncodeError::MissingIntegrityPassword());
                        }
                    } else {
                        // Return pre-submitted data
                        attr.clone()
                    }
                }
                _ => {
                    // If a MessageIntegrity attribute has been already addded,
                    // no other attributes (except Fingerprint) can be added after it
                    if msg_integrity_present {
                        return Err(MessageEncodeError::AttributeAfterIntegrity());
                    }

                    attr.clone()
                }
            };

            // Encode and write the attribute
            let encoded_attr = processed_attr.encode(self.header.transaction_id)?;
            cursor.write_all(&encoded_attr)?;
        }

        // Update the encoded message length
        self.set_message_length(cursor.get_mut(), 0);

        // Return the encoded data
        Ok(cursor.get_ref().to_vec())
    }
}