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(&mut 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(&mut cursor.get_mut(), 0);

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