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())
}
}