stun_agent/
client.rs

1use crate::events::{StunTransactionError, StuntClientEvent, TransactionEventHandler};
2use crate::fingerprint::{add_fingerprint_attribute, validate_fingerprint};
3use crate::integrity::IntegrityError;
4use crate::lt_cred_mech::LongTermCredentialClient;
5use crate::message::{create_stun_message, StunAttributes};
6use crate::rtt::{RttCalcuator, DEFAULT_GRANULARITY};
7use crate::st_cred_mech::ShortTermCredentialClient;
8use crate::timeout::{RtoManager, StunMessageTimeout, DEFAULT_RC, DEFAULT_RM, DEFAULT_RTO};
9use crate::{CredentialMechanism, StunAgentError, StunPacket};
10use log::{debug, info, warn};
11use std::collections::HashMap;
12use std::time::{Duration, Instant};
13use stun_rs::attributes::stun::UserName;
14use stun_rs::error::StunEncodeError;
15use stun_rs::{
16    HMACKey, MessageClass, MessageDecoder, MessageEncoder, MessageMethod, StunMessage,
17    TransactionId,
18};
19
20// Maximum number of outstanding transactions to the same server
21pub const DEFAULT_MAX_TRANSACTIONS: usize = 10;
22
23/// Description of the transport reliability, for STUN protocol
24/// communication. It can be reliable or unreliable depending on
25/// whether this is a `UDP` or `TCP` connection.
26#[derive(Debug)]
27pub enum TransportReliability {
28    /// Reliable transport, such as `TCP`, where [`Duration`] represents the maximum
29    /// time to wait for a response.
30    Reliable(Duration),
31    /// Unreliable transport, such as `UDP`, where the [`RttConfig`] contains the
32    /// parameters to calculate the re-transmission timeout.
33    Unreliable(RttConfig),
34}
35
36/// When using unreliable transport, such as `UDP`, the re-transmission timeout
37/// is calculated using the following parameters.
38/// The `RTO` is an estimate of the round-trip time (`RTT`) and is computed as described
39/// in [`RFC6298`](`https://datatracker.ietf.org/doc/html/rfc6298`), with two exceptions.
40/// First, the initial value for `RTO` SHOULD be greater than or equal to 500 ms.
41/// Second, the value of `RTO` SHOULD NOT be rounded up to the nearest second. Rather,
42/// a 1 ms accuracy SHOULD be maintained.
43#[derive(Debug)]
44pub struct RttConfig {
45    /// Initial re-transmission timeout. Default is 500 ms.
46    pub rto: Duration,
47    /// The clock granularity to use for the `RTT` calculation. Default is 1 ms.
48    pub granularity: Duration,
49    /// The last re-transmission multiplier. Default is 16.
50    pub rm: u32,
51    /// Re-transmission counter. Default is 7.
52    pub rc: u32,
53}
54
55impl Default for RttConfig {
56    fn default() -> Self {
57        Self {
58            rto: DEFAULT_RTO,
59            granularity: DEFAULT_GRANULARITY,
60            rm: DEFAULT_RM,
61            rc: DEFAULT_RC,
62        }
63    }
64}
65
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67enum StunClientMessageClass {
68    Request,
69    Indication,
70}
71
72#[derive(Debug)]
73struct StunClientParameters {
74    user_name: Option<String>,
75    password: Option<String>,
76    mechanism: Option<CredentialMechanism>,
77    reliability: TransportReliability,
78    fingerprint: bool,
79    max_transactions: usize,
80}
81
82/// Builder for the STUN client. It allows to configure the client
83/// with the required parameters for the STUN usage.
84/// ```rust
85/// # use stun_agent::{CredentialMechanism, RttConfig, StunAgentError, StunClienteBuilder, TransportReliability};
86/// # fn main() -> Result<(), StunAgentError> {
87/// // Next example shows how to create a STUN client that uses
88/// // long-term credentials to authenticate with the server over
89/// // an unreliable transport and mandates the use of the FINGERPRINT
90/// // attribute.
91/// let client = StunClienteBuilder::new(
92///     TransportReliability::Unreliable(RttConfig::default()))
93///         .with_mechanism("user", "password", CredentialMechanism::LongTerm)
94///         .with_fingerprint()
95///         .build()?;
96/// # Ok(())
97/// # }
98/// ```
99#[derive(Debug)]
100pub struct StunClienteBuilder(StunClientParameters);
101
102impl StunClienteBuilder {
103    /// Creates a new STUN client builder with the given [`TransportReliability`].
104    pub fn new(reliability: TransportReliability) -> StunClienteBuilder {
105        Self(StunClientParameters {
106            user_name: None,
107            password: None,
108            mechanism: None,
109            reliability,
110            fingerprint: false,
111            max_transactions: DEFAULT_MAX_TRANSACTIONS,
112        })
113    }
114
115    /// Sets the maximum number of outstanding transactions to the same server.
116    /// The default value is 10.
117    pub fn with_max_transactions(mut self, max_transactions: usize) -> Self {
118        self.0.max_transactions = max_transactions;
119        self
120    }
121
122    /// Sets the credentials for the STUN client.
123    pub fn with_mechanism<U, P>(
124        mut self,
125        user_name: U,
126        password: P,
127        mechanism: CredentialMechanism,
128    ) -> Self
129    where
130        U: Into<String>,
131        P: Into<String>,
132    {
133        self.0.user_name = Some(user_name.into());
134        self.0.password = Some(password.into());
135        self.0.mechanism = Some(mechanism);
136        self
137    }
138
139    /// Sets the use of the FINGERPRINT attribute in the STUN messages. The
140    /// FINGERPRINT mechanism is not backwards compatible with
141    /// [`RFC3489`](https://datatracker.ietf.org/doc/html/rfc3489) and
142    /// cannot be used in environments where such compatibility is required.
143    pub fn with_fingerprint(mut self) -> Self {
144        self.0.fingerprint = true;
145        self
146    }
147
148    /// Builds the STUN client with the given parameters.
149    pub fn build(self) -> Result<StunClient, StunAgentError> {
150        StunClient::new(self.0)
151    }
152}
153
154#[derive(Debug)]
155enum CredentialMechanismClient {
156    ShortTerm(ShortTermCredentialClient),
157    LongTerm(LongTermCredentialClient),
158}
159
160impl CredentialMechanismClient {
161    fn prepare_request(&mut self, attributes: &mut StunAttributes) -> Result<(), StunAgentError> {
162        match self {
163            CredentialMechanismClient::ShortTerm(mechanism) => {
164                mechanism.add_attributes(attributes);
165                Ok(())
166            }
167            CredentialMechanismClient::LongTerm(mechanism) => mechanism.prepare_request(attributes),
168        }
169    }
170
171    fn prepare_indication(
172        &mut self,
173        attributes: &mut StunAttributes,
174    ) -> Result<(), StunAgentError> {
175        match self {
176            CredentialMechanismClient::ShortTerm(mechanism) => {
177                mechanism.add_attributes(attributes);
178                Ok(())
179            }
180            CredentialMechanismClient::LongTerm(mechanism) => {
181                mechanism.prepare_indication(attributes)
182            }
183        }
184    }
185
186    fn recv_message(
187        &mut self,
188        raw_data: &[u8],
189        message: &StunMessage,
190    ) -> Result<(), IntegrityError> {
191        match self {
192            CredentialMechanismClient::ShortTerm(mechanism) => {
193                mechanism.recv_message(raw_data, message)
194            }
195            CredentialMechanismClient::LongTerm(mechanism) => {
196                mechanism.recv_message(raw_data, message)
197            }
198        }
199    }
200
201    fn signal_protection_violated_on_timeout(&mut self, transaction_id: &TransactionId) -> bool {
202        match self {
203            CredentialMechanismClient::ShortTerm(mechanism) => {
204                mechanism.signal_protection_violated_on_timeout(transaction_id)
205            }
206            CredentialMechanismClient::LongTerm(mechanism) => {
207                mechanism.signal_protection_violated_on_timeout(transaction_id)
208            }
209        }
210    }
211}
212
213#[derive(Debug)]
214struct StunTransaction {
215    instant: Option<Instant>,
216    packet: StunPacket,
217    rtos: RtoManager,
218}
219
220#[derive(Debug)]
221struct RttHandler {
222    rtt: RttCalcuator,
223    rm: u32,
224    rc: u32,
225    last_request: Option<Instant>,
226}
227
228#[derive(Debug)]
229enum StunRttCalcuator {
230    Reliable(Duration),
231    Unreliable(RttHandler),
232}
233
234impl From<TransportReliability> for StunRttCalcuator {
235    fn from(reliability: TransportReliability) -> Self {
236        match reliability {
237            TransportReliability::Reliable(timeout) => StunRttCalcuator::Reliable(timeout),
238            TransportReliability::Unreliable(config) => StunRttCalcuator::Unreliable(RttHandler {
239                rtt: RttCalcuator::new(config.rto, config.granularity),
240                rm: config.rm,
241                rc: config.rc,
242                last_request: None,
243            }),
244        }
245    }
246}
247
248/// A STUN client is an entity that sends STUN requests and receives STUN responses
249/// and STUN indications. A STUN client can also send indications.
250///
251/// # [`StunClient`]
252///
253/// This is the main entity used to interact with the STUN server.
254/// The [`StunClient`] provides the tools required to implement different
255/// STUN [usages](https://datatracker.ietf.org/doc/html/rfc8489#section-13)
256/// over the STUN protocol in an easy and efficient way.
257///
258/// # API considerations
259///
260/// Since the [`StunClient`] abstains from performing any I/O operations, the controller assumes
261/// responsibility for managing input and output buffers, timeouts, and client-generated
262/// events. The implementation of this controller is entirely at the user's discretion and
263/// does not enforce the use of any specific I/O stack or asynchronous framework. This
264/// abstraction imposes certain guidelines to ensure the protocol's proper functioning.
265/// Consequently, users must consider the following technical aspects:
266/// - The controller must capture and handle any events that the client may generate after
267/// interacting with the library.
268/// - The controller must handle the input and output buffers that the client will use to send and
269/// receive data from the server.
270/// - Timing management falls under the controller's jurisdiction, as the client lacks internal
271/// time-handling mechanisms. The controller must define transaction timeouts and inform the client
272/// upon their expiration. For supporting timed events, the API exposes an [`Instant`] parameter
273/// to the controller, facilitating specification of event occurrence times.
274///
275/// # Design considerations
276///
277/// Most Sans I/O implementations are structured around a state machine that responds to events
278/// generated by both the client and the server. Each event triggers the generation of output buffers,
279/// timers, or additional events. This foundational concept is illustrated in the following API:
280///
281/// ```no_run
282/// # use stun_agent::StuntClientEvent;
283/// # fn handle_data(in_bytes: &[u8]) -> Vec<StuntClientEvent> { vec![] }
284/// # fn perform_action() -> Vec<u8> { vec![] }
285/// # let in_bytes = [];
286/// let events = handle_data(&in_bytes);
287/// let out_bytes = perform_action();
288/// ```
289///
290/// However, the STUN requirements introduce complexity to the API. The aforementioned API alone
291/// does not suffice to manage STUN intricacies. For instance, the `handle_data` function might fail
292/// and trigger events even in case of failures. The STUN client needs to manage these events and
293/// generate further events for the controller. This implementation could have been realized as follows:
294///
295/// ```no_run
296/// # use stun_agent::StuntClientEvent;
297/// # use stun_agent::StunAgentError;
298/// fn handle_data(in_bytes: &[u8])
299///     -> Result<Vec<StuntClientEvent>, (StunAgentError, Vec<StuntClientEvent>)> {
300///    // Implementation
301/// #   Ok(vec![])
302/// }
303/// ```
304///
305/// The design of this API necessitates that the caller manages both errors and the events they generate.
306/// This approach can lead to increased complexity and maintenance challenges in the caller's code.
307/// For instance, the caller may employ a match expression when invoking the function to handle both
308/// success outcomes and the errors and resulting events in case of failure:
309///
310/// ```no_run
311/// # use stun_agent::StuntClientEvent;
312/// # use stun_agent::StunAgentError;
313/// # fn handle_data(in_bytes: &[u8]) -> Result<Vec<StuntClientEvent>, (StunAgentError, Vec<StuntClientEvent>)> { Ok(vec![]) }
314/// # fn handle_events(events: Vec<StuntClientEvent>) {}
315/// # fn handle_error(error: StunAgentError) {}
316/// # let in_bytes = [];
317/// let response = handle_data(&in_bytes);
318/// match response {
319///    Ok(events) => {
320///       handle_events(events);
321///   },
322///   Err((error, events)) => {
323///       handle_error(error);
324///       handle_events(events);
325///   },
326/// }
327/// ```
328///
329/// As observed, managing events in both success and failure scenarios indicates a sub-optimal design.
330/// Consequently, the STUN client API is structured to enable the caller to pull [`events`](`Self::events`) generated by
331/// the client. While this approach offers a more ergonomic event handling mechanism, it requires the
332/// caller to actively retrieve and process events from the client.
333///
334/// ```no_run
335/// # use stun_agent::StunAgentError;
336/// # type ClientData = u8;
337/// fn handle_data(in_bytes: &[u8]) -> Result<ClientData, StunAgentError> {
338///   // Implementation
339/// # Ok(ClientData::default())
340/// }
341/// ```
342///
343/// And the controller's code would look like this:
344///
345/// ```no_run
346/// # use stun_agent::StunAgentError;
347/// # type ClientData = u8;
348/// # fn handle_data(in_bytes: &[u8]) -> Result<ClientData, StunAgentError> {Ok(ClientData::default())}
349/// # fn pull_events() -> Vec<u8> { vec![] }
350/// # fn main() -> Result<(), StunAgentError> {
351/// # let in_bytes = [];
352/// let data = handle_data(&in_bytes)?;
353/// // Now we can pull events from the client
354/// let events = pull_events();
355/// # Ok(())
356/// # }
357/// ```
358///
359/// Moreover, this type of API not only facilitates the retrieval of events but also allows for the
360/// retrieval of data generated by the client. For instance, the [`send_request`](`Self::send_request`)
361/// method returns the [`TransactionId`] of the request, which the controller can use to manage outgoing
362/// transactions.
363///
364/// <div class="warning">
365///
366/// Events are overwritten whenever a new operation is performed on the client. Therefore, the controller
367/// must ensure that all events are processed before initiating any new operations. In multi-threaded
368/// environments, the controller must also synchronize operations and event retrieval to maintain
369/// consistency and prevent data loss.
370///
371/// </div>
372///
373/// ## Input and Output
374///
375/// The STUN client does not perform any I/O operations. Instead, the controller is responsible for
376/// managing input and output buffers. Memory allocation is delegated to the controller, which must
377/// provide the buffers used by the client. This approach reduces the client's memory footprint and
378/// enhances performance by enabling more sophisticated memory management strategies, such as memory
379/// pools, where buffers can be reused to minimize memory allocation overhead.
380///
381/// ## Timing Management
382///
383/// The STUN client does not manage timing internally. Instead, the controller is responsible for setting
384/// timeouts and managing transaction timing. The API provides an [`Instant`] parameter to the controller,
385/// allowing it to specify event occurrence times. Timing consistency across operations is crucial,
386/// meaning that time must monotonically increase to ensure the proper functioning of the client.
387///
388/// Exposing the [`Instant`] parameter in the API might seem counter intuitive, as it requires the controller
389/// to manage time. However, this design choice ensures that the client remains agnostic to time
390/// management, granting the controller full control over the internal state machine. This approach
391/// facilitates comprehensive testing of complex scenarios by enabling deterministic time control without
392/// the need to mock time.
393///
394/// ## Timeouts
395///
396/// Timeouts specify the maximum duration the client will wait for an event to occur. The STUN client
397/// uses timeouts to manage transactions and prevent indefinite waiting for responses. If a response
398/// is not received within the designated timeout period, the client generates a timeout event, marking
399/// the transaction as failed. Timeouts are also employed to manage re-transmissions of requests sent
400/// over unreliable transports. When the client needs to set a timeout for a re-transmission, it generates
401/// a [`RestransmissionTimeOut`](`crate::StuntClientEvent::RestransmissionTimeOut`) event, which is then
402/// notified to the controller when the events are pulled.
403///
404/// If multiple timeouts are scheduled, the client will only notify the controller of the most recent
405/// timeout. This approach allows the controller to manage timeouts more efficiently, ensuring that
406/// only one timeout needs to be handled at a time.
407///
408/// Managing timeouts is the responsibility of the controller; the STUN client will only provide the
409/// timeout duration. If the timeout is not canceled, the controller must call the
410/// [`on_timeout`](`crate::StunClient::on_timeout`) method to inform the client that the timeout
411/// has been reached.
412///
413/// Timeouts are identified by a [`TransactionId`]. When a timeout is canceled for any reason, the
414/// client will notify the controller either by setting a new timeout with a different [`TransactionId`]
415/// or by not setting any timeout event at all.
416///
417/// # Usage
418///
419/// The following example demonstrates how to create a STUN client and send a
420/// BINDING indication to a STUN server.
421///
422/// ```rust
423/// # use stun_agent::{RttConfig, StunAttributes, StunClienteBuilder, StuntClientEvent, TransportReliability};
424/// # use stun_rs::methods::BINDING;
425/// # use std::time::Instant;
426///
427/// // We use a client builder to create a STUN client, for this example,
428/// // the client will be used over an unreliable transport such as UDP.
429/// // This client will no use any credential mechanism nor the FINGERPRINT
430/// // attributes. Besides, we configure the default parameters for the
431/// // re-transmission timeout.
432/// let mut client = StunClienteBuilder::new(
433///     TransportReliability::Unreliable(RttConfig::default()))
434///     .build()
435///     .unwrap();
436///
437/// // We create a STUN BINDING indication to send to the server.
438/// // According to the RFC8489, the BINDING indications does not require
439/// // any attributes.
440/// let mut attributes = StunAttributes::default();
441///
442/// // Since this is a library implementation without direct I/O operations,
443/// // no input or output will be handled by the stack. Instead, we need to
444/// // access the output buffer event provided by the client to send the data
445/// // through the socket.
446/// // Besides, no buffer allocations will be performed by the library, so the
447/// // controller must provide the buffer that will be used to send the data.
448/// // This allow the library to reduce the memory footprint and improve the
449/// // performance, being flexible to allow more complex usages of memory such
450/// // as memory pools where buffers can be reused to minimize the memory
451/// // allocation overhead.
452/// let buffer = vec![0; 1024];
453/// client.send_indication(BINDING, attributes, buffer).unwrap();
454///
455/// // Pull events from the client
456/// let events = client.events();
457///
458/// // Only one output packect event is expected. This event must contain the
459/// // buffer that will be sent to the server. Because indications do not require
460/// // a response, no timeouts will be set for this transaction.
461/// assert_eq!(events.len(), 1);
462/// let mut iter = events.iter();
463///
464/// // Next event already contains the buffer that needs to be send to the server.
465/// let StuntClientEvent::OutputPacket(buffer) = iter
466///     .next()
467///     .expect("Expected event")
468/// else {
469///     panic!("Expected OutputBuffer event");
470/// };
471/// ```
472///
473/// In the following example we are going to use the STUN client to send a BINDING request to a STUN server.
474/// Requests require a response from the server, so the client will set a timeout for the transaction.
475/// The response must arrive before the timeout is reached, otherwise the client will generate a timeout event
476/// and will mark the transaction as failed.
477/// ```rust
478/// # use stun_agent::{RttConfig, StunAttributes, StunClienteBuilder, StuntClientEvent, TransportReliability};
479/// # use stun_rs::methods::BINDING;
480/// # use std::time::Instant;
481/// # let mut client = StunClienteBuilder::new(TransportReliability::Unreliable(RttConfig::default()))
482/// #    .build().unwrap();
483///
484/// // We create a STUN BINDING request to send to the server.
485/// // According to the RFC8489, the BINDING request does not require
486/// // any attributes.
487/// let instant = std::time::Instant::now();
488/// let mut attributes = StunAttributes::default();
489/// let buffer = vec![0; 1024];
490/// let transaction_id = client
491///     .send_request(BINDING, attributes, buffer, instant)
492///     .unwrap();
493///
494/// // Pull events from the client
495/// let events = client.events();
496///
497/// // Two events are expected, the first one is the output buffer event
498/// // and the second one is the timeout event.
499/// assert_eq!(events.len(), 2);
500/// let mut iter = events.iter();
501/// // Next event already contains the buffer that needs to be send to the server.
502/// let StuntClientEvent::OutputPacket(buffer) = iter
503///     .next()
504///     .expect("Expected event")
505/// else {
506///     panic!("Expected OutputBuffer event");
507/// };
508/// // Next event indicates that the user must set a timeout for the transaction
509/// // identified by the transaction_id.
510/// let StuntClientEvent::RestransmissionTimeOut((id, duration)) = iter
511///     .next()
512///     .expect("Expected event")
513/// else {
514///     panic!("Expected RestransmissionTimeOut event");
515/// };
516/// assert_eq!(id, &transaction_id);
517///
518/// // Now the controller should set a timout of `duration` for the transaction
519/// // identified by `id`. After the timeout is reached, the controller must call
520/// // the `on_timeout` method to notify the client that the time has expired.
521///
522/// // We re going to simulate the timeout event by calling the `on_timeout` method.
523/// let instant = instant + *duration;
524/// client.on_timeout(instant);
525///
526/// // Pull events from the client
527/// let events = client.events();
528///
529/// // Two events are expected, the first one is the retransmission of the requests,
530/// // and the second one is the new timeout set for the transaction.
531/// assert_eq!(events.len(), 2);
532/// let mut iter = events.iter();
533///
534/// // Next event contains the buffer that needs to be retransmitted.
535/// let StuntClientEvent::OutputPacket(buffer) = iter
536///     .next()
537///     .expect("Expected event")
538/// else {
539///         panic!("Expected OutputBuffer event");
540/// };
541/// let StuntClientEvent::RestransmissionTimeOut((id, duration)) = iter
542///     .next()
543///     .expect("Expected event")
544/// else {
545///    panic!("Expected RestransmissionTimeOut event");
546/// };
547/// assert_eq!(id, &transaction_id);
548/// ```
549///
550/// When sending over an unreliable transport, the client SHOULD re-transmit a STUN request
551/// message starting with an interval of `RTO` ("Re-transmission `TimeOut`"), doubling after
552/// each re-transmission until a final timeout is reached. By default, if the controller does
553/// not set a different value, the default timeout is 39500 ms for both, reliable and not
554/// reliable transports. If the client has not received a response after that time, the client
555/// will consider the transaction to have timed out, and an event of type
556/// [`TransactionFailed`](crate::StuntClientEvent::TransactionFailed) will be generated the
557/// next time that events were pulled with the error
558/// [`TimedOut`](crate::StunTransactionError::TimedOut) for the transaction.
559///
560/// To finish, the next example shows how to handle buffers received from the server. Raw buffers will
561/// be processed by the client to generate events that can be pulled by the controller.
562///```rust
563/// # use stun_agent::{RttConfig, StunAttributes, StunClienteBuilder, StuntClientEvent, TransportReliability};
564/// # use stun_rs::methods::BINDING;
565/// # use std::time::Instant;
566/// # use stun_rs::MessageClass::Indication;
567/// # let mut client = StunClienteBuilder::new(TransportReliability::Unreliable(RttConfig::default()))
568/// #    .build().unwrap();
569/// // Buffer received from the server
570/// let buffer = [
571///    0x00, 0x11, 0x00, 0x00, // BINDING Indication type and message length
572///    0x21, 0x12, 0xA4, 0x42, // Magic cookie
573///    0xB8, 0xC2, 0x8E, 0x1A, // }
574///    0x41, 0x05, 0x18, 0x56, // }  Transaction ID
575///    0x3E, 0xFC, 0xCF, 0x5D, // }
576/// ];
577///
578/// // Process buffer
579/// client.on_buffer_recv(&buffer, Instant::now()).unwrap();
580///
581/// // Pull events from the client
582/// let events = client.events();
583///
584/// // There must be only one events with the STUN message received
585/// assert_eq!(events.len(), 1);
586///
587/// let mut iter = events.iter();
588/// let StuntClientEvent::StunMessageReceived(msg) = iter
589///     .next()
590///     .expect("Expected event")
591/// else {
592///    panic!("Expected StunMessageReceived event");
593/// };
594/// assert_eq!(msg.method(), BINDING);
595/// assert_eq!(msg.class(), Indication);
596///
597/// // No attributes in the message
598/// assert_eq!(msg.attributes().len(), 0);
599///```
600
601#[derive(Debug)]
602pub struct StunClient {
603    mechanism: Option<CredentialMechanismClient>,
604    encoder: MessageEncoder,
605    decoder: MessageDecoder,
606    use_fingerprint: bool,
607    timeouts: StunMessageTimeout,
608    rtt: StunRttCalcuator,
609    transactions: HashMap<TransactionId, StunTransaction>,
610    transaction_events: TransactionEventHandler,
611    max_transactions: usize,
612}
613
614impl StunClient {
615    fn new(params: StunClientParameters) -> Result<Self, StunAgentError> {
616        let rtt = StunRttCalcuator::from(params.reliability);
617        let is_reliable = matches!(rtt, StunRttCalcuator::Reliable(_));
618
619        let mechanism = match params.mechanism {
620            Some(value) => {
621                let user_name = params.user_name.ok_or_else(|| {
622                    StunAgentError::InternalError(String::from("User name is required"))
623                })?;
624                let password = params.password.ok_or_else(|| {
625                    StunAgentError::InternalError(String::from("Password is required"))
626                })?;
627                let user_name = UserName::new(user_name).map_err(|e| {
628                    StunAgentError::InternalError(format!("Failed to create user name: {}", e))
629                })?;
630                match value {
631                    CredentialMechanism::ShortTerm(integrity) => Some(
632                        CredentialMechanismClient::ShortTerm(ShortTermCredentialClient::new(
633                            user_name,
634                            HMACKey::new_short_term(password).map_err(|e| {
635                                StunAgentError::InternalError(format!(
636                                    "Failed to create HMAC key: {}",
637                                    e
638                                ))
639                            })?,
640                            integrity,
641                            is_reliable,
642                        )),
643                    ),
644                    CredentialMechanism::LongTerm => Some(CredentialMechanismClient::LongTerm(
645                        LongTermCredentialClient::new(user_name, password, is_reliable),
646                    )),
647                }
648            }
649            None => None,
650        };
651
652        Ok(Self {
653            mechanism,
654            encoder: Default::default(),
655            decoder: Default::default(),
656            use_fingerprint: params.fingerprint,
657            timeouts: StunMessageTimeout::default(),
658            rtt,
659            transactions: Default::default(),
660            transaction_events: Default::default(),
661            max_transactions: params.max_transactions,
662        })
663    }
664
665    fn prepare_request(&mut self, attributes: &mut StunAttributes) -> Result<(), StunAgentError> {
666        prepare_stun_message(
667            StunClientMessageClass::Request,
668            attributes,
669            self.mechanism.as_mut(),
670            self.use_fingerprint,
671        )
672    }
673
674    fn prepare_indication(
675        &mut self,
676        attributes: &mut StunAttributes,
677    ) -> Result<(), StunAgentError> {
678        prepare_stun_message(
679            StunClientMessageClass::Indication,
680            attributes,
681            self.mechanism.as_mut(),
682            self.use_fingerprint,
683        )
684    }
685
686    fn set_timeout(
687        &mut self,
688        transaction_id: TransactionId,
689        instant: Instant,
690    ) -> Result<RtoManager, StunAgentError> {
691        let mut rto_manager = match self.rtt {
692            StunRttCalcuator::Reliable(timeout) => RtoManager::new(timeout, 1, 1),
693            StunRttCalcuator::Unreliable(ref mut handler) => {
694                if let Some(last_request) = handler.last_request {
695                    if instant - last_request > Duration::from_secs(600) {
696                        debug!(
697                            "Current RTT value {}ms staled caused by inactivity. Resetting.",
698                            handler.rtt.rto().as_millis()
699                        );
700                        handler.rtt.reset();
701                    }
702                }
703                handler.last_request = Some(instant);
704                RtoManager::new(handler.rtt.rto(), handler.rm, handler.rc)
705            }
706        };
707
708        let timeout = rto_manager.next_rto(instant).ok_or_else(|| {
709            StunAgentError::InternalError(String::from("Can not calculate next RTO"))
710        })?;
711        self.timeouts.add(instant, timeout, transaction_id);
712        debug!("[{:?}] Set timeout {:?}", transaction_id, timeout);
713        Ok(rto_manager)
714    }
715
716    fn transaction_finished(&mut self, transaction_id: &TransactionId, instant: Instant) {
717        self.timeouts.remove(transaction_id);
718        let Some(transaction) = self.transactions.remove(transaction_id) else {
719            debug!("[{:?}] Not found", transaction_id);
720            return;
721        };
722
723        let Some(sent_instant) = transaction.instant else {
724            // This means that there was a retransmission
725            return;
726        };
727        if let StunRttCalcuator::Unreliable(handler) = &mut self.rtt {
728            let new_rtt = instant - sent_instant;
729            debug!(
730                "[{:?}] RTT calculation: sent={:?}, recv={:?}, rtt={:?}",
731                transaction_id, sent_instant, instant, new_rtt
732            );
733            handler.rtt.update(new_rtt);
734        }
735    }
736
737    /// Creates a STUN request.
738    ///
739    /// # Arguments
740    /// * `method` - The STUN [`MessageMethod`] to use.
741    /// * `attributes` - The [`StunAttributes`] to include in the request.
742    /// * `buffer` - The buffer to send with the request.
743    /// * `instant` - The instant when the request is sent.
744    /// # Returns
745    /// The [`TransactionId`] of the request on success. Otherwise, a [`StunAgentError`] is returned.
746    /// <div class="warning">
747    ///
748    /// After calling this method, the user should invoke [events](`Self::events`) to retrieve the events
749    /// generated by the agent.
750    ///
751    /// </div>
752    pub fn send_request(
753        &mut self,
754        method: MessageMethod,
755        mut attributes: StunAttributes,
756        buffer: Vec<u8>,
757        instant: Instant,
758    ) -> Result<TransactionId, StunAgentError> {
759        if self.transactions.len() >= self.max_transactions {
760            return Err(StunAgentError::MaxOutstandingRequestsReached);
761        }
762
763        self.prepare_request(&mut attributes)?;
764        let msg = create_stun_message(method, MessageClass::Request, None, attributes);
765        let packet = encode_buffer(&self.encoder, &msg, buffer).map_err(|e| {
766            StunAgentError::InternalError(format!("Failed to encode request message: {}", e))
767        })?;
768
769        let transaction = StunTransaction {
770            instant: Some(instant),
771            packet: packet.clone(),
772            rtos: self.set_timeout(*msg.transaction_id(), instant)?,
773        };
774        self.transactions.insert(*msg.transaction_id(), transaction);
775
776        let mut events = self.transaction_events.init();
777        events.push(StuntClientEvent::OutputPacket(packet));
778
779        // Add the most recent timout event if any
780        if let Some((id, left)) = self.timeouts.next_timeout(instant) {
781            events.push(StuntClientEvent::RestransmissionTimeOut((id, left)));
782        }
783
784        Ok(*msg.transaction_id())
785    }
786
787    /// Creates a STUN indication.
788    /// # Arguments
789    /// * `method` - The STUN [`MessageMethod`] to use.
790    /// * `attributes` - The [`StunAttributes`] to include in the indication.
791    /// * `buffer` - The buffer to send with the indication.
792    /// # Returns
793    /// The [`TransactionId`] of the indication on success. Otherwise, a [`StunAgentError`] is returned.
794    /// <div class="warning">
795    ///
796    /// After calling this method, the user should invoke [events](`Self::events`) to retrieve the events
797    /// generated by the agent.
798    ///
799    /// </div>
800    pub fn send_indication(
801        &mut self,
802        method: MessageMethod,
803        mut attributes: StunAttributes,
804        buffer: Vec<u8>,
805    ) -> Result<TransactionId, StunAgentError> {
806        self.prepare_indication(&mut attributes)?;
807        let msg = create_stun_message(method, MessageClass::Indication, None, attributes);
808        let packet = encode_buffer(&self.encoder, &msg, buffer).map_err(|e| {
809            StunAgentError::InternalError(format!("Failed to encode indication message: {}", e))
810        })?;
811
812        let mut events = self.transaction_events.init();
813        events.push(StuntClientEvent::OutputPacket(packet));
814
815        Ok(*msg.transaction_id())
816    }
817
818    /// Called when a buffer is received from the server.
819    /// # Arguments
820    /// * `buffer` - The buffer received from the server.
821    /// * `instant` - The instant when the buffer was received.
822    /// # Returns
823    /// A [`StunAgentError`] if the buffer is invalid or the transaction is discarded.
824    /// In the case when STUN is being    multiplexed with another protocol, an error
825    /// may indicate that this is not really a STUN message; in this case, the agent
826    /// should try to parse the message as a different protocol.
827    /// <div class="warning">
828    ///
829    /// After calling this method, the user should invoke [events](`Self::events`) to retrieve the events
830    /// generated by the agent.
831    ///
832    /// </div>
833    pub fn on_buffer_recv(
834        &mut self,
835        buffer: &[u8],
836        instant: Instant,
837    ) -> Result<(), StunAgentError> {
838        let (msg, _) = self.decoder.decode(buffer).map_err(|e| {
839            StunAgentError::InternalError(format!("Failed to decode message: {}", e))
840        })?;
841
842        match msg.class() {
843            MessageClass::Request => {
844                // A STUN client can receive STUN responses and STUN indications.
845                debug!(
846                    "Received STUN request with {:?}. Discarding.",
847                    msg.transaction_id()
848                );
849                return Err(StunAgentError::Discarded);
850            }
851            MessageClass::Indication => {
852                debug!("Received STUN indication with {:?}", msg.transaction_id());
853            }
854            MessageClass::SuccessResponse | MessageClass::ErrorResponse => {
855                // Check this is an outstanding transaction
856                if !self.transactions.contains_key(msg.transaction_id()) {
857                    debug!(
858                        "Received response with no matching {:?}. Discarding.",
859                        msg.transaction_id()
860                    );
861                    return Err(StunAgentError::Discarded);
862                }
863            }
864        }
865
866        // Validate fingerprint attribute
867        if self.use_fingerprint && !validate_fingerprint(buffer, &msg)? {
868            debug!(
869                "[{:?}] Fingerprint validation failed. Discarding.",
870                msg.transaction_id()
871            );
872            return Err(StunAgentError::Discarded);
873        }
874
875        // Validate message integrity
876        let mut integrity_event = None;
877        if let Some(mechanism) = &mut self.mechanism {
878            if let Err(e) = mechanism.recv_message(buffer, &msg) {
879                integrity_event = process_integrity_error(e, msg.transaction_id())?;
880            }
881        }
882
883        if msg.class() != MessageClass::Indication {
884            // finish outgoing transaction
885            self.transaction_finished(msg.transaction_id(), instant);
886        }
887
888        let mut events = self.transaction_events.init();
889        match integrity_event {
890            Some(event) => {
891                // notify the integrity issue
892                events.push(event);
893            }
894            None => {
895                // Notify the user about the received message
896                events.push(StuntClientEvent::StunMessageReceived(msg));
897            }
898        }
899
900        Ok(())
901    }
902
903    /// Called when a timeout event occurs.
904    /// # Arguments
905    /// * `instant` - The instant when the timeout event occurred.
906    /// <div class="warning">
907    ///
908    /// After calling this method, the user should invoke [events](`Self::events`) to retrieve the events
909    /// generated by the agent.
910    ///
911    /// </div>
912    pub fn on_timeout(&mut self, instant: Instant) {
913        let timed_out = self.timeouts.check(instant);
914        let mut events = self.transaction_events.init();
915
916        for transaction_id in timed_out {
917            if let Some(transaction) = self.transactions.get_mut(&transaction_id) {
918                match transaction.rtos.next_rto(instant) {
919                    Some(rto) => {
920                        // Cancel rtt calculation on retransmission
921                        transaction.instant = None;
922                        self.timeouts.add(instant, rto, transaction_id);
923                        debug!("set timeout {:?} for transaction {:?}", rto, transaction_id);
924                        events.push(StuntClientEvent::OutputPacket(transaction.packet.clone()));
925                    }
926                    None => {
927                        let protection_violated = self.mechanism.as_mut().map_or(false, |m| {
928                            m.signal_protection_violated_on_timeout(&transaction_id)
929                        });
930                        let event = if protection_violated {
931                            StuntClientEvent::TransactionFailed((
932                                transaction_id,
933                                StunTransactionError::ProtectionViolated,
934                            ))
935                        } else {
936                            StuntClientEvent::TransactionFailed((
937                                transaction_id,
938                                StunTransactionError::TimedOut,
939                            ))
940                        };
941                        info!(
942                            "Transaction {:?} timed out. Event: {:?}",
943                            transaction_id, event
944                        );
945                        events.push(event);
946                    }
947                }
948            } else {
949                warn!("Transaction {:?} not found", transaction_id);
950            }
951        }
952
953        // Add the most recent timout event if any
954        if let Some((id, left)) = self.timeouts.next_timeout(instant) {
955            events.push(StuntClientEvent::RestransmissionTimeOut((id, left)));
956        }
957    }
958
959    /// Returns the events generated by the agent.
960    /// This method should be called after any interaction with the agent.
961    /// The events notify the user about the status of the transactions.
962    /// Note that no state is maintained between interactions with the agent.
963    /// Therefore, the user should call this method to retrieve the events as
964    /// soon as an operation is completed. Otherwise, the events may be lost
965    /// if a new operation is performed.
966    pub fn events(&mut self) -> Vec<StuntClientEvent> {
967        self.transaction_events.events()
968    }
969}
970
971fn process_integrity_error(
972    error: IntegrityError,
973    transaction_id: &TransactionId,
974) -> Result<Option<StuntClientEvent>, StunAgentError> {
975    match error {
976        IntegrityError::ProtectionViolated => Ok(Some(StuntClientEvent::TransactionFailed((
977            *transaction_id,
978            StunTransactionError::ProtectionViolated,
979        )))),
980        IntegrityError::Retry => Ok(Some(StuntClientEvent::Retry(*transaction_id))),
981        IntegrityError::NotRetryable => Ok(Some(StuntClientEvent::TransactionFailed((
982            *transaction_id,
983            StunTransactionError::DoNotRetry,
984        )))),
985        IntegrityError::Discarded => {
986            // Transaction was discarded. Retransmission will continue.
987            Err(StunAgentError::Discarded)
988        }
989    }
990}
991
992fn prepare_stun_message(
993    class: StunClientMessageClass,
994    attributes: &mut StunAttributes,
995    mechanism: Option<&mut CredentialMechanismClient>,
996    use_fingerprint: bool,
997) -> Result<(), StunAgentError> {
998    if let Some(mechanism) = mechanism {
999        match class {
1000            StunClientMessageClass::Request => mechanism.prepare_request(attributes)?,
1001            StunClientMessageClass::Indication => mechanism.prepare_indication(attributes)?,
1002        }
1003    }
1004
1005    if use_fingerprint {
1006        add_fingerprint_attribute(attributes);
1007    }
1008
1009    Ok(())
1010}
1011
1012fn encode_buffer(
1013    encoder: &MessageEncoder,
1014    msg: &StunMessage,
1015    mut buffer: Vec<u8>,
1016) -> Result<StunPacket, StunEncodeError> {
1017    let size = encoder.encode(&mut buffer, msg)?;
1018    Ok(StunPacket::new(buffer, size))
1019}
1020
1021#[cfg(test)]
1022mod stun_client_tests {
1023    use super::*;
1024
1025    fn init_logging() {
1026        let _ = env_logger::builder().is_test(true).try_init();
1027    }
1028
1029    #[test]
1030    fn test_stun_client_builder() {
1031        init_logging();
1032
1033        let client =
1034            StunClienteBuilder::new(TransportReliability::Reliable(Duration::from_secs(5)))
1035                .with_max_transactions(5)
1036                .with_mechanism("user", "password", CredentialMechanism::ShortTerm(None))
1037                .with_fingerprint()
1038                .build()
1039                .expect("Could not create STUN client");
1040
1041        assert_eq!(client.max_transactions, 5);
1042        assert!(matches!(
1043            client.mechanism,
1044            Some(CredentialMechanismClient::ShortTerm(_))
1045        ));
1046        assert!(client.use_fingerprint);
1047
1048        let error = StunClienteBuilder::new(TransportReliability::Reliable(Duration::from_secs(5)))
1049            .with_max_transactions(5)
1050            .with_mechanism(
1051                "bad\u{0009}user",
1052                "password",
1053                CredentialMechanism::ShortTerm(None),
1054            )
1055            .with_fingerprint()
1056            .build()
1057            .expect_err("Should not create STUN client");
1058        assert!(matches!(error, StunAgentError::InternalError(_)));
1059
1060        let error = StunClienteBuilder::new(TransportReliability::Reliable(Duration::from_secs(5)))
1061            .with_max_transactions(5)
1062            .with_mechanism(
1063                "user",
1064                "bad\u{0009}password",
1065                CredentialMechanism::ShortTerm(None),
1066            )
1067            .with_fingerprint()
1068            .build()
1069            .expect_err("Should not create STUN client");
1070        assert!(matches!(error, StunAgentError::InternalError(_)));
1071    }
1072
1073    #[test]
1074    fn test_stun_client_constructor() {
1075        init_logging();
1076
1077        let client = StunClient::new(StunClientParameters {
1078            user_name: Some(String::from("user")),
1079            password: Some(String::from("password")),
1080            mechanism: Some(CredentialMechanism::ShortTerm(None)),
1081            reliability: TransportReliability::Reliable(Duration::from_secs(5)),
1082            fingerprint: true,
1083            max_transactions: 5,
1084        })
1085        .expect("Could not create STUN client");
1086        assert_eq!(client.max_transactions, 5);
1087        assert!(matches!(
1088            client.mechanism,
1089            Some(CredentialMechanismClient::ShortTerm(_))
1090        ));
1091        assert!(client.use_fingerprint);
1092
1093        let error = StunClient::new(StunClientParameters {
1094            user_name: Some(String::from("bad\u{0009}user")),
1095            password: Some(String::from("password")),
1096            mechanism: Some(CredentialMechanism::ShortTerm(None)),
1097            reliability: TransportReliability::Reliable(Duration::from_secs(5)),
1098            fingerprint: true,
1099            max_transactions: 5,
1100        })
1101        .expect_err("Should not create STUN client");
1102        assert!(matches!(error, StunAgentError::InternalError(_)));
1103
1104        let error = StunClient::new(StunClientParameters {
1105            user_name: Some(String::from("user")),
1106            password: Some(String::from("bad\u{0009}password")),
1107            mechanism: Some(CredentialMechanism::ShortTerm(None)),
1108            reliability: TransportReliability::Reliable(Duration::from_secs(5)),
1109            fingerprint: true,
1110            max_transactions: 5,
1111        })
1112        .expect_err("Should not create STUN client");
1113        assert!(matches!(error, StunAgentError::InternalError(_)));
1114
1115        let error = StunClient::new(StunClientParameters {
1116            user_name: None,
1117            password: Some(String::from("password")),
1118            mechanism: Some(CredentialMechanism::ShortTerm(None)),
1119            reliability: TransportReliability::Reliable(Duration::from_secs(5)),
1120            fingerprint: true,
1121            max_transactions: 5,
1122        })
1123        .expect_err("Should not create STUN client");
1124        assert!(matches!(error, StunAgentError::InternalError(_)));
1125
1126        let error = StunClient::new(StunClientParameters {
1127            user_name: Some(String::from("user")),
1128            password: None,
1129            mechanism: Some(CredentialMechanism::ShortTerm(None)),
1130            reliability: TransportReliability::Reliable(Duration::from_secs(5)),
1131            fingerprint: true,
1132            max_transactions: 5,
1133        })
1134        .expect_err("Should not create STUN client");
1135        assert!(matches!(error, StunAgentError::InternalError(_)));
1136    }
1137
1138    #[test]
1139    fn test_stun_client_transaction_finished_unknown_transaction_id() {
1140        init_logging();
1141
1142        let mut client =
1143            StunClienteBuilder::new(TransportReliability::Reliable(Duration::from_secs(5)))
1144                .with_max_transactions(5)
1145                .with_mechanism("user", "password", CredentialMechanism::ShortTerm(None))
1146                .with_fingerprint()
1147                .build()
1148                .expect("Could not create STUN client");
1149        assert_eq!(client.transactions.len(), 0);
1150
1151        let transanction_id = TransactionId::default();
1152        client.transaction_finished(&transanction_id, Instant::now());
1153
1154        assert_eq!(client.transactions.len(), 0);
1155    }
1156}