1use crate::origin;
2#[cfg(test)]
3use crate::runtime::Timers;
4use crate::time::{Clock, Instant};
5use crate::{
6 ALPN_14, ALPN_15, ALPN_16, ALPN_17, ALPN_18, ALPN_19, ALPN_20, ALPN_21, ALPN_22, ALPN_LITE, ALPN_LITE_03,
7 ALPN_LITE_04, ALPN_LITE_05, ALPN_LITE_06, ALPN_LITE_07_WIP, Consume, Error, NEGOTIATED, Session, Version, Versions,
8 coding::{self, Decode, Encode, Stream},
9 ietf, lite, setup, stats,
10};
11
12#[derive(Default, Clone)]
14pub struct Client {
15 publish: Option<origin::Consumer>,
16 subscribe: Option<origin::Producer>,
17 stats: stats::Session,
18 versions: Versions,
19 setup_path: Option<String>,
20 setup_authority: Option<String>,
21 cost: Option<u64>,
22 peer_hop: Option<crate::Hop>,
23}
24
25impl Client {
26 pub fn new() -> Self {
28 Default::default()
29 }
30
31 pub fn with_publisher(mut self, publish: impl Consume<origin::Consumer>) -> Self {
35 self.publish = Some(publish.consume());
36 self
37 }
38
39 pub fn with_subscriber(mut self, subscribe: origin::Producer) -> Self {
42 self.subscribe = Some(subscribe);
43 self
44 }
45
46 pub fn with_stats(mut self, stats: stats::Session) -> Self {
51 self.stats = stats;
52 self
53 }
54
55 pub fn with_origin(self, origin: origin::Producer) -> Self {
60 self.with_publisher(&origin).with_subscriber(origin)
61 }
62
63 pub fn with_versions(mut self, versions: Versions) -> Self {
66 self.versions = versions;
67 self
68 }
69
70 pub fn with_path(mut self, path: impl Into<String>) -> Self {
82 self.setup_path = Some(path.into());
83 self
84 }
85
86 pub fn with_authority(mut self, authority: impl Into<String>) -> Self {
88 self.setup_authority = Some(authority.into());
89 self
90 }
91
92 pub fn with_cost(mut self, cost: u64) -> Self {
107 self.cost = Some(cost);
108 self
109 }
110
111 pub fn with_peer_hop(mut self, hop: crate::Hop) -> Self {
128 self.peer_hop = Some(hop);
129 self
130 }
131
132 fn assigned_hop(&self) -> crate::Hop {
138 self.peer_hop.unwrap_or_else(crate::Hop::random)
139 }
140
141 fn origins(&self, peer_hop: crate::Hop) -> (Option<origin::Consumer>, Option<origin::Producer>) {
152 if self.publish.is_none() && self.subscribe.is_none() {
153 tracing::warn!("not publishing or consuming anything");
154 }
155 let publish = self.publish.clone().map(|origin| origin.with_stats(self.stats.clone()));
156 let subscribe = self
157 .subscribe
158 .clone()
159 .map(|origin| origin.with_stats(self.stats.clone()));
160 let publish = publish.map(|origin| origin.excluding(peer_hop));
161 (publish, subscribe)
162 }
163
164 fn start_lite<S>(
167 &self,
168 runtime: Clock,
169 session: S,
170 version: lite::Version,
171 peer_hop: crate::Hop,
172 ) -> Result<(Session, crate::Driver<S>), Error>
173 where
174 S: crate::transport::poll::Session,
175 {
176 let (publish, subscribe) = self.origins(peer_hop);
177
178 let our_setup = if version.has_setup_stream() {
184 lite::Setup {
185 probe: lite::ProbeLevel::detect(&session),
186 path: self.setup_path.clone(),
187 role: lite::Role::from_origins(self.publish.is_some(), self.subscribe.is_some()),
188 cost: self.cost,
189 hop: None,
191 }
192 } else {
193 lite::Setup::default()
194 };
195
196 let start = lite::start(lite::Config {
197 runtime: runtime.clone(),
198 session: session.clone(),
199 setup_stream: None,
200 publish,
201 subscribe,
202 peer_hop: Some(peer_hop),
203 version,
204 our_setup,
205 peer_setup: None,
206 })?;
207
208 Ok(Session::new(
209 runtime,
210 session,
211 version.into(),
212 start.recv_bandwidth,
213 crate::driver::Protocol::Lite(Box::new(start.driver)),
214 start.goaway,
215 ))
216 }
217
218 pub async fn connect_lite<S>(&self, now: Instant, session: S) -> Result<(Session, crate::Driver<S>), Error>
228 where
229 S: crate::transport::poll::Session,
230 {
231 let runtime = Clock::new(now);
232 let version = match session.protocol() {
233 Some(ALPN_LITE_07_WIP) => lite::Version::Lite07,
234 Some(ALPN_LITE_06) => lite::Version::Lite06,
235 Some(ALPN_LITE_05) => lite::Version::Lite05,
236 Some(ALPN_LITE_04) => lite::Version::Lite04,
237 Some(ALPN_LITE_03) => lite::Version::Lite03,
238 _ => return Err(Error::Version),
239 };
240 self.versions.select(Version::Lite(version)).ok_or(Error::Version)?;
241 self.start_lite(runtime, session, version, self.assigned_hop())
242 }
243
244 pub async fn connect<S>(&self, now: Instant, mut session: S) -> Result<(Session, crate::Driver<S>), Error>
248 where
249 S: crate::transport::poll::Boxable,
250 {
251 let runtime = Clock::new(now);
252 let peer_hop = self.assigned_hop();
253 let (publish, subscribe) = self.origins(peer_hop);
254
255 let (encoding, supported) = match session.protocol() {
258 Some(alpn @ (ALPN_22 | ALPN_21 | ALPN_20 | ALPN_19 | ALPN_18 | ALPN_17)) => {
259 let draft = match alpn {
260 ALPN_22 => ietf::Version::Draft22,
261 ALPN_21 => ietf::Version::Draft21,
262 ALPN_20 => ietf::Version::Draft20,
263 ALPN_19 => ietf::Version::Draft19,
264 ALPN_18 => ietf::Version::Draft18,
265 _ => ietf::Version::Draft17,
266 };
267
268 let v = self.versions.select(Version::Ietf(draft)).ok_or(Error::Version)?;
269
270 let (protocol, goaway) = ietf::start(ietf::Config {
273 runtime: runtime.clone(),
274 session: session.clone(),
275 setup: None,
276 request_id_max: None,
277 client: true,
278 publish: publish.clone(),
279 subscribe: subscribe.clone(),
280 peer_hop: Some(peer_hop),
281 cost: self.cost,
282 version: draft,
283 path: self.setup_path.clone(),
284 authority: self.setup_authority.clone(),
285 peer_setup_stream: None,
286 peer_declared: None,
287 early_unis: Vec::new(),
288 })?;
289
290 tracing::debug!(version = ?v, "connected");
291 return Ok(Session::new(
292 runtime,
293 session,
294 v,
295 None,
296 crate::driver::Protocol::Ietf(protocol),
297 goaway,
298 ));
299 }
300 Some(ALPN_16) => {
301 let v = self
302 .versions
303 .select(Version::Ietf(ietf::Version::Draft16))
304 .ok_or(Error::Version)?;
305 (v, v.into())
306 }
307 Some(ALPN_15) => {
308 let v = self
309 .versions
310 .select(Version::Ietf(ietf::Version::Draft15))
311 .ok_or(Error::Version)?;
312 (v, v.into())
313 }
314 Some(ALPN_14) => {
315 let v = self
316 .versions
317 .select(Version::Ietf(ietf::Version::Draft14))
318 .ok_or(Error::Version)?;
319 (v, v.into())
320 }
321 Some(alpn @ (ALPN_LITE_05 | ALPN_LITE_06 | ALPN_LITE_07_WIP)) => {
322 let version = match alpn {
323 ALPN_LITE_07_WIP => lite::Version::Lite07,
324 ALPN_LITE_06 => lite::Version::Lite06,
325 _ => lite::Version::Lite05,
326 };
327 self.versions.select(Version::Lite(version)).ok_or(Error::Version)?;
328 return self.start_lite(runtime, session, version, peer_hop);
329 }
330 Some(ALPN_LITE_04) => {
331 self.versions
332 .select(Version::Lite(lite::Version::Lite04))
333 .ok_or(Error::Version)?;
334 return self.start_lite(runtime, session, lite::Version::Lite04, peer_hop);
335 }
336 Some(ALPN_LITE_03) => {
337 self.versions
338 .select(Version::Lite(lite::Version::Lite03))
339 .ok_or(Error::Version)?;
340 return self.start_lite(runtime, session, lite::Version::Lite03, peer_hop);
341 }
342 Some(ALPN_LITE) | None => {
343 let supported = self.versions.filter(&NEGOTIATED.into()).ok_or(Error::Version)?;
344 (Version::Ietf(ietf::Version::Draft14), supported)
345 }
346 Some(p) => return Err(Error::UnknownAlpn(p.to_string())),
347 };
348
349 let mut stream = Stream::open(&mut session, encoding).await?;
350
351 let ietf_encoding = ietf::Version::try_from(encoding).map_err(|_| Error::Version)?;
353
354 let mut parameters = ietf::Parameters::default();
355 parameters.set_varint(ietf::ParameterVarInt::MaxRequestId, u32::MAX as u64);
356 parameters.set_bytes(ietf::ParameterBytes::Implementation, b"moq-lite-rs".to_vec());
357 if let Some(path) = &self.setup_path {
359 parameters.set_bytes(ietf::ParameterBytes::Path, path.clone().into_bytes());
360 }
361 if let Some(authority) = &self.setup_authority {
362 parameters.set_bytes(ietf::ParameterBytes::Authority, authority.clone().into_bytes());
363 }
364 ietf::solicit::into_setup(&mut parameters, ietf_encoding);
365 ietf::hidden::into_setup(&mut parameters, ietf_encoding);
366 let parameters = parameters.encode_bytes(ietf_encoding)?;
367
368 let client = setup::Client {
369 versions: supported.clone().into(),
370 parameters,
371 };
372
373 stream.writer.encode(&client).await?;
374
375 let mut server: setup::Server = stream.reader.decode().await?;
376
377 let version = supported
378 .iter()
379 .find(|v| coding::Version::from(**v) == server.version)
380 .copied()
381 .ok_or(Error::Version)?;
382
383 let (recv_bw, protocol, goaway) = match version {
384 Version::Lite(v) => {
385 let stream = stream.with_version(v);
386 let start = lite::start(lite::Config {
387 runtime: runtime.clone(),
388 session: session.clone(),
389 setup_stream: Some(stream),
390 publish: publish.clone(),
391 subscribe: subscribe.clone(),
392 peer_hop: Some(peer_hop),
393 version: v,
394 our_setup: lite::Setup::default(),
397 peer_setup: None,
398 })?;
399
400 (
401 start.recv_bandwidth,
402 crate::driver::Protocol::Lite(Box::new(start.driver)),
403 start.goaway,
404 )
405 }
406 Version::Ietf(v) => {
407 let parameters = ietf::Parameters::decode(&mut server.parameters, v)?;
410 let request_id_max = parameters
411 .get_varint(ietf::ParameterVarInt::MaxRequestId)
412 .map(ietf::RequestId);
413 let peer_declared = ietf::peer::Peer {
414 solicit: ietf::solicit::from_setup(¶meters, v)?,
415 hidden: ietf::hidden::from_setup(¶meters, v),
416 ..Default::default()
417 };
418
419 let stream = stream.with_version(v);
420 let (protocol, goaway) = ietf::start(ietf::Config {
422 runtime: runtime.clone(),
423 session: session.clone(),
424 setup: Some(stream),
425 request_id_max,
426 client: true,
427 publish: publish.clone(),
428 subscribe: subscribe.clone(),
429 peer_hop: Some(peer_hop),
430 cost: self.cost,
431 version: v,
432 path: None,
433 authority: None,
434 peer_setup_stream: None,
435 peer_declared: Some(peer_declared),
436 early_unis: Vec::new(),
437 })?;
438 (None, crate::driver::Protocol::Ietf(protocol), goaway)
439 }
440 };
441
442 Ok(Session::new(runtime, session, version, recv_bw, protocol, goaway))
443 }
444}
445
446#[cfg(test)]
447mod tests {
448 use super::*;
449 use crate::model::ProduceTest;
450 use std::{
451 collections::VecDeque,
452 sync::{Arc, Mutex},
453 };
454
455 use std::task::{Context, Poll};
456
457 use crate::SessionError;
458 use crate::coding::{Decode, Encode};
459 use bytes::{BufMut, Bytes};
460
461 #[derive(Debug, Clone, Default)]
462 struct FakeError;
463
464 impl std::fmt::Display for FakeError {
465 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
466 write!(f, "fake transport error")
467 }
468 }
469
470 impl std::error::Error for FakeError {}
471
472 impl web_transport_trait::Error for FakeError {
473 fn session_error(&self) -> Option<(u32, String)> {
474 Some((0, "closed".to_string()))
475 }
476 }
477
478 #[derive(Clone, Default)]
479 struct FakeSession {
480 state: Arc<FakeSessionState>,
481 park: kio::Park,
483 }
484
485 #[derive(Default)]
486 struct FakeSessionState {
487 protocol: Option<&'static str>,
488 control_stream: Mutex<Option<(FakeSendStream, FakeRecvStream)>>,
489 close_events: Mutex<Vec<(u32, String)>>,
490 closed: kio::Fan,
491 control_writes: Arc<Mutex<Vec<u8>>>,
492 send_rate: Mutex<Option<u64>>,
493 bytes_sent: Mutex<Option<u64>>,
494 }
495
496 impl FakeSession {
497 fn new(protocol: Option<&'static str>, server_control_bytes: Vec<u8>) -> Self {
498 let writes = Arc::new(Mutex::new(Vec::new()));
499 let send = FakeSendStream { writes: writes.clone() };
500 let recv = FakeRecvStream {
501 data: VecDeque::from(server_control_bytes),
502 };
503 let state = FakeSessionState {
504 protocol,
505 control_stream: Mutex::new(Some((send, recv))),
506 close_events: Mutex::new(Vec::new()),
507 closed: kio::Fan::default(),
508 control_writes: writes,
509 send_rate: Mutex::new(None),
510 bytes_sent: Mutex::new(None),
511 };
512 Self {
513 state: Arc::new(state),
514 park: kio::Park::default(),
515 }
516 }
517
518 fn set_send_rate(&self, rate: Option<u64>) {
519 *self.state.send_rate.lock().unwrap() = rate;
520 }
521
522 fn set_bytes_sent(&self, bytes: Option<u64>) {
523 *self.state.bytes_sent.lock().unwrap() = bytes;
524 }
525
526 fn control_writes(&self) -> Vec<u8> {
527 self.state.control_writes.lock().unwrap().clone()
528 }
529
530 async fn wait_for_first_close(&self) -> (u32, String) {
531 kio::wait(|waiter| {
532 self.state.closed.register(waiter);
533 match self.state.close_events.lock().unwrap().first().cloned() {
534 Some(close) => std::task::Poll::Ready(close),
535 None => std::task::Poll::Pending,
536 }
537 })
538 .await
539 }
540 }
541
542 impl web_transport_trait::poll::Session for FakeSession {
543 type SendStream = FakeSendStream;
544 type RecvStream = FakeRecvStream;
545 type Error = FakeError;
546
547 fn poll_accept_uni(&mut self, _cx: &mut Context<'_>) -> Poll<Result<Self::RecvStream, Self::Error>> {
548 Poll::Pending
549 }
550
551 fn poll_accept_bi(
552 &mut self,
553 _cx: &mut Context<'_>,
554 ) -> Poll<Result<(Self::SendStream, Self::RecvStream), Self::Error>> {
555 Poll::Pending
556 }
557
558 fn poll_open_bi(
559 &mut self,
560 _cx: &mut Context<'_>,
561 ) -> Poll<Result<(Self::SendStream, Self::RecvStream), Self::Error>> {
562 Poll::Ready(self.state.control_stream.lock().unwrap().take().ok_or(FakeError))
563 }
564
565 fn poll_open_uni(&mut self, _cx: &mut Context<'_>) -> Poll<Result<Self::SendStream, Self::Error>> {
566 Poll::Pending
567 }
568
569 fn poll_send_datagram(&mut self, _cx: &mut Context<'_>, _payload: &[u8]) -> Poll<Result<(), Self::Error>> {
570 Poll::Ready(Ok(()))
571 }
572
573 fn poll_recv_datagram(&mut self, _cx: &mut Context<'_>) -> Poll<Result<Bytes, Self::Error>> {
574 Poll::Pending
575 }
576
577 fn max_datagram_size(&self) -> usize {
578 1200
579 }
580
581 fn protocol(&self) -> Option<&str> {
582 self.state.protocol
583 }
584
585 fn close(&mut self, code: u32, reason: &str) {
586 self.state.close_events.lock().unwrap().push((code, reason.to_string()));
587 self.state.closed.wake();
588 }
589
590 fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Self::Error> {
591 self.state.closed.register(self.park.hold(cx));
593 match self.state.close_events.lock().unwrap().is_empty() {
594 false => Poll::Ready(FakeError),
595 true => Poll::Pending,
596 }
597 }
598
599 fn stats(&self) -> impl web_transport_trait::Stats {
600 FakeStats {
601 send_rate: *self.state.send_rate.lock().unwrap(),
602 bytes_sent: *self.state.bytes_sent.lock().unwrap(),
603 }
604 }
605 }
606
607 struct FakeStats {
608 send_rate: Option<u64>,
609 bytes_sent: Option<u64>,
610 }
611
612 impl web_transport_trait::Stats for FakeStats {
613 fn estimated_send_rate(&self) -> Option<u64> {
614 self.send_rate
615 }
616
617 fn bytes_sent(&self) -> Option<u64> {
618 self.bytes_sent
619 }
620 }
621
622 #[derive(Clone, Default)]
623 struct FakeSendStream {
624 writes: Arc<Mutex<Vec<u8>>>,
625 }
626
627 impl web_transport_trait::poll::SendStream for FakeSendStream {
628 type Error = FakeError;
629
630 fn poll_write(&mut self, _cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Self::Error>> {
631 self.writes.lock().unwrap().put_slice(buf);
632 Poll::Ready(Ok(buf.len()))
633 }
634
635 fn set_priority(&mut self, _order: u8) {}
636
637 fn finish(&mut self) -> Result<(), Self::Error> {
638 Ok(())
639 }
640
641 fn reset(&mut self, _code: u32) {}
642
643 fn poll_closed(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
644 Poll::Ready(Ok(()))
645 }
646 }
647
648 struct FakeRecvStream {
649 data: VecDeque<u8>,
650 }
651
652 impl web_transport_trait::poll::RecvStream for FakeRecvStream {
653 type Error = FakeError;
654
655 fn poll_read(&mut self, _cx: &mut Context<'_>, dst: &mut [u8]) -> Poll<Result<Option<usize>, Self::Error>> {
656 if self.data.is_empty() {
657 return Poll::Ready(Ok(None));
658 }
659
660 let size = dst.len().min(self.data.len());
661 for slot in dst.iter_mut().take(size) {
662 *slot = self.data.pop_front().unwrap();
663 }
664 Poll::Ready(Ok(Some(size)))
665 }
666
667 fn stop(&mut self, _code: u32) {}
668
669 fn poll_closed(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
670 Poll::Ready(Ok(()))
671 }
672 }
673
674 fn mock_server_setup(negotiated: Version) -> Vec<u8> {
675 let mut encoded = Vec::new();
676 let server = setup::Server {
677 version: negotiated.into(),
678 parameters: Bytes::new(),
679 };
680 server
681 .encode(&mut encoded, Version::Ietf(ietf::Version::Draft14))
682 .unwrap();
683
684 let info = lite::SessionInfo { bitrate: Some(1) };
686 let lite_v = lite::Version::try_from(negotiated).unwrap();
687 info.encode(&mut encoded, lite_v).unwrap();
688
689 encoded
690 }
691
692 async fn run_alpn_lite_fallback_case(protocol: Option<&'static str>) {
693 let fake = FakeSession::new(protocol, mock_server_setup(Version::Lite(lite::Version::Lite01)));
694 let client = Client::new().with_versions(
695 [
696 Version::Lite(lite::Version::Lite03),
697 Version::Lite(lite::Version::Lite02),
698 Version::Lite(lite::Version::Lite01),
699 Version::Ietf(ietf::Version::Draft14),
700 ]
701 .into(),
702 );
703
704 let (_session, driver) = client
706 .connect(tokio::time::Instant::now().into_std(), fake.clone())
707 .await
708 .unwrap();
709 tokio::spawn(crate::time::run(driver));
710
711 let mut setup_bytes = Bytes::from(fake.control_writes());
713 let setup = setup::Client::decode(&mut setup_bytes, Version::Ietf(ietf::Version::Draft14)).unwrap();
714 let advertised: Vec<Version> = setup.versions.iter().map(|v| Version::try_from(*v).unwrap()).collect();
715 assert_eq!(
716 advertised,
717 vec![
718 Version::Lite(lite::Version::Lite02),
719 Version::Lite(lite::Version::Lite01),
720 Version::Ietf(ietf::Version::Draft14),
721 ]
722 );
723
724 let (code, _) = fake.wait_for_first_close().await;
731 assert_ne!(code, SessionError::Version.to_code(), "SessionInfo failed to decode");
733 }
734
735 #[tokio::test(start_paused = true)]
741 async fn connect_does_not_wait_for_the_peer_to_announce() {
742 let gate = kio::Producer::new(true);
744 let transport = crate::lite::test_transport::SinkSession::gated_bi(gate.consume())
745 .with_protocol(crate::version::ALPN_LITE_05);
746
747 let origin = crate::origin::Config::new(crate::Hop::new(1).unwrap()).produce();
749 let client = Client::new()
750 .with_versions([Version::Lite(lite::Version::Lite05)].into())
751 .with_subscriber(origin);
752
753 let (_session, _driver) = tokio::time::timeout(
756 std::time::Duration::from_secs(30),
757 client.connect(tokio::time::Instant::now().into_std(), transport),
758 )
759 .await
760 .expect("connect waited on a peer that never announced")
761 .expect("connect failed");
762 }
763
764 #[tokio::test(start_paused = true)]
767 async fn draft14_setup_carries_the_authority() {
768 let fake = FakeSession::new(Some(ALPN_LITE), mock_server_setup(Version::Lite(lite::Version::Lite01)));
769 let client = Client::new()
770 .with_versions(
771 [
772 Version::Lite(lite::Version::Lite01),
773 Version::Ietf(ietf::Version::Draft14),
774 ]
775 .into(),
776 )
777 .with_path("/anon")
778 .with_authority("relay.example.com:4443");
779
780 let (_session, driver) = client
781 .connect(tokio::time::Instant::now().into_std(), fake.clone())
782 .await
783 .unwrap();
784 tokio::spawn(crate::time::run(driver));
785
786 let mut setup_bytes = Bytes::from(fake.control_writes());
787 let setup = setup::Client::decode(&mut setup_bytes, Version::Ietf(ietf::Version::Draft14)).unwrap();
788 let mut parameters = setup.parameters;
789 let parameters = ietf::Parameters::decode(&mut parameters, ietf::Version::Draft14).unwrap();
790 assert_eq!(
791 parameters.get_bytes(ietf::ParameterBytes::Authority),
792 Some(b"relay.example.com:4443".as_ref())
793 );
794 assert_eq!(
795 parameters.get_bytes(ietf::ParameterBytes::Path),
796 Some(b"/anon".as_ref())
797 );
798 }
799
800 #[tokio::test(start_paused = true)]
801 async fn alpn_lite_falls_back_to_draft14_and_switches_version_post_setup() {
802 run_alpn_lite_fallback_case(Some(ALPN_LITE)).await;
803 }
804
805 #[tokio::test(start_paused = true)]
806 async fn no_alpn_falls_back_to_draft14_and_switches_version_post_setup() {
807 run_alpn_lite_fallback_case(None).await;
808 }
809
810 #[test]
813 fn driver_is_caller_polled_and_holds_no_session() {
814 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
815 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
816
817 let runtime = crate::runtime::Test::new();
818 let (session, mut driver) = futures::executor::block_on(client.connect(runtime.now(), fake.clone())).unwrap();
819 assert_eq!(session.version(), Version::Lite(lite::Version::Lite04));
820
821 assert!(driver.poll(runtime.now(), &kio::Waiter::noop()).is_ok());
823
824 drop(session);
827 assert!(fake.state.close_events.lock().unwrap().is_empty());
828 let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
829 assert_eq!(
830 fake.state.close_events.lock().unwrap()[0].0,
831 SessionError::Cancel.to_code()
832 );
833 }
834
835 #[test]
839 fn session_clones_share_the_close() {
840 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
841 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
842
843 let runtime = crate::runtime::Test::new();
844 let (session, mut driver) = futures::executor::block_on(client.connect(runtime.now(), fake.clone())).unwrap();
845 let clone = session.clone();
846
847 drop(session);
849 assert!(fake.state.close_events.lock().unwrap().is_empty());
850 let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
851 assert!(fake.state.close_events.lock().unwrap().is_empty());
852
853 clone.abort(Error::Cancel);
854 let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
855 assert_eq!(
856 fake.state.close_events.lock().unwrap()[0].0,
857 SessionError::Cancel.to_code()
858 );
859
860 let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
863 futures::executor::block_on(clone.closed());
864
865 let closes = fake.state.close_events.lock().unwrap().len();
867 drop(clone);
868 let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
869 assert_eq!(fake.state.close_events.lock().unwrap().len(), closes);
870 }
871
872 #[test]
876 fn dropped_driver_resolves_closed() {
877 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
878 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
879
880 let runtime = crate::runtime::Test::new();
881 let (session, driver) = futures::executor::block_on(client.connect(runtime.now(), fake.clone())).unwrap();
882
883 drop(driver);
884 assert!(matches!(futures::executor::block_on(session.closed()), Error::Cancel));
885 }
886
887 #[derive(Clone)]
891 struct LocalSession {
892 inner: FakeSession,
893 _local: std::rc::Rc<()>,
894 }
895
896 struct LocalSend {
897 inner: FakeSendStream,
898 _local: std::rc::Rc<()>,
899 }
900
901 struct LocalRecv {
902 inner: FakeRecvStream,
903 _local: std::rc::Rc<()>,
904 }
905
906 impl web_transport_trait::poll::Session for LocalSession {
907 type SendStream = LocalSend;
908 type RecvStream = LocalRecv;
909 type Error = FakeError;
910
911 fn poll_accept_uni(&mut self, cx: &mut Context<'_>) -> Poll<Result<Self::RecvStream, Self::Error>> {
912 self.inner.poll_accept_uni(cx).map_ok(|stream| LocalRecv {
913 inner: stream,
914 _local: self._local.clone(),
915 })
916 }
917
918 fn poll_accept_bi(
919 &mut self,
920 cx: &mut Context<'_>,
921 ) -> Poll<Result<(Self::SendStream, Self::RecvStream), Self::Error>> {
922 self.inner.poll_accept_bi(cx).map_ok(|(send, recv)| {
923 (
924 LocalSend {
925 inner: send,
926 _local: self._local.clone(),
927 },
928 LocalRecv {
929 inner: recv,
930 _local: self._local.clone(),
931 },
932 )
933 })
934 }
935
936 fn poll_open_bi(
937 &mut self,
938 cx: &mut Context<'_>,
939 ) -> Poll<Result<(Self::SendStream, Self::RecvStream), Self::Error>> {
940 self.inner.poll_open_bi(cx).map_ok(|(send, recv)| {
941 (
942 LocalSend {
943 inner: send,
944 _local: self._local.clone(),
945 },
946 LocalRecv {
947 inner: recv,
948 _local: self._local.clone(),
949 },
950 )
951 })
952 }
953
954 fn poll_open_uni(&mut self, cx: &mut Context<'_>) -> Poll<Result<Self::SendStream, Self::Error>> {
955 self.inner.poll_open_uni(cx).map_ok(|stream| LocalSend {
956 inner: stream,
957 _local: self._local.clone(),
958 })
959 }
960
961 fn poll_send_datagram(&mut self, cx: &mut Context<'_>, payload: &[u8]) -> Poll<Result<(), Self::Error>> {
962 self.inner.poll_send_datagram(cx, payload)
963 }
964
965 fn poll_recv_datagram(&mut self, cx: &mut Context<'_>) -> Poll<Result<Bytes, Self::Error>> {
966 self.inner.poll_recv_datagram(cx)
967 }
968
969 fn max_datagram_size(&self) -> usize {
970 self.inner.max_datagram_size()
971 }
972
973 fn protocol(&self) -> Option<&str> {
974 self.inner.protocol()
975 }
976
977 fn close(&mut self, code: u32, reason: &str) {
978 self.inner.close(code, reason);
979 }
980
981 fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Self::Error> {
982 self.inner.poll_closed(cx)
983 }
984
985 fn stats(&self) -> impl web_transport_trait::Stats {
986 self.inner.stats()
987 }
988 }
989
990 impl web_transport_trait::poll::SendStream for LocalSend {
991 type Error = FakeError;
992
993 fn poll_write(&mut self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Self::Error>> {
994 self.inner.poll_write(cx, buf)
995 }
996
997 fn set_priority(&mut self, order: u8) {
998 self.inner.set_priority(order);
999 }
1000
1001 fn finish(&mut self) -> Result<(), Self::Error> {
1002 self.inner.finish()
1003 }
1004
1005 fn reset(&mut self, code: u32) {
1006 self.inner.reset(code);
1007 }
1008
1009 fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
1010 web_transport_trait::poll::SendStream::poll_closed(&mut self.inner, cx)
1011 }
1012 }
1013
1014 impl web_transport_trait::poll::RecvStream for LocalRecv {
1015 type Error = FakeError;
1016
1017 fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll<Result<Option<usize>, Self::Error>> {
1018 self.inner.poll_read(cx, dst)
1019 }
1020
1021 fn stop(&mut self, code: u32) {
1022 self.inner.stop(code);
1023 }
1024
1025 fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
1026 web_transport_trait::poll::RecvStream::poll_closed(&mut self.inner, cx)
1027 }
1028 }
1029
1030 #[test]
1035 fn connect_lite_over_a_send_less_transport() {
1036 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1037 let local = LocalSession {
1038 inner: fake.clone(),
1039 _local: std::rc::Rc::new(()),
1040 };
1041 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1042
1043 let runtime = crate::runtime::Test::new();
1044 let (session, mut driver) = futures::executor::block_on(client.connect_lite(runtime.now(), local)).unwrap();
1045 assert!(driver.poll(runtime.now(), &kio::Waiter::noop()).is_ok());
1046
1047 fn assert_send_sync<T: Send + Sync>(_: &T) {}
1048 assert_send_sync(&session);
1049
1050 session.abort(Error::Cancel);
1051 let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
1052 assert_eq!(
1053 fake.state.close_events.lock().unwrap()[0].0,
1054 SessionError::Cancel.to_code()
1055 );
1056 }
1057
1058 #[test]
1061 fn accept_lite_over_a_send_less_transport() {
1062 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1063 let local = LocalSession {
1064 inner: fake.clone(),
1065 _local: std::rc::Rc::new(()),
1066 };
1067 let server = crate::Server::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1068
1069 let runtime = crate::runtime::Test::new();
1070 let (session, mut driver) = futures::executor::block_on(server.accept_lite(runtime.now(), local)).unwrap();
1071 assert_eq!(session.version(), Version::Lite(lite::Version::Lite04));
1072 assert!(driver.poll(runtime.now(), &kio::Waiter::noop()).is_ok());
1073
1074 drop(session);
1075 assert!(fake.state.close_events.lock().unwrap().is_empty());
1076 let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
1077 assert_eq!(
1078 fake.state.close_events.lock().unwrap()[0].0,
1079 SessionError::Cancel.to_code()
1080 );
1081 }
1082
1083 #[test]
1086 fn connect_lite_refuses_ietf_alpns() {
1087 let fake = FakeSession::new(Some(ALPN_19), Vec::new());
1088 let local = LocalSession {
1089 inner: fake,
1090 _local: std::rc::Rc::new(()),
1091 };
1092 let client = Client::new();
1093 let runtime = crate::runtime::Test::new();
1094 let result = futures::executor::block_on(client.connect_lite(runtime.now(), local));
1095 assert!(matches!(result, Err(Error::Version)));
1096 }
1097
1098 #[tokio::test(start_paused = true)]
1102 async fn stats_reads_prime_the_sampler() {
1103 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1104 fake.set_send_rate(Some(1_000_000));
1105
1106 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1107 let (session, driver) = client
1108 .connect(tokio::time::Instant::now().into_std(), fake.clone())
1109 .await
1110 .unwrap();
1111 tokio::spawn(crate::time::run(driver));
1112
1113 assert_eq!(
1115 session.stats().estimated_send_rate,
1116 Some(crate::bandwidth::Rate::from_bps(1_000_000))
1117 );
1118
1119 fake.set_send_rate(Some(2_000_000));
1122 while session.stats().estimated_send_rate != Some(crate::bandwidth::Rate::from_bps(2_000_000)) {
1123 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1124 }
1125 }
1126
1127 #[tokio::test(start_paused = true)]
1134 async fn stats_capture_the_final_counters() {
1135 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1136 fake.set_send_rate(None);
1137 fake.set_bytes_sent(Some(0));
1138
1139 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1140 let (session, driver) = client
1141 .connect(tokio::time::Instant::now().into_std(), fake.clone())
1142 .await
1143 .unwrap();
1144 tokio::spawn(crate::time::run(driver));
1145 assert!(
1146 session.send_bandwidth().is_none(),
1147 "no send-rate estimate, so nothing samples on its own"
1148 );
1149
1150 fake.set_bytes_sent(Some(4242));
1151
1152 session.abort(Error::Cancel);
1153 session.closed().await;
1154
1155 assert_eq!(
1156 session.stats().bytes_sent,
1157 Some(4242),
1158 "the closing snapshot must carry the session's final counters"
1159 );
1160 }
1161
1162 #[tokio::test(start_paused = true)]
1166 async fn send_bandwidth_samples_while_the_driver_runs() {
1167 let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1168 fake.set_send_rate(Some(1_000_000));
1169
1170 let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1171 let (session, driver) = client
1172 .connect(tokio::time::Instant::now().into_std(), fake.clone())
1173 .await
1174 .unwrap();
1175 tokio::spawn(crate::time::run(driver));
1176
1177 let mut bandwidth = session.send_bandwidth().expect("backend reports an estimate");
1178 assert_eq!(
1179 bandwidth.changed().await.unwrap(),
1180 Some(crate::bandwidth::Rate::from_bps(1_000_000))
1181 );
1182
1183 fake.set_send_rate(Some(2_000_000));
1185 assert_eq!(
1186 bandwidth.changed().await.unwrap(),
1187 Some(crate::bandwidth::Rate::from_bps(2_000_000))
1188 );
1189 }
1190}