Skip to main content

rtc_stun/
client.rs

1//! A Sans-I/O STUN client.
2//!
3//! Sends Binding requests and matches the responses, applying the retransmission schedule the RFC
4//! specifies (an initial RTO, doubling per retry) so a lost request on a UDP path is retried
5//! rather than lost. Build one with [`ClientBuilder`](crate::client::ClientBuilder); drive it with datagrams and time.
6use bytes::BytesMut;
7use shared::error::*;
8use std::collections::{HashMap, VecDeque};
9use std::io::BufReader;
10use std::net::SocketAddr;
11use std::ops::Add;
12use std::time::{Duration, Instant};
13
14use crate::agent::*;
15use crate::message::*;
16use shared::{TaggedBytesMut, TransportContext, TransportMessage, TransportProtocol};
17
18const DEFAULT_TIMEOUT_RATE: Duration = Duration::from_millis(5);
19const DEFAULT_RTO: Duration = Duration::from_millis(300);
20const DEFAULT_MAX_ATTEMPTS: u32 = 7;
21const DEFAULT_MAX_BUFFER_SIZE: usize = 8;
22
23/// A [`Message`] together with the instant the caller is sending it at.
24///
25/// `Rin` is a [`TaggedBytesMut`], which carries a timestamp; without this the write channel
26/// would not, and `handle_write` would have to ask the clock for one. A STUN client's first
27/// action is typically a write before any read, so a retained instant is not an option here —
28/// it would still be the construction seed.
29pub struct TaggedMessage {
30    /// When the caller is sending this message.
31    pub now: Instant,
32    /// The STUN message to send.
33    pub message: Message,
34}
35
36/// ClientTransaction represents transaction in progress.
37/// If transaction is succeed or failed, f will be called
38/// provided by event.
39/// Concurrent access is invalid.
40#[derive(Debug, Clone)]
41pub struct ClientTransaction {
42    id: TransactionId,
43    attempt: u32,
44    start: Instant,
45    rto: Duration,
46    raw: Vec<u8>,
47}
48
49impl ClientTransaction {
50    pub(crate) fn next_timeout(&self, now: Instant) -> Instant {
51        now.add((self.attempt + 1) * self.rto)
52    }
53}
54
55struct ClientSettings {
56    buffer_size: usize,
57    rto: Duration,
58    rto_rate: Duration,
59    max_attempts: u32,
60    closed: bool,
61}
62
63impl Default for ClientSettings {
64    fn default() -> Self {
65        ClientSettings {
66            buffer_size: DEFAULT_MAX_BUFFER_SIZE,
67            rto: DEFAULT_RTO,
68            rto_rate: DEFAULT_TIMEOUT_RATE,
69            max_attempts: DEFAULT_MAX_ATTEMPTS,
70            closed: false,
71        }
72    }
73}
74
75#[derive(Default)]
76/// Builds a [`Client`] with a chosen transaction timeout, retransmission schedule and
77/// handler.
78pub struct ClientBuilder {
79    settings: ClientSettings,
80}
81
82impl ClientBuilder {
83    /// with_rto sets client RTO as defined in STUN RFC.
84    pub fn with_rto(mut self, rto: Duration) -> Self {
85        self.settings.rto = rto;
86        self
87    }
88
89    /// with_timeout_rate sets RTO timer minimum resolution.
90    pub fn with_timeout_rate(mut self, d: Duration) -> Self {
91        self.settings.rto_rate = d;
92        self
93    }
94
95    /// with_buffer_size sets buffer size.
96    pub fn with_buffer_size(mut self, buffer_size: usize) -> Self {
97        self.settings.buffer_size = buffer_size;
98        self
99    }
100
101    /// with_no_retransmit disables retransmissions and sets RTO to
102    /// DEFAULT_MAX_ATTEMPTS * DEFAULT_RTO which will be effectively time out
103    /// if not set.
104    /// Useful for TCP connections where transport handles RTO.
105    pub fn with_no_retransmit(mut self) -> Self {
106        self.settings.max_attempts = 0;
107        if self.settings.rto == Duration::from_secs(0) {
108            self.settings.rto = DEFAULT_MAX_ATTEMPTS * DEFAULT_RTO;
109        }
110        self
111    }
112
113    /// A builder with the RFC's default timings.
114    pub fn new() -> Self {
115        ClientBuilder {
116            settings: ClientSettings::default(),
117        }
118    }
119
120    /// Builds the client for the given local and remote addresses.
121    ///
122    /// # Errors
123    ///
124    /// Fails if the configured timings are inconsistent.
125    pub fn build(
126        self,
127        now: Instant,
128        local: SocketAddr,
129        remote: SocketAddr,
130        protocol: TransportProtocol,
131    ) -> Result<Client> {
132        Ok(Client::new(now, local, remote, protocol, self.settings))
133    }
134}
135
136/// Client simulates "connection" to STUN server.
137pub struct Client {
138    local: SocketAddr,
139    remote: SocketAddr,
140    transport_protocol: TransportProtocol,
141    agent: Agent,
142    settings: ClientSettings,
143    transactions: HashMap<TransactionId, ClientTransaction>,
144    transmits: VecDeque<TransportMessage<BytesMut>>,
145
146    /// The newest instant a caller has supplied, seeded at construction.
147    ///
148    /// `poll_event` schedules retransmissions, which needs a deadline, but a poll is a drain
149    /// and receives no instant. The retransmission is caused by the timeout the caller reported
150    /// through `handle_timeout`, so that instant is the right one to schedule against.
151    now: Instant,
152}
153
154impl Client {
155    fn new(
156        now: Instant,
157        local: SocketAddr,
158        remote: SocketAddr,
159        transport_protocol: TransportProtocol,
160        settings: ClientSettings,
161    ) -> Self {
162        Self {
163            local,
164            remote,
165            transport_protocol,
166            agent: Agent::new(),
167            settings,
168            transactions: HashMap::new(),
169            transmits: VecDeque::new(),
170            now,
171        }
172    }
173
174    /// Records the newest instant a caller has supplied.
175    ///
176    /// `max` rather than assignment: an outbound message carries the instant of the input that
177    /// caused it, so a caller can legitimately present an older one than the newest seen.
178    fn observe(&mut self, now: Instant) {
179        self.now = now.max(self.now);
180    }
181
182    /// The address this client sends from.
183    pub fn local_addr(&self) -> SocketAddr {
184        self.local
185    }
186
187    /// The STUN server this client talks to.
188    pub fn peer_addr(&self) -> SocketAddr {
189        self.remote
190    }
191}
192
193impl sansio::Protocol<TaggedBytesMut, TaggedMessage, ()> for Client {
194    type Rout = ();
195    type Wout = TaggedBytesMut;
196    type Eout = StunEvent;
197    type Error = Error;
198    type Time = Instant;
199
200    fn handle_read(&mut self, msg: TaggedBytesMut) -> Result<()> {
201        self.observe(msg.now);
202        let mut stun_msg = Message::new();
203        let mut reader = BufReader::new(&msg.message[..]);
204        stun_msg.read_from(&mut reader)?;
205        self.agent.handle_event(ClientAgent::Process(stun_msg))
206    }
207
208    fn poll_read(&mut self) -> Option<Self::Rout> {
209        None
210    }
211
212    fn handle_write(&mut self, msg: TaggedMessage) -> Result<()> {
213        if self.settings.closed {
214            return Err(Error::ErrClientClosed);
215        }
216
217        let now = msg.now;
218        self.observe(now);
219        let m = msg.message;
220        let payload = BytesMut::from(&m.raw[..]);
221
222        let ct = ClientTransaction {
223            id: m.transaction_id,
224            attempt: 0,
225            start: now,
226            rto: self.settings.rto,
227            raw: m.raw,
228        };
229        let deadline = ct.next_timeout(ct.start);
230        self.transactions.entry(ct.id).or_insert(ct);
231        self.agent
232            .handle_event(ClientAgent::Start(m.transaction_id, deadline))?;
233
234        self.transmits.push_back(TransportMessage {
235            now,
236            transport: TransportContext {
237                local_addr: self.local,
238                peer_addr: self.remote,
239                ecn: None,
240                transport_protocol: self.transport_protocol,
241            },
242            message: payload,
243        });
244
245        Ok(())
246    }
247
248    /// Returns packets to transmit
249    ///
250    /// It should be polled for transmit after:
251    /// - the application performed some I/O
252    /// - a call was made to `handle_read`
253    /// - a call was made to `handle_write`
254    /// - a call was made to `handle_timeout`
255    fn poll_write(&mut self) -> Option<Self::Wout> {
256        self.transmits.pop_front()
257    }
258
259    fn poll_event(&mut self) -> Option<Self::Eout> {
260        while let Some(event) = self.agent.poll_event() {
261            let mut ct = if self.transactions.contains_key(&event.id) {
262                self.transactions.remove(&event.id).unwrap()
263            } else {
264                continue;
265            };
266
267            if let StunEvent::Message(_) = &event.evt {
268                return Some(event.evt);
269            }
270            if ct.attempt >= self.settings.max_attempts {
271                return Some(event.evt);
272            }
273
274            // Doing re-transmission.
275            ct.attempt += 1;
276
277            let payload = BytesMut::from(&ct.raw[..]);
278            let timeout = ct.next_timeout(self.now);
279            let id = ct.id;
280
281            // Starting client transaction.
282            self.transactions.entry(ct.id).or_insert(ct);
283
284            // Starting agent transaction.
285            if self
286                .agent
287                .handle_event(ClientAgent::Start(id, timeout))
288                .is_err()
289            {
290                self.transactions.remove(&id);
291                return Some(event.evt);
292            }
293
294            // Writing message to connection again.
295            self.transmits.push_back(TransportMessage {
296                now: self.now,
297                transport: TransportContext {
298                    local_addr: self.local,
299                    peer_addr: self.remote,
300                    ecn: None,
301                    transport_protocol: self.transport_protocol,
302                },
303                message: payload,
304            });
305        }
306
307        None
308    }
309
310    fn poll_timeout(&mut self) -> Option<Self::Time> {
311        self.agent.poll_timeout()
312    }
313
314    fn handle_timeout(&mut self, now: Instant) -> Result<()> {
315        self.observe(now);
316        self.agent.handle_event(ClientAgent::Collect(now))
317    }
318
319    fn close(&mut self) -> Result<()> {
320        if self.settings.closed {
321            return Err(Error::ErrClientClosed);
322        }
323        self.settings.closed = true;
324        self.agent.handle_event(ClientAgent::Close)
325    }
326}
327
328#[cfg(test)]
329mod client_test {
330    use super::*;
331    use sansio::Protocol;
332
333    fn addrs() -> (SocketAddr, SocketAddr) {
334        (
335            "127.0.0.1:5000".parse().unwrap(),
336            "127.0.0.1:3478".parse().unwrap(),
337        )
338    }
339
340    fn binding_request() -> Message {
341        let mut msg = Message::new();
342        msg.build(&[Box::<TransactionId>::default(), Box::new(BINDING_REQUEST)])
343            .expect("a binding request encodes");
344        msg
345    }
346
347    /// The transaction deadline is computed from the instant the caller supplied to
348    /// `handle_write`, not from an ambient reading, so a retransmission can be observed by
349    /// arithmetic on a base instant with no wall-clock time passing and no sleeping.
350    #[test]
351    fn test_transaction_retransmits_on_injected_time() -> Result<()> {
352        let base = Instant::now();
353        let t = |millis| base + Duration::from_millis(millis);
354
355        let (local, remote) = addrs();
356        let mut client = ClientBuilder::new()
357            .with_rto(Duration::from_millis(100))
358            .build(t(0), local, remote, TransportProtocol::UDP)?;
359
360        // The write is stamped at t(10), so the first attempt's deadline is t(10) + 1 * rto.
361        client.handle_write(TaggedMessage {
362            now: t(10),
363            message: binding_request(),
364        })?;
365
366        let transmit = client.poll_write().expect("the request is queued");
367        assert_eq!(
368            transmit.now,
369            t(10),
370            "the request carries the caller's instant, not an ambient reading"
371        );
372        assert!(client.poll_write().is_none());
373
374        assert_eq!(
375            client.poll_timeout(),
376            Some(t(110)),
377            "the deadline is one RTO after the instant the caller wrote at"
378        );
379
380        // Before the deadline nothing is retransmitted. Note the agent collects on
381        // `deadline < now`, *strictly* — so arriving exactly at the deadline is not yet
382        // late. A virtual clock advanced by exactly one RTO therefore needs one more tick,
383        // which is worth knowing before writing a `clock.advance(rto)` test against it.
384        client.handle_timeout(t(110))?;
385        while client.poll_event().is_some() {}
386        assert!(
387            client.poll_write().is_none(),
388            "the deadline is not yet past at exactly the deadline"
389        );
390
391        // Past the deadline the request goes out again, stamped with that same instant.
392        client.handle_timeout(t(111))?;
393        while client.poll_event().is_some() {}
394        let retransmit = client.poll_write().expect("the request is retransmitted");
395        assert_eq!(retransmit.now, t(111));
396        assert_eq!(
397            retransmit.message, transmit.message,
398            "a retransmission repeats the original request verbatim"
399        );
400
401        // The second attempt backs off to two RTOs from the instant it was scheduled at.
402        assert_eq!(client.poll_timeout(), Some(t(311)));
403
404        client.close()
405    }
406}