1use crate::origin;
2use crate::{
3 ALPN_14, ALPN_15, ALPN_16, ALPN_17, ALPN_18, ALPN_19, ALPN_LITE, ALPN_LITE_03, ALPN_LITE_04, ALPN_LITE_05,
4 ALPN_LITE_06_WIP, Consume, Driver, Error, NEGOTIATED, Session, Version, Versions,
5 coding::{self, Decode, Encode, Stream},
6 ietf, lite, setup, stats,
7};
8
9#[derive(Default, Clone)]
11pub struct Client {
12 publish: Option<origin::Consumer>,
13 subscribe: Option<origin::Producer>,
14 stats: stats::Session,
15 versions: Versions,
16 setup_path: Option<String>,
17 cost: Option<u64>,
18 peer_origin: Option<crate::Origin>,
19}
20
21impl Client {
22 pub fn new() -> Self {
24 Default::default()
25 }
26
27 pub fn with_publisher(mut self, publish: impl Consume<origin::Consumer>) -> Self {
31 self.publish = Some(publish.consume());
32 self
33 }
34
35 pub fn with_subscriber(mut self, subscribe: origin::Producer) -> Self {
38 self.subscribe = Some(subscribe);
39 self
40 }
41
42 pub fn with_stats(mut self, stats: stats::Session) -> Self {
47 self.stats = stats;
48 self
49 }
50
51 pub fn with_origin(self, origin: origin::Producer) -> Self {
56 self.with_publisher(&origin).with_subscriber(origin)
57 }
58
59 pub fn with_versions(mut self, versions: Versions) -> Self {
62 self.versions = versions;
63 self
64 }
65
66 pub fn with_path(mut self, path: impl Into<String>) -> Self {
78 self.setup_path = Some(path.into());
79 self
80 }
81
82 pub fn with_cost(mut self, cost: u64) -> Self {
94 self.cost = Some(cost);
95 self
96 }
97
98 pub fn with_peer_origin(mut self, origin: crate::Origin) -> Self {
117 self.peer_origin = Some(origin);
118 self
119 }
120
121 pub async fn connect<S: web_transport_trait::Session>(&self, session: S) -> Result<(Session, Driver), Error> {
125 if self.publish.is_none() && self.subscribe.is_none() {
126 tracing::warn!("not publishing or consuming anything");
127 }
128
129 let publish = self.publish.clone().map(|origin| origin.with_stats(self.stats.clone()));
134 let subscribe = self
135 .subscribe
136 .clone()
137 .map(|origin| origin.with_stats(self.stats.clone()));
138
139 let publish = match self.peer_origin {
144 Some(peer) => publish.map(|origin| origin.excluding(peer)),
145 None => publish,
146 };
147
148 let (encoding, supported) = match session.protocol() {
151 Some(ALPN_19) => {
152 let v = self
153 .versions
154 .select(Version::Ietf(ietf::Version::Draft19))
155 .ok_or(Error::Version)?;
156
157 let protocol = ietf::start(
159 session.clone(),
160 None,
161 None,
162 true,
163 publish.clone(),
164 subscribe.clone(),
165 self.peer_origin,
166 ietf::Version::Draft19,
167 self.setup_path.clone(),
168 None,
169 )?;
170
171 tracing::debug!(version = ?v, "connected");
172 return Ok(Session::new(session, v, None, protocol));
173 }
174 Some(ALPN_18) => {
175 let v = self
176 .versions
177 .select(Version::Ietf(ietf::Version::Draft18))
178 .ok_or(Error::Version)?;
179
180 let protocol = ietf::start(
183 session.clone(),
184 None,
185 None,
186 true,
187 publish.clone(),
188 subscribe.clone(),
189 self.peer_origin,
190 ietf::Version::Draft18,
191 self.setup_path.clone(),
192 None,
193 )?;
194
195 tracing::debug!(version = ?v, "connected");
196 return Ok(Session::new(session, v, None, protocol));
197 }
198 Some(ALPN_17) => {
199 let v = self
200 .versions
201 .select(Version::Ietf(ietf::Version::Draft17))
202 .ok_or(Error::Version)?;
203
204 let protocol = ietf::start(
207 session.clone(),
208 None,
209 None,
210 true,
211 publish.clone(),
212 subscribe.clone(),
213 self.peer_origin,
214 ietf::Version::Draft17,
215 self.setup_path.clone(),
216 None,
217 )?;
218
219 tracing::debug!(version = ?v, "connected");
220 return Ok(Session::new(session, v, None, protocol));
221 }
222 Some(ALPN_16) => {
223 let v = self
224 .versions
225 .select(Version::Ietf(ietf::Version::Draft16))
226 .ok_or(Error::Version)?;
227 (v, v.into())
228 }
229 Some(ALPN_15) => {
230 let v = self
231 .versions
232 .select(Version::Ietf(ietf::Version::Draft15))
233 .ok_or(Error::Version)?;
234 (v, v.into())
235 }
236 Some(ALPN_14) => {
237 let v = self
238 .versions
239 .select(Version::Ietf(ietf::Version::Draft14))
240 .ok_or(Error::Version)?;
241 (v, v.into())
242 }
243 Some(alpn @ (ALPN_LITE_05 | ALPN_LITE_06_WIP)) => {
244 let version = match alpn {
245 ALPN_LITE_06_WIP => lite::Version::Lite06Wip,
246 _ => lite::Version::Lite05,
247 };
248 self.versions.select(Version::Lite(version)).ok_or(Error::Version)?;
249
250 let our_setup = lite::Setup {
255 probe: lite::ProbeLevel::Report,
256 path: self.setup_path.clone(),
257 role: lite::Role::from_origins(self.publish.is_some(), self.subscribe.is_some()),
258 cost: self.cost,
259 origin: None,
261 };
262
263 let start = lite::start(
264 session.clone(),
265 None,
266 publish.clone(),
267 subscribe.clone(),
268 self.peer_origin,
269 version,
270 our_setup,
271 None,
272 )?;
273
274 let (session, mut driver) = Session::new(session, version.into(), start.recv_bandwidth, start.driver);
278 driver.wait_ready(|waiter| start.connecting.poll_ready(waiter)).await;
279
280 return Ok((session, driver));
281 }
282 Some(ALPN_LITE_04) => {
283 self.versions
284 .select(Version::Lite(lite::Version::Lite04))
285 .ok_or(Error::Version)?;
286
287 let start = lite::start(
288 session.clone(),
289 None,
290 publish.clone(),
291 subscribe.clone(),
292 self.peer_origin,
293 lite::Version::Lite04,
294 lite::Setup::default(),
295 None,
296 )?;
297
298 let (session, mut driver) = Session::new(
300 session,
301 lite::Version::Lite04.into(),
302 start.recv_bandwidth,
303 start.driver,
304 );
305 driver.wait_ready(|waiter| start.connecting.poll_ready(waiter)).await;
306
307 return Ok((session, driver));
308 }
309 Some(ALPN_LITE_03) => {
310 self.versions
311 .select(Version::Lite(lite::Version::Lite03))
312 .ok_or(Error::Version)?;
313
314 let start = lite::start(
316 session.clone(),
317 None,
318 publish.clone(),
319 subscribe.clone(),
320 self.peer_origin,
321 lite::Version::Lite03,
322 lite::Setup::default(),
323 None,
324 )?;
325
326 let (session, mut driver) = Session::new(
328 session,
329 lite::Version::Lite03.into(),
330 start.recv_bandwidth,
331 start.driver,
332 );
333 driver.wait_ready(|waiter| start.connecting.poll_ready(waiter)).await;
334
335 return Ok((session, driver));
336 }
337 Some(ALPN_LITE) | None => {
338 let supported = self.versions.filter(&NEGOTIATED.into()).ok_or(Error::Version)?;
339 (Version::Ietf(ietf::Version::Draft14), supported)
340 }
341 Some(p) => return Err(Error::UnknownAlpn(p.to_string())),
342 };
343
344 let mut stream = Stream::open(&session, encoding).await?;
345
346 let ietf_encoding = ietf::Version::try_from(encoding).map_err(|_| Error::Version)?;
348
349 let mut parameters = ietf::Parameters::default();
350 parameters.set_varint(ietf::ParameterVarInt::MaxRequestId, u32::MAX as u64);
351 parameters.set_bytes(ietf::ParameterBytes::Implementation, b"moq-lite-rs".to_vec());
352 if let Some(path) = &self.setup_path {
354 parameters.set_bytes(ietf::ParameterBytes::Path, path.clone().into_bytes());
355 }
356 let parameters = parameters.encode_bytes(ietf_encoding)?;
357
358 let client = setup::Client {
359 versions: supported.clone().into(),
360 parameters,
361 };
362
363 stream.writer.encode(&client).await?;
364
365 let mut server: setup::Server = stream.reader.decode().await?;
366
367 let version = supported
368 .iter()
369 .find(|v| coding::Version::from(**v) == server.version)
370 .copied()
371 .ok_or(Error::Version)?;
372
373 let (recv_bw, protocol, connecting) = match version {
374 Version::Lite(v) => {
375 let stream = stream.with_version(v);
376 let start = lite::start(
377 session.clone(),
378 Some(stream),
379 publish.clone(),
380 subscribe.clone(),
381 self.peer_origin,
382 v,
383 lite::Setup::default(),
386 None,
387 )?;
388
389 (start.recv_bandwidth, start.driver, Some(start.connecting))
390 }
391 Version::Ietf(v) => {
392 let parameters = ietf::Parameters::decode(&mut server.parameters, v)?;
394 let request_id_max = parameters
395 .get_varint(ietf::ParameterVarInt::MaxRequestId)
396 .map(ietf::RequestId);
397
398 let stream = stream.with_version(v);
399 let protocol = ietf::start(
401 session.clone(),
402 Some(stream),
403 request_id_max,
404 true,
405 publish.clone(),
406 subscribe.clone(),
407 self.peer_origin,
408 v,
409 None,
410 None,
411 )?;
412 (None, protocol, None)
413 }
414 };
415
416 let (session, mut driver) = Session::new(session, version, recv_bw, protocol);
417 if let Some(connecting) = connecting {
418 driver.wait_ready(|waiter| connecting.poll_ready(waiter)).await;
421 }
422
423 Ok((session, driver))
424 }
425}
426
427#[cfg(test)]
428mod tests {
429 use super::*;
430 use std::{
431 collections::VecDeque,
432 sync::{Arc, Mutex},
433 };
434
435 use crate::coding::{Decode, Encode};
436 use bytes::{BufMut, Bytes};
437
438 #[derive(Debug, Clone, Default)]
439 struct FakeError;
440
441 impl std::fmt::Display for FakeError {
442 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
443 write!(f, "fake transport error")
444 }
445 }
446
447 impl std::error::Error for FakeError {}
448
449 impl web_transport_trait::Error for FakeError {
450 fn session_error(&self) -> Option<(u32, String)> {
451 Some((0, "closed".to_string()))
452 }
453 }
454
455 #[derive(Clone, Default)]
456 struct FakeSession {
457 state: Arc<FakeSessionState>,
458 }
459
460 #[derive(Default)]
461 struct FakeSessionState {
462 protocol: Option<&'static str>,
463 control_stream: Mutex<Option<(FakeSendStream, FakeRecvStream)>>,
464 close_events: Mutex<Vec<(u32, String)>>,
465 close_notify: tokio::sync::Notify,
466 control_writes: Arc<Mutex<Vec<u8>>>,
467 send_rate: Mutex<Option<u64>>,
468 }
469
470 impl FakeSession {
471 fn new(protocol: Option<&'static str>, server_control_bytes: Vec<u8>) -> Self {
472 let writes = Arc::new(Mutex::new(Vec::new()));
473 let send = FakeSendStream { writes: writes.clone() };
474 let recv = FakeRecvStream {
475 data: VecDeque::from(server_control_bytes),
476 };
477 let state = FakeSessionState {
478 protocol,
479 control_stream: Mutex::new(Some((send, recv))),
480 close_events: Mutex::new(Vec::new()),
481 close_notify: tokio::sync::Notify::new(),
482 control_writes: writes,
483 send_rate: Mutex::new(None),
484 };
485 Self { state: Arc::new(state) }
486 }
487
488 fn set_send_rate(&self, rate: Option<u64>) {
489 *self.state.send_rate.lock().unwrap() = rate;
490 }
491
492 fn control_writes(&self) -> Vec<u8> {
493 self.state.control_writes.lock().unwrap().clone()
494 }
495
496 async fn wait_for_first_close(&self) -> (u32, String) {
497 loop {
498 let notified = self.state.close_notify.notified();
499 if let Some(close) = self.state.close_events.lock().unwrap().first().cloned() {
500 return close;
501 }
502 notified.await;
503 }
504 }
505 }
506
507 impl web_transport_trait::Session for FakeSession {
508 type SendStream = FakeSendStream;
509 type RecvStream = FakeRecvStream;
510 type Error = FakeError;
511
512 async fn accept_uni(&self) -> Result<Self::RecvStream, Self::Error> {
513 std::future::pending().await
514 }
515
516 async fn accept_bi(&self) -> Result<(Self::SendStream, Self::RecvStream), Self::Error> {
517 std::future::pending().await
518 }
519
520 async fn open_bi(&self) -> Result<(Self::SendStream, Self::RecvStream), Self::Error> {
521 self.state.control_stream.lock().unwrap().take().ok_or(FakeError)
522 }
523
524 async fn open_uni(&self) -> Result<Self::SendStream, Self::Error> {
525 std::future::pending().await
526 }
527
528 fn send_datagram(&self, _payload: Bytes) -> Result<(), Self::Error> {
529 Ok(())
530 }
531
532 async fn recv_datagram(&self) -> Result<Bytes, Self::Error> {
533 std::future::pending().await
534 }
535
536 fn max_datagram_size(&self) -> usize {
537 1200
538 }
539
540 fn protocol(&self) -> Option<&str> {
541 self.state.protocol
542 }
543
544 fn close(&self, code: u32, reason: &str) {
545 self.state.close_events.lock().unwrap().push((code, reason.to_string()));
546 self.state.close_notify.notify_waiters();
547 }
548
549 async fn closed(&self) -> Self::Error {
550 loop {
551 let notified = self.state.close_notify.notified();
552 if !self.state.close_events.lock().unwrap().is_empty() {
553 return FakeError;
554 }
555 notified.await;
556 }
557 }
558
559 fn stats(&self) -> impl web_transport_trait::Stats {
560 FakeStats {
561 send_rate: *self.state.send_rate.lock().unwrap(),
562 }
563 }
564 }
565
566 struct FakeStats {
567 send_rate: Option<u64>,
568 }
569
570 impl web_transport_trait::Stats for FakeStats {
571 fn estimated_send_rate(&self) -> Option<u64> {
572 self.send_rate
573 }
574 }
575
576 #[derive(Clone, Default)]
577 struct FakeSendStream {
578 writes: Arc<Mutex<Vec<u8>>>,
579 }
580
581 impl web_transport_trait::SendStream for FakeSendStream {
582 type Error = FakeError;
583
584 async fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
585 self.writes.lock().unwrap().put_slice(buf);
586 Ok(buf.len())
587 }
588
589 fn set_priority(&mut self, _order: u8) {}
590
591 fn finish(&mut self) -> Result<(), Self::Error> {
592 Ok(())
593 }
594
595 fn reset(&mut self, _code: u32) {}
596
597 async fn closed(&mut self) -> Result<(), Self::Error> {
598 Ok(())
599 }
600 }
601
602 struct FakeRecvStream {
603 data: VecDeque<u8>,
604 }
605
606 impl web_transport_trait::RecvStream for FakeRecvStream {
607 type Error = FakeError;
608
609 async fn read(&mut self, dst: &mut [u8]) -> Result<Option<usize>, Self::Error> {
610 if self.data.is_empty() {
611 return Ok(None);
612 }
613
614 let size = dst.len().min(self.data.len());
615 for slot in dst.iter_mut().take(size) {
616 *slot = self.data.pop_front().unwrap();
617 }
618 Ok(Some(size))
619 }
620
621 fn stop(&mut self, _code: u32) {}
622
623 async fn closed(&mut self) -> Result<(), Self::Error> {
624 Ok(())
625 }
626 }
627
628 fn mock_server_setup(negotiated: Version) -> Vec<u8> {
629 let mut encoded = Vec::new();
630 let server = setup::Server {
631 version: negotiated.into(),
632 parameters: Bytes::new(),
633 };
634 server
635 .encode(&mut encoded, Version::Ietf(ietf::Version::Draft14))
636 .unwrap();
637
638 let info = lite::SessionInfo { bitrate: Some(1) };
640 let lite_v = lite::Version::try_from(negotiated).unwrap();
641 info.encode(&mut encoded, lite_v).unwrap();
642
643 encoded
644 }
645
646 async fn run_alpn_lite_fallback_case(protocol: Option<&'static str>) {
647 let fake = FakeSession::new(protocol, mock_server_setup(Version::Lite(lite::Version::Lite01)));
648 let client = Client::new().with_versions(
649 [
650 Version::Lite(lite::Version::Lite03),
651 Version::Lite(lite::Version::Lite02),
652 Version::Lite(lite::Version::Lite01),
653 Version::Ietf(ietf::Version::Draft14),
654 ]
655 .into(),
656 );
657
658 let _connection = client.connect(fake.clone()).await.unwrap();
659
660 let mut setup_bytes = Bytes::from(fake.control_writes());
662 let setup = setup::Client::decode(&mut setup_bytes, Version::Ietf(ietf::Version::Draft14)).unwrap();
663 let advertised: Vec<Version> = setup.versions.iter().map(|v| Version::try_from(*v).unwrap()).collect();
664 assert_eq!(
665 advertised,
666 vec![
667 Version::Lite(lite::Version::Lite02),
668 Version::Lite(lite::Version::Lite01),
669 Version::Ietf(ietf::Version::Draft14),
670 ]
671 );
672
673 let (code, _) = fake.wait_for_first_close().await;
680 assert_ne!(code, Error::Version.to_code(), "SessionInfo failed to decode");
681 }
682
683 #[tokio::test(start_paused = true)]
684 async fn alpn_lite_falls_back_to_draft14_and_switches_version_post_setup() {
685 run_alpn_lite_fallback_case(Some(ALPN_LITE)).await;
686 }
687
688 #[tokio::test(start_paused = true)]
689 async fn no_alpn_falls_back_to_draft14_and_switches_version_post_setup() {
690 run_alpn_lite_fallback_case(None).await;
691 }
692
693 #[test]
701 fn driver_is_caller_polled_and_holds_no_session() {
702 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
703 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
704
705 let (session, mut driver) = futures::executor::block_on(client.connect(fake.clone())).unwrap();
706 assert_eq!(session.version(), Version::Lite(lite::Version::Lite04));
707
708 assert!(driver.poll(&kio::Waiter::noop()).is_pending());
710
711 let mut context = std::task::Context::from_waker(std::task::Waker::noop());
713 assert!(std::future::Future::poll(std::pin::Pin::new(&mut driver), &mut context).is_pending());
714
715 drop(session);
718 assert_eq!(fake.state.close_events.lock().unwrap()[0].0, Error::Cancel.to_code());
719 }
720
721 #[test]
724 fn session_clones_share_the_close() {
725 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
726 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
727
728 let (session, _driver) = futures::executor::block_on(client.connect(fake.clone())).unwrap();
729 let clone = session.clone();
730
731 drop(session);
733 assert!(fake.state.close_events.lock().unwrap().is_empty());
734
735 clone.abort(Error::Cancel);
736 assert_eq!(fake.state.close_events.lock().unwrap()[0].0, Error::Cancel.to_code());
737
738 drop(clone);
740 assert_eq!(fake.state.close_events.lock().unwrap().len(), 1);
741 }
742
743 #[tokio::test(start_paused = true)]
747 async fn send_bandwidth_samples_while_the_driver_runs() {
748 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
749 fake.set_send_rate(Some(1_000_000));
750
751 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
752 let (session, driver) = client.connect(fake.clone()).await.unwrap();
753 tokio::spawn(driver);
754
755 let mut bandwidth = session.send_bandwidth().expect("backend reports an estimate");
756 assert_eq!(bandwidth.changed().await.unwrap(), Some(1_000_000));
757
758 fake.set_send_rate(Some(2_000_000));
760 assert_eq!(bandwidth.changed().await.unwrap(), Some(2_000_000));
761 }
762}