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