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