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 {
265 probe: lite::ProbeLevel::Report,
266 path: self.setup_path.clone(),
267 role: lite::Role::from_origins(self.publish.is_some(), self.subscribe.is_some()),
268 cost: self.cost,
269 origin: None,
271 };
272
273 let start = lite::start(lite::Config {
274 session: session.clone(),
275 setup_stream: None,
276 publish: publish.clone(),
277 subscribe: subscribe.clone(),
278 peer_origin: self.peer_origin,
279 version,
280 our_setup,
281 peer_setup: None,
282 })?;
283
284 return Ok(Session::new(
285 session,
286 version.into(),
287 start.recv_bandwidth,
288 start.driver,
289 ));
290 }
291 Some(ALPN_LITE_04) => {
292 self.versions
293 .select(Version::Lite(lite::Version::Lite04))
294 .ok_or(Error::Version)?;
295
296 let start = lite::start(lite::Config {
297 session: session.clone(),
298 setup_stream: None,
299 publish: publish.clone(),
300 subscribe: subscribe.clone(),
301 peer_origin: self.peer_origin,
302 version: lite::Version::Lite04,
303 our_setup: lite::Setup::default(),
304 peer_setup: None,
305 })?;
306
307 return Ok(Session::new(
308 session,
309 lite::Version::Lite04.into(),
310 start.recv_bandwidth,
311 start.driver,
312 ));
313 }
314 Some(ALPN_LITE_03) => {
315 self.versions
316 .select(Version::Lite(lite::Version::Lite03))
317 .ok_or(Error::Version)?;
318
319 let start = lite::start(lite::Config {
321 session: session.clone(),
322 setup_stream: None,
323 publish: publish.clone(),
324 subscribe: subscribe.clone(),
325 peer_origin: self.peer_origin,
326 version: lite::Version::Lite03,
327 our_setup: lite::Setup::default(),
328 peer_setup: None,
329 })?;
330
331 return Ok(Session::new(
332 session,
333 lite::Version::Lite03.into(),
334 start.recv_bandwidth,
335 start.driver,
336 ));
337 }
338 Some(ALPN_LITE) | None => {
339 let supported = self.versions.filter(&NEGOTIATED.into()).ok_or(Error::Version)?;
340 (Version::Ietf(ietf::Version::Draft14), supported)
341 }
342 Some(p) => return Err(Error::UnknownAlpn(p.to_string())),
343 };
344
345 let mut stream = Stream::open(&session, encoding).await?;
346
347 let ietf_encoding = ietf::Version::try_from(encoding).map_err(|_| Error::Version)?;
349
350 let mut parameters = ietf::Parameters::default();
351 parameters.set_varint(ietf::ParameterVarInt::MaxRequestId, u32::MAX as u64);
352 parameters.set_bytes(ietf::ParameterBytes::Implementation, b"moq-lite-rs".to_vec());
353 if let Some(path) = &self.setup_path {
355 parameters.set_bytes(ietf::ParameterBytes::Path, path.clone().into_bytes());
356 }
357 ietf::solicit::into_setup(&mut parameters, ietf_encoding);
358 let parameters = parameters.encode_bytes(ietf_encoding)?;
359
360 let client = setup::Client {
361 versions: supported.clone().into(),
362 parameters,
363 };
364
365 stream.writer.encode(&client).await?;
366
367 let mut server: setup::Server = stream.reader.decode().await?;
368
369 let version = supported
370 .iter()
371 .find(|v| coding::Version::from(**v) == server.version)
372 .copied()
373 .ok_or(Error::Version)?;
374
375 let (recv_bw, protocol) = match version {
376 Version::Lite(v) => {
377 let stream = stream.with_version(v);
378 let start = lite::start(lite::Config {
379 session: session.clone(),
380 setup_stream: Some(stream),
381 publish: publish.clone(),
382 subscribe: subscribe.clone(),
383 peer_origin: self.peer_origin,
384 version: v,
385 our_setup: lite::Setup::default(),
388 peer_setup: None,
389 })?;
390
391 (start.recv_bandwidth, start.driver)
392 }
393 Version::Ietf(v) => {
394 let parameters = ietf::Parameters::decode(&mut server.parameters, v)?;
397 let request_id_max = parameters
398 .get_varint(ietf::ParameterVarInt::MaxRequestId)
399 .map(ietf::RequestId);
400 let peer_declared = ietf::peer::Peer {
401 solicit: ietf::solicit::from_setup(¶meters, v)?,
402 ..Default::default()
403 };
404
405 let stream = stream.with_version(v);
406 let protocol = ietf::start(ietf::Config {
408 session: session.clone(),
409 setup: Some(stream),
410 request_id_max,
411 client: true,
412 publish: publish.clone(),
413 subscribe: subscribe.clone(),
414 peer_origin: self.peer_origin,
415 cost: self.cost,
416 version: v,
417 path: None,
418 peer_setup_stream: None,
419 peer_declared: Some(peer_declared),
420 })?;
421 (None, protocol)
422 }
423 };
424
425 Ok(Session::new(session, version, recv_bw, protocol))
426 }
427}
428
429#[cfg(test)]
430mod tests {
431 use super::*;
432 use std::{
433 collections::VecDeque,
434 sync::{Arc, Mutex},
435 };
436
437 use crate::coding::{Decode, Encode};
438 use bytes::{BufMut, Bytes};
439
440 #[derive(Debug, Clone, Default)]
441 struct FakeError;
442
443 impl std::fmt::Display for FakeError {
444 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
445 write!(f, "fake transport error")
446 }
447 }
448
449 impl std::error::Error for FakeError {}
450
451 impl web_transport_trait::Error for FakeError {
452 fn session_error(&self) -> Option<(u32, String)> {
453 Some((0, "closed".to_string()))
454 }
455 }
456
457 #[derive(Clone, Default)]
458 struct FakeSession {
459 state: Arc<FakeSessionState>,
460 }
461
462 #[derive(Default)]
463 struct FakeSessionState {
464 protocol: Option<&'static str>,
465 control_stream: Mutex<Option<(FakeSendStream, FakeRecvStream)>>,
466 close_events: Mutex<Vec<(u32, String)>>,
467 close_notify: tokio::sync::Notify,
468 control_writes: Arc<Mutex<Vec<u8>>>,
469 send_rate: Mutex<Option<u64>>,
470 }
471
472 impl FakeSession {
473 fn new(protocol: Option<&'static str>, server_control_bytes: Vec<u8>) -> Self {
474 let writes = Arc::new(Mutex::new(Vec::new()));
475 let send = FakeSendStream { writes: writes.clone() };
476 let recv = FakeRecvStream {
477 data: VecDeque::from(server_control_bytes),
478 };
479 let state = FakeSessionState {
480 protocol,
481 control_stream: Mutex::new(Some((send, recv))),
482 close_events: Mutex::new(Vec::new()),
483 close_notify: tokio::sync::Notify::new(),
484 control_writes: writes,
485 send_rate: Mutex::new(None),
486 };
487 Self { state: Arc::new(state) }
488 }
489
490 fn set_send_rate(&self, rate: Option<u64>) {
491 *self.state.send_rate.lock().unwrap() = rate;
492 }
493
494 fn control_writes(&self) -> Vec<u8> {
495 self.state.control_writes.lock().unwrap().clone()
496 }
497
498 async fn wait_for_first_close(&self) -> (u32, String) {
499 loop {
500 let notified = self.state.close_notify.notified();
501 if let Some(close) = self.state.close_events.lock().unwrap().first().cloned() {
502 return close;
503 }
504 notified.await;
505 }
506 }
507 }
508
509 impl web_transport_trait::Session for FakeSession {
510 type SendStream = FakeSendStream;
511 type RecvStream = FakeRecvStream;
512 type Error = FakeError;
513
514 async fn accept_uni(&self) -> Result<Self::RecvStream, Self::Error> {
515 std::future::pending().await
516 }
517
518 async fn accept_bi(&self) -> Result<(Self::SendStream, Self::RecvStream), Self::Error> {
519 std::future::pending().await
520 }
521
522 async fn open_bi(&self) -> Result<(Self::SendStream, Self::RecvStream), Self::Error> {
523 self.state.control_stream.lock().unwrap().take().ok_or(FakeError)
524 }
525
526 async fn open_uni(&self) -> Result<Self::SendStream, Self::Error> {
527 std::future::pending().await
528 }
529
530 fn send_datagram(&self, _payload: Bytes) -> Result<(), Self::Error> {
531 Ok(())
532 }
533
534 async fn recv_datagram(&self) -> Result<Bytes, Self::Error> {
535 std::future::pending().await
536 }
537
538 fn max_datagram_size(&self) -> usize {
539 1200
540 }
541
542 fn protocol(&self) -> Option<&str> {
543 self.state.protocol
544 }
545
546 fn close(&self, code: u32, reason: &str) {
547 self.state.close_events.lock().unwrap().push((code, reason.to_string()));
548 self.state.close_notify.notify_waiters();
549 }
550
551 async fn closed(&self) -> Self::Error {
552 loop {
553 let notified = self.state.close_notify.notified();
554 if !self.state.close_events.lock().unwrap().is_empty() {
555 return FakeError;
556 }
557 notified.await;
558 }
559 }
560
561 fn stats(&self) -> impl web_transport_trait::Stats {
562 FakeStats {
563 send_rate: *self.state.send_rate.lock().unwrap(),
564 }
565 }
566 }
567
568 struct FakeStats {
569 send_rate: Option<u64>,
570 }
571
572 impl web_transport_trait::Stats for FakeStats {
573 fn estimated_send_rate(&self) -> Option<u64> {
574 self.send_rate
575 }
576 }
577
578 #[derive(Clone, Default)]
579 struct FakeSendStream {
580 writes: Arc<Mutex<Vec<u8>>>,
581 }
582
583 impl web_transport_trait::SendStream for FakeSendStream {
584 type Error = FakeError;
585
586 async fn write(&mut self, buf: &[u8]) -> Result<usize, Self::Error> {
587 self.writes.lock().unwrap().put_slice(buf);
588 Ok(buf.len())
589 }
590
591 fn set_priority(&mut self, _order: u8) {}
592
593 fn finish(&mut self) -> Result<(), Self::Error> {
594 Ok(())
595 }
596
597 fn reset(&mut self, _code: u32) {}
598
599 async fn closed(&mut self) -> Result<(), Self::Error> {
600 Ok(())
601 }
602 }
603
604 struct FakeRecvStream {
605 data: VecDeque<u8>,
606 }
607
608 impl web_transport_trait::RecvStream for FakeRecvStream {
609 type Error = FakeError;
610
611 async fn read(&mut self, dst: &mut [u8]) -> Result<Option<usize>, Self::Error> {
612 if self.data.is_empty() {
613 return Ok(None);
614 }
615
616 let size = dst.len().min(self.data.len());
617 for slot in dst.iter_mut().take(size) {
618 *slot = self.data.pop_front().unwrap();
619 }
620 Ok(Some(size))
621 }
622
623 fn stop(&mut self, _code: u32) {}
624
625 async fn closed(&mut self) -> Result<(), Self::Error> {
626 Ok(())
627 }
628 }
629
630 fn mock_server_setup(negotiated: Version) -> Vec<u8> {
631 let mut encoded = Vec::new();
632 let server = setup::Server {
633 version: negotiated.into(),
634 parameters: Bytes::new(),
635 };
636 server
637 .encode(&mut encoded, Version::Ietf(ietf::Version::Draft14))
638 .unwrap();
639
640 let info = lite::SessionInfo { bitrate: Some(1) };
642 let lite_v = lite::Version::try_from(negotiated).unwrap();
643 info.encode(&mut encoded, lite_v).unwrap();
644
645 encoded
646 }
647
648 async fn run_alpn_lite_fallback_case(protocol: Option<&'static str>) {
649 let fake = FakeSession::new(protocol, mock_server_setup(Version::Lite(lite::Version::Lite01)));
650 let client = Client::new().with_versions(
651 [
652 Version::Lite(lite::Version::Lite03),
653 Version::Lite(lite::Version::Lite02),
654 Version::Lite(lite::Version::Lite01),
655 Version::Ietf(ietf::Version::Draft14),
656 ]
657 .into(),
658 );
659
660 let (_session, driver) = client.connect(fake.clone()).await.unwrap();
663 let _driver = tokio::spawn(driver);
664
665 let mut setup_bytes = Bytes::from(fake.control_writes());
667 let setup = setup::Client::decode(&mut setup_bytes, Version::Ietf(ietf::Version::Draft14)).unwrap();
668 let advertised: Vec<Version> = setup.versions.iter().map(|v| Version::try_from(*v).unwrap()).collect();
669 assert_eq!(
670 advertised,
671 vec![
672 Version::Lite(lite::Version::Lite02),
673 Version::Lite(lite::Version::Lite01),
674 Version::Ietf(ietf::Version::Draft14),
675 ]
676 );
677
678 let (code, _) = fake.wait_for_first_close().await;
685 assert_ne!(code, Error::Version.to_code(), "SessionInfo failed to decode");
686 }
687
688 #[tokio::test(start_paused = true)]
694 async fn connect_does_not_wait_for_the_peer_to_announce() {
695 let gate = kio::Producer::new(true);
697 let transport = crate::lite::test_transport::SinkSession::gated_bi(gate.consume())
698 .with_protocol(crate::version::ALPN_LITE_05);
699
700 let origin = crate::origin::Info::new(crate::Origin::new(1).unwrap()).produce();
702 let client = Client::new()
703 .with_versions([Version::Lite(lite::Version::Lite05)].into())
704 .with_subscriber(origin);
705
706 tokio::time::timeout(std::time::Duration::from_secs(30), client.connect(transport))
709 .await
710 .expect("connect waited on a peer that never announced")
711 .expect("connect failed");
712 }
713
714 #[tokio::test(start_paused = true)]
715 async fn alpn_lite_falls_back_to_draft14_and_switches_version_post_setup() {
716 run_alpn_lite_fallback_case(Some(ALPN_LITE)).await;
717 }
718
719 #[tokio::test(start_paused = true)]
720 async fn no_alpn_falls_back_to_draft14_and_switches_version_post_setup() {
721 run_alpn_lite_fallback_case(None).await;
722 }
723
724 #[test]
732 fn driver_is_caller_polled_and_holds_no_session() {
733 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
734 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
735
736 let (session, mut driver) = futures::executor::block_on(client.connect(fake.clone())).unwrap();
737 assert_eq!(session.version(), Version::Lite(lite::Version::Lite04));
738
739 assert!(driver.poll(&kio::Waiter::noop()).is_pending());
741
742 let mut context = std::task::Context::from_waker(std::task::Waker::noop());
744 assert!(std::future::Future::poll(std::pin::Pin::new(&mut driver), &mut context).is_pending());
745
746 drop(session);
749 assert_eq!(fake.state.close_events.lock().unwrap()[0].0, Error::Cancel.to_code());
750 }
751
752 #[test]
755 fn session_clones_share_the_close() {
756 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
757 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
758
759 let (session, _driver) = futures::executor::block_on(client.connect(fake.clone())).unwrap();
760 let clone = session.clone();
761
762 drop(session);
764 assert!(fake.state.close_events.lock().unwrap().is_empty());
765
766 clone.abort(Error::Cancel);
767 assert_eq!(fake.state.close_events.lock().unwrap()[0].0, Error::Cancel.to_code());
768
769 drop(clone);
771 assert_eq!(fake.state.close_events.lock().unwrap().len(), 1);
772 }
773
774 #[tokio::test(start_paused = true)]
778 async fn send_bandwidth_samples_while_the_driver_runs() {
779 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
780 fake.set_send_rate(Some(1_000_000));
781
782 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
783 let (session, driver) = client.connect(fake.clone()).await.unwrap();
784 tokio::spawn(driver);
785
786 let mut bandwidth = session.send_bandwidth().expect("backend reports an estimate");
787 assert_eq!(bandwidth.changed().await.unwrap(), Some(1_000_000));
788
789 fake.set_send_rate(Some(2_000_000));
791 assert_eq!(bandwidth.changed().await.unwrap(), Some(2_000_000));
792 }
793}