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 {
97 self.cost = Some(cost);
98 self
99 }
100
101 pub fn with_peer_origin(mut self, origin: crate::Origin) -> Self {
121 self.peer_origin = Some(origin);
122 self
123 }
124
125 pub async fn connect<S: web_transport_trait::Session>(&self, session: S) -> Result<(Session, Driver), Error> {
129 if self.publish.is_none() && self.subscribe.is_none() {
130 tracing::warn!("not publishing or consuming anything");
131 }
132
133 let publish = self.publish.clone().map(|origin| origin.with_stats(self.stats.clone()));
138 let subscribe = self
139 .subscribe
140 .clone()
141 .map(|origin| origin.with_stats(self.stats.clone()));
142
143 let publish = match self.peer_origin {
148 Some(peer) => publish.map(|origin| origin.excluding(peer)),
149 None => publish,
150 };
151
152 let (encoding, supported) = match session.protocol() {
155 Some(ALPN_19) => {
156 let v = self
157 .versions
158 .select(Version::Ietf(ietf::Version::Draft19))
159 .ok_or(Error::Version)?;
160
161 let protocol = ietf::start(ietf::Config {
163 session: session.clone(),
164 setup: None,
165 request_id_max: None,
166 client: true,
167 publish: publish.clone(),
168 subscribe: subscribe.clone(),
169 peer_origin: self.peer_origin,
170 cost: self.cost,
171 version: ietf::Version::Draft19,
172 path: self.setup_path.clone(),
173 peer_setup_stream: None,
174 peer_declared: None,
175 })?;
176
177 tracing::debug!(version = ?v, "connected");
178 return Ok(Session::new(session, v, None, protocol));
179 }
180 Some(ALPN_18) => {
181 let v = self
182 .versions
183 .select(Version::Ietf(ietf::Version::Draft18))
184 .ok_or(Error::Version)?;
185
186 let protocol = ietf::start(ietf::Config {
189 session: session.clone(),
190 setup: None,
191 request_id_max: None,
192 client: true,
193 publish: publish.clone(),
194 subscribe: subscribe.clone(),
195 peer_origin: self.peer_origin,
196 cost: self.cost,
197 version: ietf::Version::Draft18,
198 path: self.setup_path.clone(),
199 peer_setup_stream: None,
200 peer_declared: None,
201 })?;
202
203 tracing::debug!(version = ?v, "connected");
204 return Ok(Session::new(session, v, None, protocol));
205 }
206 Some(ALPN_17) => {
207 let v = self
208 .versions
209 .select(Version::Ietf(ietf::Version::Draft17))
210 .ok_or(Error::Version)?;
211
212 let protocol = ietf::start(ietf::Config {
215 session: session.clone(),
216 setup: None,
217 request_id_max: None,
218 client: true,
219 publish: publish.clone(),
220 subscribe: subscribe.clone(),
221 peer_origin: self.peer_origin,
222 cost: self.cost,
223 version: ietf::Version::Draft17,
224 path: self.setup_path.clone(),
225 peer_setup_stream: None,
226 peer_declared: None,
227 })?;
228
229 tracing::debug!(version = ?v, "connected");
230 return Ok(Session::new(session, v, None, protocol));
231 }
232 Some(ALPN_16) => {
233 let v = self
234 .versions
235 .select(Version::Ietf(ietf::Version::Draft16))
236 .ok_or(Error::Version)?;
237 (v, v.into())
238 }
239 Some(ALPN_15) => {
240 let v = self
241 .versions
242 .select(Version::Ietf(ietf::Version::Draft15))
243 .ok_or(Error::Version)?;
244 (v, v.into())
245 }
246 Some(ALPN_14) => {
247 let v = self
248 .versions
249 .select(Version::Ietf(ietf::Version::Draft14))
250 .ok_or(Error::Version)?;
251 (v, v.into())
252 }
253 Some(alpn @ (ALPN_LITE_05 | ALPN_LITE_06_WIP)) => {
254 let version = match alpn {
255 ALPN_LITE_06_WIP => lite::Version::Lite06Wip,
256 _ => lite::Version::Lite05,
257 };
258 self.versions.select(Version::Lite(version)).ok_or(Error::Version)?;
259
260 let our_setup = lite::Setup {
266 probe: lite::ProbeLevel::detect(&session),
267 path: self.setup_path.clone(),
268 role: lite::Role::from_origins(self.publish.is_some(), self.subscribe.is_some()),
269 cost: self.cost,
270 origin: None,
272 };
273
274 let start = lite::start(lite::Config {
275 session: session.clone(),
276 setup_stream: None,
277 publish: publish.clone(),
278 subscribe: subscribe.clone(),
279 peer_origin: self.peer_origin,
280 version,
281 our_setup,
282 peer_setup: None,
283 })?;
284
285 return Ok(Session::new(
286 session,
287 version.into(),
288 start.recv_bandwidth,
289 start.driver,
290 ));
291 }
292 Some(ALPN_LITE_04) => {
293 self.versions
294 .select(Version::Lite(lite::Version::Lite04))
295 .ok_or(Error::Version)?;
296
297 let start = lite::start(lite::Config {
298 session: session.clone(),
299 setup_stream: None,
300 publish: publish.clone(),
301 subscribe: subscribe.clone(),
302 peer_origin: self.peer_origin,
303 version: lite::Version::Lite04,
304 our_setup: lite::Setup::default(),
305 peer_setup: None,
306 })?;
307
308 return Ok(Session::new(
309 session,
310 lite::Version::Lite04.into(),
311 start.recv_bandwidth,
312 start.driver,
313 ));
314 }
315 Some(ALPN_LITE_03) => {
316 self.versions
317 .select(Version::Lite(lite::Version::Lite03))
318 .ok_or(Error::Version)?;
319
320 let start = lite::start(lite::Config {
322 session: session.clone(),
323 setup_stream: None,
324 publish: publish.clone(),
325 subscribe: subscribe.clone(),
326 peer_origin: self.peer_origin,
327 version: lite::Version::Lite03,
328 our_setup: lite::Setup::default(),
329 peer_setup: None,
330 })?;
331
332 return Ok(Session::new(
333 session,
334 lite::Version::Lite03.into(),
335 start.recv_bandwidth,
336 start.driver,
337 ));
338 }
339 Some(ALPN_LITE) | None => {
340 let supported = self.versions.filter(&NEGOTIATED.into()).ok_or(Error::Version)?;
341 (Version::Ietf(ietf::Version::Draft14), supported)
342 }
343 Some(p) => return Err(Error::UnknownAlpn(p.to_string())),
344 };
345
346 let mut stream = Stream::open(&session, encoding).await?;
347
348 let ietf_encoding = ietf::Version::try_from(encoding).map_err(|_| Error::Version)?;
350
351 let mut parameters = ietf::Parameters::default();
352 parameters.set_varint(ietf::ParameterVarInt::MaxRequestId, u32::MAX as u64);
353 parameters.set_bytes(ietf::ParameterBytes::Implementation, b"moq-lite-rs".to_vec());
354 if let Some(path) = &self.setup_path {
356 parameters.set_bytes(ietf::ParameterBytes::Path, path.clone().into_bytes());
357 }
358 ietf::solicit::into_setup(&mut parameters, ietf_encoding);
359 let parameters = parameters.encode_bytes(ietf_encoding)?;
360
361 let client = setup::Client {
362 versions: supported.clone().into(),
363 parameters,
364 };
365
366 stream.writer.encode(&client).await?;
367
368 let mut server: setup::Server = stream.reader.decode().await?;
369
370 let version = supported
371 .iter()
372 .find(|v| coding::Version::from(**v) == server.version)
373 .copied()
374 .ok_or(Error::Version)?;
375
376 let (recv_bw, protocol) = match version {
377 Version::Lite(v) => {
378 let stream = stream.with_version(v);
379 let start = lite::start(lite::Config {
380 session: session.clone(),
381 setup_stream: Some(stream),
382 publish: publish.clone(),
383 subscribe: subscribe.clone(),
384 peer_origin: self.peer_origin,
385 version: v,
386 our_setup: lite::Setup::default(),
389 peer_setup: None,
390 })?;
391
392 (start.recv_bandwidth, start.driver)
393 }
394 Version::Ietf(v) => {
395 let parameters = ietf::Parameters::decode(&mut server.parameters, v)?;
398 let request_id_max = parameters
399 .get_varint(ietf::ParameterVarInt::MaxRequestId)
400 .map(ietf::RequestId);
401 let peer_declared = ietf::peer::Peer {
402 solicit: ietf::solicit::from_setup(¶meters, v)?,
403 ..Default::default()
404 };
405
406 let stream = stream.with_version(v);
407 let protocol = ietf::start(ietf::Config {
409 session: session.clone(),
410 setup: Some(stream),
411 request_id_max,
412 client: true,
413 publish: publish.clone(),
414 subscribe: subscribe.clone(),
415 peer_origin: self.peer_origin,
416 cost: self.cost,
417 version: v,
418 path: None,
419 peer_setup_stream: None,
420 peer_declared: Some(peer_declared),
421 })?;
422 (None, protocol)
423 }
424 };
425
426 Ok(Session::new(session, version, recv_bw, protocol))
427 }
428}
429
430#[cfg(test)]
431mod tests {
432 use super::*;
433 use std::{
434 collections::VecDeque,
435 sync::{Arc, Mutex},
436 };
437
438 use crate::coding::{Decode, Encode};
439 use bytes::{BufMut, Bytes};
440
441 #[derive(Debug, Clone, Default)]
442 struct FakeError;
443
444 impl std::fmt::Display for FakeError {
445 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
446 write!(f, "fake transport error")
447 }
448 }
449
450 impl std::error::Error for FakeError {}
451
452 impl web_transport_trait::Error for FakeError {
453 fn session_error(&self) -> Option<(u32, String)> {
454 Some((0, "closed".to_string()))
455 }
456 }
457
458 #[derive(Clone, Default)]
459 struct FakeSession {
460 state: Arc<FakeSessionState>,
461 }
462
463 #[derive(Default)]
464 struct FakeSessionState {
465 protocol: Option<&'static str>,
466 control_stream: Mutex<Option<(FakeSendStream, FakeRecvStream)>>,
467 close_events: Mutex<Vec<(u32, String)>>,
468 close_notify: tokio::sync::Notify,
469 control_writes: Arc<Mutex<Vec<u8>>>,
470 send_rate: Mutex<Option<u64>>,
471 }
472
473 impl FakeSession {
474 fn new(protocol: Option<&'static str>, server_control_bytes: Vec<u8>) -> Self {
475 let writes = Arc::new(Mutex::new(Vec::new()));
476 let send = FakeSendStream { writes: writes.clone() };
477 let recv = FakeRecvStream {
478 data: VecDeque::from(server_control_bytes),
479 };
480 let state = FakeSessionState {
481 protocol,
482 control_stream: Mutex::new(Some((send, recv))),
483 close_events: Mutex::new(Vec::new()),
484 close_notify: tokio::sync::Notify::new(),
485 control_writes: writes,
486 send_rate: Mutex::new(None),
487 };
488 Self { state: Arc::new(state) }
489 }
490
491 fn set_send_rate(&self, rate: Option<u64>) {
492 *self.state.send_rate.lock().unwrap() = rate;
493 }
494
495 fn control_writes(&self) -> Vec<u8> {
496 self.state.control_writes.lock().unwrap().clone()
497 }
498
499 async fn wait_for_first_close(&self) -> (u32, String) {
500 loop {
501 let notified = self.state.close_notify.notified();
502 if let Some(close) = self.state.close_events.lock().unwrap().first().cloned() {
503 return close;
504 }
505 notified.await;
506 }
507 }
508 }
509
510 impl web_transport_trait::Session for FakeSession {
511 type SendStream = FakeSendStream;
512 type RecvStream = FakeRecvStream;
513 type Error = FakeError;
514
515 async fn accept_uni(&self) -> Result<Self::RecvStream, Self::Error> {
516 std::future::pending().await
517 }
518
519 async fn accept_bi(&self) -> Result<(Self::SendStream, Self::RecvStream), Self::Error> {
520 std::future::pending().await
521 }
522
523 async fn open_bi(&self) -> Result<(Self::SendStream, Self::RecvStream), Self::Error> {
524 self.state.control_stream.lock().unwrap().take().ok_or(FakeError)
525 }
526
527 async fn open_uni(&self) -> Result<Self::SendStream, Self::Error> {
528 std::future::pending().await
529 }
530
531 fn send_datagram(&self, _payload: Bytes) -> Result<(), Self::Error> {
532 Ok(())
533 }
534
535 async fn recv_datagram(&self) -> Result<Bytes, Self::Error> {
536 std::future::pending().await
537 }
538
539 fn max_datagram_size(&self) -> usize {
540 1200
541 }
542
543 fn protocol(&self) -> Option<&str> {
544 self.state.protocol
545 }
546
547 fn close(&self, code: u32, reason: &str) {
548 self.state.close_events.lock().unwrap().push((code, reason.to_string()));
549 self.state.close_notify.notify_waiters();
550 }
551
552 async fn closed(&self) -> Self::Error {
553 loop {
554 let notified = self.state.close_notify.notified();
555 if !self.state.close_events.lock().unwrap().is_empty() {
556 return FakeError;
557 }
558 notified.await;
559 }
560 }
561
562 fn stats(&self) -> impl web_transport_trait::Stats {
563 FakeStats {
564 send_rate: *self.state.send_rate.lock().unwrap(),
565 }
566 }
567 }
568
569 struct FakeStats {
570 send_rate: Option<u64>,
571 }
572
573 impl web_transport_trait::Stats for FakeStats {
574 fn estimated_send_rate(&self) -> Option<u64> {
575 self.send_rate
576 }
577 }
578
579 #[derive(Clone, Default)]
580 struct FakeSendStream {
581 writes: Arc<Mutex<Vec<u8>>>,
582 }
583
584 impl web_transport_trait::SendStream for FakeSendStream {
585 type Error = FakeError;
586
587 async fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
588 self.writes.lock().unwrap().put_slice(buf);
589 Ok(buf.len())
590 }
591
592 fn set_priority(&mut self, _order: u8) {}
593
594 fn finish(&mut self) -> Result<(), Self::Error> {
595 Ok(())
596 }
597
598 fn reset(&mut self, _code: u32) {}
599
600 async fn closed(&mut self) -> Result<(), Self::Error> {
601 Ok(())
602 }
603 }
604
605 struct FakeRecvStream {
606 data: VecDeque<u8>,
607 }
608
609 impl web_transport_trait::RecvStream for FakeRecvStream {
610 type Error = FakeError;
611
612 async fn read(&mut self, dst: &mut [u8]) -> Result<Option<usize>, Self::Error> {
613 if self.data.is_empty() {
614 return Ok(None);
615 }
616
617 let size = dst.len().min(self.data.len());
618 for slot in dst.iter_mut().take(size) {
619 *slot = self.data.pop_front().unwrap();
620 }
621 Ok(Some(size))
622 }
623
624 fn stop(&mut self, _code: u32) {}
625
626 async fn closed(&mut self) -> Result<(), Self::Error> {
627 Ok(())
628 }
629 }
630
631 fn mock_server_setup(negotiated: Version) -> Vec<u8> {
632 let mut encoded = Vec::new();
633 let server = setup::Server {
634 version: negotiated.into(),
635 parameters: Bytes::new(),
636 };
637 server
638 .encode(&mut encoded, Version::Ietf(ietf::Version::Draft14))
639 .unwrap();
640
641 let info = lite::SessionInfo { bitrate: Some(1) };
643 let lite_v = lite::Version::try_from(negotiated).unwrap();
644 info.encode(&mut encoded, lite_v).unwrap();
645
646 encoded
647 }
648
649 async fn run_alpn_lite_fallback_case(protocol: Option<&'static str>) {
650 let fake = FakeSession::new(protocol, mock_server_setup(Version::Lite(lite::Version::Lite01)));
651 let client = Client::new().with_versions(
652 [
653 Version::Lite(lite::Version::Lite03),
654 Version::Lite(lite::Version::Lite02),
655 Version::Lite(lite::Version::Lite01),
656 Version::Ietf(ietf::Version::Draft14),
657 ]
658 .into(),
659 );
660
661 let (_session, driver) = client.connect(fake.clone()).await.unwrap();
664 let _driver = tokio::spawn(driver);
665
666 let mut setup_bytes = Bytes::from(fake.control_writes());
668 let setup = setup::Client::decode(&mut setup_bytes, Version::Ietf(ietf::Version::Draft14)).unwrap();
669 let advertised: Vec<Version> = setup.versions.iter().map(|v| Version::try_from(*v).unwrap()).collect();
670 assert_eq!(
671 advertised,
672 vec![
673 Version::Lite(lite::Version::Lite02),
674 Version::Lite(lite::Version::Lite01),
675 Version::Ietf(ietf::Version::Draft14),
676 ]
677 );
678
679 let (code, _) = fake.wait_for_first_close().await;
686 assert_ne!(code, Error::Version.to_code(), "SessionInfo failed to decode");
687 }
688
689 #[tokio::test(start_paused = true)]
695 async fn connect_does_not_wait_for_the_peer_to_announce() {
696 let gate = kio::Producer::new(true);
698 let transport = crate::lite::test_transport::SinkSession::gated_bi(gate.consume())
699 .with_protocol(crate::version::ALPN_LITE_05);
700
701 let origin = crate::origin::Info::new(crate::Origin::new(1).unwrap()).produce();
703 let client = Client::new()
704 .with_versions([Version::Lite(lite::Version::Lite05)].into())
705 .with_subscriber(origin);
706
707 tokio::time::timeout(std::time::Duration::from_secs(30), client.connect(transport))
710 .await
711 .expect("connect waited on a peer that never announced")
712 .expect("connect failed");
713 }
714
715 #[tokio::test(start_paused = true)]
716 async fn alpn_lite_falls_back_to_draft14_and_switches_version_post_setup() {
717 run_alpn_lite_fallback_case(Some(ALPN_LITE)).await;
718 }
719
720 #[tokio::test(start_paused = true)]
721 async fn no_alpn_falls_back_to_draft14_and_switches_version_post_setup() {
722 run_alpn_lite_fallback_case(None).await;
723 }
724
725 #[test]
733 fn driver_is_caller_polled_and_holds_no_session() {
734 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
735 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
736
737 let (session, mut driver) = futures::executor::block_on(client.connect(fake.clone())).unwrap();
738 assert_eq!(session.version(), Version::Lite(lite::Version::Lite04));
739
740 assert!(driver.poll(&kio::Waiter::noop()).is_pending());
742
743 let mut context = std::task::Context::from_waker(std::task::Waker::noop());
745 assert!(std::future::Future::poll(std::pin::Pin::new(&mut driver), &mut context).is_pending());
746
747 drop(session);
750 assert_eq!(fake.state.close_events.lock().unwrap()[0].0, Error::Cancel.to_code());
751 }
752
753 #[test]
756 fn session_clones_share_the_close() {
757 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
758 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
759
760 let (session, _driver) = futures::executor::block_on(client.connect(fake.clone())).unwrap();
761 let clone = session.clone();
762
763 drop(session);
765 assert!(fake.state.close_events.lock().unwrap().is_empty());
766
767 clone.abort(Error::Cancel);
768 assert_eq!(fake.state.close_events.lock().unwrap()[0].0, Error::Cancel.to_code());
769
770 drop(clone);
772 assert_eq!(fake.state.close_events.lock().unwrap().len(), 1);
773 }
774
775 #[tokio::test(start_paused = true)]
779 async fn send_bandwidth_samples_while_the_driver_runs() {
780 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
781 fake.set_send_rate(Some(1_000_000));
782
783 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
784 let (session, driver) = client.connect(fake.clone()).await.unwrap();
785 tokio::spawn(driver);
786
787 let mut bandwidth = session.send_bandwidth().expect("backend reports an estimate");
788 assert_eq!(bandwidth.changed().await.unwrap(), Some(1_000_000));
789
790 fake.set_send_rate(Some(2_000_000));
792 assert_eq!(bandwidth.changed().await.unwrap(), Some(2_000_000));
793 }
794}