1use 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
23pub struct TaggedMessage {
30 pub now: Instant,
32 pub message: Message,
34}
35
36#[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)]
76pub struct ClientBuilder {
79 settings: ClientSettings,
80}
81
82impl ClientBuilder {
83 pub fn with_rto(mut self, rto: Duration) -> Self {
85 self.settings.rto = rto;
86 self
87 }
88
89 pub fn with_timeout_rate(mut self, d: Duration) -> Self {
91 self.settings.rto_rate = d;
92 self
93 }
94
95 pub fn with_buffer_size(mut self, buffer_size: usize) -> Self {
97 self.settings.buffer_size = buffer_size;
98 self
99 }
100
101 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 pub fn new() -> Self {
115 ClientBuilder {
116 settings: ClientSettings::default(),
117 }
118 }
119
120 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
136pub 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 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 fn observe(&mut self, now: Instant) {
179 self.now = now.max(self.now);
180 }
181
182 pub fn local_addr(&self) -> SocketAddr {
184 self.local
185 }
186
187 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 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 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 self.transactions.entry(ct.id).or_insert(ct);
283
284 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 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 #[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 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 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 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 assert_eq!(client.poll_timeout(), Some(t(311)));
403
404 client.close()
405 }
406}