stun_agent/
message.rs

1use stun_rs::{
2    MessageClass, MessageMethod, StunAttribute, StunAttributeType, StunMessage, StunMessageBuilder,
3    TransactionId,
4};
5
6/// Even though the STUN message is a collection of attributes, The [`StunAttribute`] is used to
7/// simplify the addition and removal of attributes. The
8/// [`RFC8489`](https://datatracker.ietf.org/doc/html/rfc8489) does not set neither a
9/// limit nor a specific order for the attributes that can be added to a message, nevertheless,
10/// there are certain restrictions that must be followed for the integrity and fingerprint
11/// attributes. The [`StunAttribute`] eases the manipulation of attributes
12/// while ensuring that the above restrictions are met.
13#[derive(Debug, Default, Clone)]
14pub struct StunAttributes {
15    attributes: Vec<StunAttribute>,
16    integrity: Option<StunAttribute>,
17    integrity_sha256: Option<StunAttribute>,
18    fingerprint: Option<StunAttribute>,
19}
20
21impl StunAttributes {
22    /// Adds a STUN attribute to the collection. If the attribute is already present,
23    /// it will be replaced.
24    pub fn add<T>(&mut self, attribute: T)
25    where
26        T: Into<StunAttribute>,
27    {
28        let attr = attribute.into();
29
30        if attr.is_message_integrity() {
31            self.integrity = Some(attr);
32        } else if attr.is_message_integrity_sha256() {
33            self.integrity_sha256 = Some(attr);
34        } else if attr.is_fingerprint() {
35            self.fingerprint = Some(attr);
36        } else if let Some(index) = self
37            .attributes
38            .iter()
39            .position(|a| a.attribute_type() == attr.attribute_type())
40        {
41            // If this kind of attribute is already present, replace it
42            self.attributes[index] = attr;
43        } else {
44            self.attributes.push(attr);
45        }
46    }
47
48    /// Removes a STUN attribute from the collection.
49    /// # Returns
50    /// The removed attribute if it was present.
51    pub fn remove<T>(&mut self) -> Option<StunAttribute>
52    where
53        T: StunAttributeType,
54    {
55        if let Some(attr) = &self.integrity {
56            if attr.attribute_type() == T::get_type() {
57                return self.integrity.take();
58            }
59        }
60        if let Some(attr) = &self.integrity_sha256 {
61            if attr.attribute_type() == T::get_type() {
62                return self.integrity_sha256.take();
63            }
64        }
65        if let Some(attr) = &self.fingerprint {
66            if attr.attribute_type() == T::get_type() {
67                return self.fingerprint.take();
68            }
69        }
70        if let Some(index) = self
71            .attributes
72            .iter()
73            .position(|a| a.attribute_type() == T::get_type())
74        {
75            return Some(self.attributes.remove(index));
76        }
77
78        None
79    }
80}
81
82impl From<StunAttributes> for Vec<StunAttribute> {
83    fn from(val: StunAttributes) -> Self {
84        let mut attributes = val.attributes;
85        if let Some(attr) = val.integrity {
86            attributes.push(attr);
87        }
88        if let Some(attr) = val.integrity_sha256 {
89            attributes.push(attr);
90        }
91        if let Some(attr) = val.fingerprint {
92            attributes.push(attr);
93        }
94        attributes
95    }
96}
97
98pub fn create_stun_message(
99    method: MessageMethod,
100    class: MessageClass,
101    transaction_id: Option<TransactionId>,
102    attributes: StunAttributes,
103) -> StunMessage {
104    let mut builder = StunMessageBuilder::new(method, class);
105    if let Some(transaction_id) = transaction_id {
106        builder = builder.with_transaction_id(transaction_id);
107    }
108
109    let attributes: Vec<StunAttribute> = attributes.into();
110    for attr in attributes {
111        builder = builder.with_attribute(attr);
112    }
113
114    builder.build()
115}
116
117#[cfg(test)]
118mod stun_message_tests {
119    use super::*;
120    use stun_rs::attributes::stun::{
121        ErrorCode, Fingerprint, MessageIntegrity, MessageIntegritySha256, Nonce, Realm, Software,
122        UserName,
123    };
124    use stun_rs::methods::BINDING;
125    use stun_rs::HMACKey;
126
127    const USERNAME: &str = "test-username";
128    const REALM: &str = "test-realm";
129    const NONCE: &str = "test-nonce";
130    const PASSWORD: &str = "test-password";
131
132    #[test]
133    fn test_add_stun_attribute() {
134        let mut attributes = StunAttributes::default();
135
136        attributes.add(UserName::try_from("test-username-1").expect("Failed to create username"));
137        attributes.add(UserName::try_from("test-username-2").expect("Failed to create username"));
138        attributes.add(Software::try_from("test-software-1").expect("Failed to create software"));
139        attributes.add(Software::try_from("test-software-2").expect("Failed to create software"));
140
141        let mut iter = attributes.attributes.iter();
142        let attr = iter.next().expect("Expected attribute UserName");
143        let username = attr.expect_user_name();
144        assert_eq!(username, "test-username-2");
145
146        let attr = iter.next().expect("Expected attribute Software");
147        let username = attr.expect_software();
148        assert_eq!(username, "test-software-2");
149
150        assert!(iter.next().is_none());
151    }
152
153    #[test]
154    fn test_remove_stun_attribute() {
155        let mut attributes = StunAttributes::default();
156
157        attributes.add(UserName::try_from("test-username-1").expect("Failed to create username"));
158        attributes.add(UserName::try_from("test-username-2").expect("Failed to create username"));
159        attributes.add(Software::try_from("test-software-1").expect("Failed to create software"));
160        attributes.add(Software::try_from("test-software-2").expect("Failed to create software"));
161
162        let attr = attributes
163            .remove::<UserName>()
164            .expect("Expected attribute UserName");
165        let username = attr.expect_user_name();
166        assert_eq!(username, "test-username-2");
167
168        assert!(attributes.remove::<ErrorCode>().is_none());
169
170        let mut iter = attributes.attributes.iter();
171        let attr = iter.next().expect("Expected attribute Software");
172        let username = attr.expect_software();
173        assert_eq!(username, "test-software-2");
174        assert!(iter.next().is_none());
175
176        let key = HMACKey::new_short_term("foo bar").expect("Failed to create HMACKey");
177        let mut attributes = StunAttributes::default();
178        attributes.add(MessageIntegrity::new(key.clone()));
179        attributes.add(MessageIntegritySha256::new(key));
180        attributes.add(Fingerprint::default());
181        attributes.add(UserName::try_from("test-username-1").expect("Failed to create username"));
182        attributes.add(Software::try_from("test-software-1").expect("Failed to create software"));
183        attributes.add(ErrorCode::from(
184            stun_rs::ErrorCode::new(420, "Unknown Attribute").expect("Failed to create error"),
185        ));
186
187        assert!(attributes.remove::<MessageIntegrity>().is_some());
188        assert!(attributes.remove::<MessageIntegritySha256>().is_some());
189        assert!(attributes.remove::<Fingerprint>().is_some());
190        assert!(attributes.remove::<Software>().is_some());
191        assert!(attributes.remove::<ErrorCode>().is_some());
192        assert!(attributes.remove::<UserName>().is_some());
193
194        assert!(attributes.attributes.is_empty());
195    }
196
197    #[test]
198    fn test_stun_attribute_position() {
199        let key = HMACKey::new_short_term(PASSWORD).expect("Failed to create HMACKey");
200        let username = UserName::new(USERNAME).expect("Failed to create UserName");
201        let realm = Realm::new(REALM).expect("Failed to create Realm");
202        let nonce = Nonce::new(NONCE).expect("Failed to create Nonce");
203        let integrity = MessageIntegrity::new(key.clone());
204        let integrity_sha256 = MessageIntegritySha256::new(key.clone());
205        let fingerprint = Fingerprint::default();
206
207        let mut attributes = StunAttributes::default();
208        attributes.add(integrity);
209        attributes.add(username);
210        attributes.add(fingerprint);
211        attributes.add(realm);
212        attributes.add(integrity_sha256);
213        attributes.add(nonce);
214
215        let vector: Vec<StunAttribute> = Vec::from(attributes);
216        let mut iter = vector.iter();
217        let attr = iter.next().expect("Expected attribute UserName");
218        assert!(attr.is_user_name());
219        let attr = iter.next().expect("Expected attribute Realm");
220        assert!(attr.is_realm());
221        let attr = iter.next().expect("Expected attribute Nonce");
222        assert!(attr.is_nonce());
223        let attr = iter.next().expect("Expected attribute MessageIntegrity");
224        assert!(attr.is_message_integrity());
225        let attr = iter
226            .next()
227            .expect("Expected attribute MessageIntegritySha256");
228        assert!(attr.is_message_integrity_sha256());
229        let attr = iter.next().expect("Expected attribute Fingerprint");
230        assert!(attr.is_fingerprint());
231        assert!(iter.next().is_none());
232    }
233
234    #[test]
235    fn test_create_stun_message() {
236        let key = HMACKey::new_short_term("foo bar").expect("Failed to create HMACKey");
237        let transaction_id = TransactionId::default();
238
239        let mut attributes = StunAttributes::default();
240
241        attributes.add(Fingerprint::default());
242        attributes.add(UserName::try_from("test-username-1").expect("Failed to create username"));
243        attributes.add(MessageIntegritySha256::new(key.clone()));
244        attributes.add(UserName::try_from("test-username-2").expect("Failed to create username"));
245        attributes.add(MessageIntegrity::new(key));
246        attributes.add(Software::try_from("test-software-1").expect("Failed to create software"));
247        attributes.add(Software::try_from("test-software-2").expect("Failed to create software"));
248
249        let message = create_stun_message(
250            BINDING,
251            MessageClass::Request,
252            Some(transaction_id),
253            attributes,
254        );
255
256        assert_eq!(message.method(), BINDING);
257        assert_eq!(message.class(), MessageClass::Request);
258        assert_eq!(message.transaction_id(), &transaction_id);
259
260        let attributes = message.attributes();
261        assert_eq!(attributes.len(), 5);
262
263        let mut iter = attributes.iter();
264        let attr = iter.next().expect("Expected attribute UserName");
265        let username = attr.expect_user_name();
266        assert_eq!(username, "test-username-2");
267
268        let attr = iter.next().expect("Expected attribute Software");
269        let username = attr.expect_software();
270        assert_eq!(username, "test-software-2");
271
272        let attr = iter.next().expect("Expected attribute MessageIntegrity");
273        assert!(attr.is_message_integrity());
274
275        let attr = iter
276            .next()
277            .expect("Expected attribute MessageIntegritySha256");
278        assert!(attr.is_message_integrity_sha256());
279
280        let attr = iter.next().expect("Expected attribute Fingerprint");
281        assert!(attr.is_fingerprint());
282
283        assert!(iter.next().is_none());
284    }
285}
stun_agent/message.rs

stun_agent/
message.rs
