Skip to main content

moq_net/
client.rs

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/// A MoQ client session builder.
13#[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	/// A client that neither publishes nor subscribes until configured.
27	pub fn new() -> Self {
28		Default::default()
29	}
30
31	/// Publish local broadcasts to the remote: the session reads from the given
32	/// origin (pass an [`origin::Producer`] or [`origin::Consumer`] by reference) and
33	/// forwards its announcements. Omit to publish nothing.
34	pub fn with_publisher(mut self, publish: impl Consume<origin::Consumer>) -> Self {
35		self.publish = Some(publish.consume());
36		self
37	}
38
39	/// Subscribe to remote broadcasts: the session writes the broadcasts the
40	/// remote announces into this [`origin::Producer`]. Omit to subscribe to nothing.
41	pub fn with_subscriber(mut self, subscribe: origin::Producer) -> Self {
42		self.subscribe = Some(subscribe);
43		self
44	}
45
46	/// Attach a per-connection [`stats::Session`] context. The session's publish
47	/// (egress) and subscribe (ingress) origin handles are tagged with it, so all
48	/// traffic counters are attributed through the model for this session's lifetime.
49	/// Pass [`stats::Session::default`] (a no-op context) to opt out.
50	pub fn with_stats(mut self, stats: stats::Session) -> Self {
51		self.stats = stats;
52		self
53	}
54
55	/// Set both publish and subscribe from one shared [`origin::Producer`].
56	///
57	/// Equivalent to [`with_publisher`](Self::with_publisher) and
58	/// [`with_subscriber`](Self::with_subscriber) with the same origin.
59	pub fn with_origin(self, origin: origin::Producer) -> Self {
60		self.with_publisher(&origin).with_subscriber(origin)
61	}
62
63	/// Restrict which protocol versions to offer, in preference order.
64	/// Defaults to every version this crate supports.
65	pub fn with_versions(mut self, versions: Versions) -> Self {
66		self.versions = versions;
67		self
68	}
69
70	/// Set the request path to advertise in SETUP (moq-lite-05 and newer, and
71	/// every moq-transport draft we speak).
72	///
73	/// Only for transports that carry no request URI of their own (native QUIC, qmux
74	/// over TCP/TLS, unix sockets), so the server learns which path the client wants.
75	/// Append `?` and the URI query when there is one: that is how a credential in the
76	/// query (`?jwt=`) reaches the server.
77	/// Bindings that already carry a URI (WebTransport, qmux over WebSocket) convey
78	/// the path there and MUST NOT send this; a server is entitled to treat it as a
79	/// protocol violation. An empty path is equivalent to omitting it. Ignored by
80	/// versions with no in-band request path (lite 01-04).
81	pub fn with_path(mut self, path: impl Into<String>) -> Self {
82		self.setup_path = Some(path.into());
83		self
84	}
85
86	/// Set the URI authority to advertise in SETUP (moq-transport only)
87	pub fn with_authority(mut self, authority: impl Into<String>) -> Self {
88		self.setup_authority = Some(authority.into());
89		self
90	}
91
92	/// Price this link, in the units the rest of the mesh uses (moq-lite-06+, and
93	/// `moqt-17`+ via the MoQ Cluster extension).
94	///
95	/// The dialer is the side that knows what a link costs, because it chose the peer:
96	/// use `0` for a sibling in the same datacenter and something large for another
97	/// region across a metered backbone. So this prices both directions. We add it to
98	/// the route cost of every announcement the peer sends us, and declare it in our
99	/// SETUP so the peer adds it to every announcement we send, which is what a server
100	/// accepting an anonymous connection needs: it cannot tell a sibling from a
101	/// stranger, so it has no price of its own to apply.
102	///
103	/// A price the peer declares applies only where we set none. An unpriced link costs
104	/// `1`, which makes the cost track the hop count and so reproduces plain
105	/// shortest-path routing.
106	pub fn with_cost(mut self, cost: u64) -> Self {
107		self.cost = Some(cost);
108		self
109	}
110
111	/// Assign the identity this peer's routes are attributed to, overriding the fresh
112	/// per-dial default.
113	///
114	/// A dialed session whose peer declares no hop gets a random one for that
115	/// connection, the same way an accepted session does. That keeps a route
116	/// learned from the peer off the session that learned it, and keeps a
117	/// SUBSCRIBE that arrives on it from being routed back to it. An identity the
118	/// peer declares on the wire still wins.
119	///
120	/// Pass an id only for a peer whose identity the caller has actually
121	/// established. Two dials given the same hop are treated as one endpoint:
122	/// routes learned from either are kept off both, and content arriving on
123	/// either is interchangeable with the other's. That is the point when they
124	/// really are one peer reconnecting or running redundant links, and a bug
125	/// otherwise. Derive it from the authenticated identity, never from something
126	/// coarser like the remote address.
127	pub fn with_peer_hop(mut self, hop: crate::Hop) -> Self {
128		self.peer_hop = Some(hop);
129		self
130	}
131
132	/// The hop this connection attributes the peer to when the peer declares none.
133	///
134	/// A caller-supplied id is stable across dials of this client. Otherwise each
135	/// connection gets a fresh one, matching the per-session default an accepted
136	/// session already has.
137	fn assigned_hop(&self) -> crate::Hop {
138		self.peer_hop.unwrap_or_else(crate::Hop::random)
139	}
140
141	/// The origin pair a session attaches, tagged and filtered.
142	///
143	/// Reads through the publish (egress) consumer and writes through the
144	/// subscribe (ingress) producer are attributed by the model through the
145	/// stats context; one shared context, so presence and viewer counts are
146	/// never double-attributed across the two halves. `peer_hop` is the identity
147	/// this connection attributes the peer to when the peer declares none, so
148	/// subscriptions from the peer resolve to a source whose route excludes it.
149	/// A declared identity still wins inside each publisher; announce filtering
150	/// is per-protocol and handled there too.
151	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	/// Start a lite session on an already-negotiated version: build our SETUP,
165	/// wire the origins, and return the session and its driver.
166	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		// Advertise our capabilities (we report what the transport measures; we
179		// don't pad) plus the request path on URI-less transports, and the
180		// direction we intend to use so the server can reject a token that lacks
181		// the matching scope during the handshake instead of silently carrying
182		// no media. Versions without a Setup Stream have nothing to advertise.
183		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				// Filled by `lite::start` from the attached origin handles.
190				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			Some(start.estimate),
216		))
217	}
218
219	/// Perform the MoQ handshake for moq-lite only, over any transport.
220	///
221	/// Unlike [`connect`](Self::connect) this puts no thread-affinity bound on
222	/// the transport, so a pinned `!Send` transport works and yields a `!Send`
223	/// machine that stays on its thread. The trade is protocol scope: only a
224	/// moq-lite ALPN is accepted, since the moq-transport driver still needs a
225	/// [`Boxable`](crate::transport::poll::Boxable) transport. An ietf ALPN, an
226	/// unknown one, or the legacy no-ALPN SETUP negotiation is refused with
227	/// [`Error::Version`].
228	pub async fn connect_lite<S>(&self, now: Instant, session: S) -> Result<(Session, crate::Driver<S>), Error>
229	where
230		S: crate::transport::poll::Session,
231	{
232		let runtime = Clock::new(now);
233		let version = match session.protocol() {
234			Some(ALPN_LITE_07_WIP) => lite::Version::Lite07,
235			Some(ALPN_LITE_06) => lite::Version::Lite06,
236			Some(ALPN_LITE_05) => lite::Version::Lite05,
237			Some(ALPN_LITE_04) => lite::Version::Lite04,
238			Some(ALPN_LITE_03) => lite::Version::Lite03,
239			_ => return Err(Error::Version),
240		};
241		self.versions.select(Version::Lite(version)).ok_or(Error::Version)?;
242		self.start_lite(runtime, session, version, self.assigned_hop())
243	}
244
245	/// Perform the MoQ handshake, returning the [`Session`] and its [`Driver`](crate::Driver).
246	///
247	/// Poll the returned driver with nondecreasing time, starting at `now`.
248	pub async fn connect<S>(&self, now: Instant, mut session: S) -> Result<(Session, crate::Driver<S>), Error>
249	where
250		S: crate::transport::poll::Boxable,
251	{
252		let runtime = Clock::new(now);
253		let peer_hop = self.assigned_hop();
254		let (publish, subscribe) = self.origins(peer_hop);
255
256		// If ALPN was used to negotiate the version, use the appropriate encoding.
257		// Default to IETF 14 if no ALPN was used and we'll negotiate the version later.
258		let (encoding, supported) = match session.protocol() {
259			Some(alpn @ (ALPN_22 | ALPN_21 | ALPN_20 | ALPN_19 | ALPN_18 | ALPN_17)) => {
260				let draft = match alpn {
261					ALPN_22 => ietf::Version::Draft22,
262					ALPN_21 => ietf::Version::Draft21,
263					ALPN_20 => ietf::Version::Draft20,
264					ALPN_19 => ietf::Version::Draft19,
265					ALPN_18 => ietf::Version::Draft18,
266					_ => ietf::Version::Draft17,
267				};
268
269				let v = self.versions.select(Version::Ietf(draft)).ok_or(Error::Version)?;
270
271				// Draft-17+: SETUP is exchanged by the connection driver.
272				// We advertise the request path in our SETUP for URL-less transports.
273				let (protocol, goaway) = ietf::start(ietf::Config {
274					runtime: runtime.clone(),
275					session: session.clone(),
276					setup: None,
277					request_id_max: None,
278					client: true,
279					publish: publish.clone(),
280					subscribe: subscribe.clone(),
281					peer_hop: Some(peer_hop),
282					cost: self.cost,
283					version: draft,
284					path: self.setup_path.clone(),
285					authority: self.setup_authority.clone(),
286					peer_setup_stream: None,
287					peer_declared: None,
288					early_unis: Vec::new(),
289				})?;
290
291				tracing::debug!(version = ?v, "connected");
292				return Ok(Session::new(
293					runtime,
294					session,
295					v,
296					None,
297					crate::driver::Protocol::Ietf(protocol),
298					goaway,
299					None,
300				));
301			}
302			Some(ALPN_16) => {
303				let v = self
304					.versions
305					.select(Version::Ietf(ietf::Version::Draft16))
306					.ok_or(Error::Version)?;
307				(v, v.into())
308			}
309			Some(ALPN_15) => {
310				let v = self
311					.versions
312					.select(Version::Ietf(ietf::Version::Draft15))
313					.ok_or(Error::Version)?;
314				(v, v.into())
315			}
316			Some(ALPN_14) => {
317				let v = self
318					.versions
319					.select(Version::Ietf(ietf::Version::Draft14))
320					.ok_or(Error::Version)?;
321				(v, v.into())
322			}
323			Some(alpn @ (ALPN_LITE_05 | ALPN_LITE_06 | ALPN_LITE_07_WIP)) => {
324				let version = match alpn {
325					ALPN_LITE_07_WIP => lite::Version::Lite07,
326					ALPN_LITE_06 => lite::Version::Lite06,
327					_ => lite::Version::Lite05,
328				};
329				self.versions.select(Version::Lite(version)).ok_or(Error::Version)?;
330				return self.start_lite(runtime, session, version, peer_hop);
331			}
332			Some(ALPN_LITE_04) => {
333				self.versions
334					.select(Version::Lite(lite::Version::Lite04))
335					.ok_or(Error::Version)?;
336				return self.start_lite(runtime, session, lite::Version::Lite04, peer_hop);
337			}
338			Some(ALPN_LITE_03) => {
339				self.versions
340					.select(Version::Lite(lite::Version::Lite03))
341					.ok_or(Error::Version)?;
342				return self.start_lite(runtime, session, lite::Version::Lite03, peer_hop);
343			}
344			Some(ALPN_LITE) | None => {
345				let supported = self.versions.filter(&NEGOTIATED.into()).ok_or(Error::Version)?;
346				(Version::Ietf(ietf::Version::Draft14), supported)
347			}
348			Some(p) => return Err(Error::UnknownAlpn(p.to_string())),
349		};
350
351		let mut stream = Stream::open(&mut session, encoding).await?;
352
353		// The encoding is always an IETF version for SETUP negotiation.
354		let ietf_encoding = ietf::Version::try_from(encoding).map_err(|_| Error::Version)?;
355
356		let mut parameters = ietf::Parameters::default();
357		parameters.set_varint(ietf::ParameterVarInt::MaxRequestId, u32::MAX as u64);
358		parameters.set_bytes(ietf::ParameterBytes::Implementation, b"moq-lite-rs".to_vec());
359		// Advertise the request path in-band (draft 14-16), same as the lite-05 SETUP.
360		if let Some(path) = &self.setup_path {
361			parameters.set_bytes(ietf::ParameterBytes::Path, path.clone().into_bytes());
362		}
363		if let Some(authority) = &self.setup_authority {
364			parameters.set_bytes(ietf::ParameterBytes::Authority, authority.clone().into_bytes());
365		}
366		ietf::solicit::into_setup(&mut parameters, ietf_encoding);
367		ietf::hidden::into_setup(&mut parameters, ietf_encoding);
368		let parameters = parameters.encode_bytes(ietf_encoding)?;
369
370		let client = setup::Client {
371			versions: supported.clone().into(),
372			parameters,
373		};
374
375		stream.writer.encode(&client).await?;
376
377		let mut server: setup::Server = stream.reader.decode().await?;
378
379		let version = supported
380			.iter()
381			.find(|v| coding::Version::from(**v) == server.version)
382			.copied()
383			.ok_or(Error::Version)?;
384
385		let (recv_bw, protocol, goaway, estimate) = match version {
386			Version::Lite(v) => {
387				let stream = stream.with_version(v);
388				let start = lite::start(lite::Config {
389					runtime: runtime.clone(),
390					session: session.clone(),
391					setup_stream: Some(stream),
392					publish: publish.clone(),
393					subscribe: subscribe.clone(),
394					peer_hop: Some(peer_hop),
395					version: v,
396					// This path only handles versions negotiated via the bidi SETUP exchange
397					// (pre-lite-05), which have no Setup Stream.
398					our_setup: lite::Setup::default(),
399					peer_setup: None,
400				})?;
401
402				(
403					start.recv_bandwidth,
404					crate::driver::Protocol::Lite(Box::new(start.driver)),
405					start.goaway,
406					Some(start.estimate),
407				)
408			}
409			Version::Ietf(v) => {
410				// Decode the parameters to get the initial request ID and what the server
411				// requires of us.
412				let parameters = ietf::Parameters::decode(&mut server.parameters, v)?;
413				let request_id_max = parameters
414					.get_varint(ietf::ParameterVarInt::MaxRequestId)
415					.map(ietf::RequestId);
416				let peer_declared = ietf::peer::Peer {
417					solicit: ietf::solicit::from_setup(&parameters, v)?,
418					hidden: ietf::hidden::from_setup(&parameters, v),
419					..Default::default()
420				};
421
422				let stream = stream.with_version(v);
423				// Draft 14-16: the path rode in the bidi SETUP above, not the uni one.
424				let (protocol, goaway) = ietf::start(ietf::Config {
425					runtime: runtime.clone(),
426					session: session.clone(),
427					setup: Some(stream),
428					request_id_max,
429					client: true,
430					publish: publish.clone(),
431					subscribe: subscribe.clone(),
432					peer_hop: Some(peer_hop),
433					cost: self.cost,
434					version: v,
435					path: None,
436					authority: None,
437					peer_setup_stream: None,
438					peer_declared: Some(peer_declared),
439					early_unis: Vec::new(),
440				})?;
441				(None, crate::driver::Protocol::Ietf(protocol), goaway, None)
442			}
443		};
444
445		Ok(Session::new(
446			runtime, session, version, recv_bw, protocol, goaway, estimate,
447		))
448	}
449}
450
451#[cfg(test)]
452mod tests {
453	use super::*;
454	use crate::model::ProduceTest;
455	use std::{
456		collections::VecDeque,
457		sync::{Arc, Mutex},
458	};
459
460	use std::task::{Context, Poll};
461
462	use crate::SessionError;
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		// Per-clone, so each pending poll_closed keeps its own registration live.
487		park: kio::Park,
488	}
489
490	#[derive(Default)]
491	struct FakeSessionState {
492		protocol: Option<&'static str>,
493		control_stream: Mutex<Option<(FakeSendStream, FakeRecvStream)>>,
494		close_events: Mutex<Vec<(u32, String)>>,
495		closed: kio::Fan,
496		control_writes: Arc<Mutex<Vec<u8>>>,
497		send_rate: Mutex<Option<u64>>,
498		bytes_sent: Mutex<Option<u64>>,
499	}
500
501	impl FakeSession {
502		fn new(protocol: Option<&'static str>, server_control_bytes: Vec<u8>) -> Self {
503			let writes = Arc::new(Mutex::new(Vec::new()));
504			let send = FakeSendStream { writes: writes.clone() };
505			let recv = FakeRecvStream {
506				data: VecDeque::from(server_control_bytes),
507			};
508			let state = FakeSessionState {
509				protocol,
510				control_stream: Mutex::new(Some((send, recv))),
511				close_events: Mutex::new(Vec::new()),
512				closed: kio::Fan::default(),
513				control_writes: writes,
514				send_rate: Mutex::new(None),
515				bytes_sent: Mutex::new(None),
516			};
517			Self {
518				state: Arc::new(state),
519				park: kio::Park::default(),
520			}
521		}
522
523		fn set_send_rate(&self, rate: Option<u64>) {
524			*self.state.send_rate.lock().unwrap() = rate;
525		}
526
527		fn set_bytes_sent(&self, bytes: Option<u64>) {
528			*self.state.bytes_sent.lock().unwrap() = bytes;
529		}
530
531		fn control_writes(&self) -> Vec<u8> {
532			self.state.control_writes.lock().unwrap().clone()
533		}
534
535		async fn wait_for_first_close(&self) -> (u32, String) {
536			kio::wait(|waiter| {
537				self.state.closed.register(waiter);
538				match self.state.close_events.lock().unwrap().first().cloned() {
539					Some(close) => std::task::Poll::Ready(close),
540					None => std::task::Poll::Pending,
541				}
542			})
543			.await
544		}
545	}
546
547	impl web_transport_trait::poll::Session for FakeSession {
548		type SendStream = FakeSendStream;
549		type RecvStream = FakeRecvStream;
550		type Error = FakeError;
551
552		fn poll_accept_uni(&mut self, _cx: &mut Context<'_>) -> Poll<Result<Self::RecvStream, Self::Error>> {
553			Poll::Pending
554		}
555
556		fn poll_accept_bi(
557			&mut self,
558			_cx: &mut Context<'_>,
559		) -> Poll<Result<(Self::SendStream, Self::RecvStream), Self::Error>> {
560			Poll::Pending
561		}
562
563		fn poll_open_bi(
564			&mut self,
565			_cx: &mut Context<'_>,
566		) -> Poll<Result<(Self::SendStream, Self::RecvStream), Self::Error>> {
567			Poll::Ready(self.state.control_stream.lock().unwrap().take().ok_or(FakeError))
568		}
569
570		fn poll_open_uni(&mut self, _cx: &mut Context<'_>) -> Poll<Result<Self::SendStream, Self::Error>> {
571			Poll::Pending
572		}
573
574		fn poll_send_datagram(&mut self, _cx: &mut Context<'_>, _payload: &[u8]) -> Poll<Result<(), Self::Error>> {
575			Poll::Ready(Ok(()))
576		}
577
578		fn poll_recv_datagram(&mut self, _cx: &mut Context<'_>) -> Poll<Result<Bytes, Self::Error>> {
579			Poll::Pending
580		}
581
582		fn max_datagram_size(&self) -> usize {
583			1200
584		}
585
586		fn protocol(&self) -> Option<&str> {
587			self.state.protocol
588		}
589
590		fn close(&mut self, code: u32, reason: &str) {
591			self.state.close_events.lock().unwrap().push((code, reason.to_string()));
592			self.state.closed.wake();
593		}
594
595		fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Self::Error> {
596			// Register before checking so a close racing this poll still wakes it.
597			self.state.closed.register(self.park.hold(cx));
598			match self.state.close_events.lock().unwrap().is_empty() {
599				false => Poll::Ready(FakeError),
600				true => Poll::Pending,
601			}
602		}
603
604		fn stats(&self) -> impl web_transport_trait::Stats {
605			FakeStats {
606				send_rate: *self.state.send_rate.lock().unwrap(),
607				bytes_sent: *self.state.bytes_sent.lock().unwrap(),
608			}
609		}
610	}
611
612	struct FakeStats {
613		send_rate: Option<u64>,
614		bytes_sent: Option<u64>,
615	}
616
617	impl web_transport_trait::Stats for FakeStats {
618		fn estimated_send_rate(&self) -> Option<u64> {
619			self.send_rate
620		}
621
622		fn bytes_sent(&self) -> Option<u64> {
623			self.bytes_sent
624		}
625	}
626
627	#[derive(Clone, Default)]
628	struct FakeSendStream {
629		writes: Arc<Mutex<Vec<u8>>>,
630	}
631
632	impl web_transport_trait::poll::SendStream for FakeSendStream {
633		type Error = FakeError;
634
635		fn poll_write(&mut self, _cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Self::Error>> {
636			self.writes.lock().unwrap().put_slice(buf);
637			Poll::Ready(Ok(buf.len()))
638		}
639
640		fn set_priority(&mut self, _order: u8) {}
641
642		fn finish(&mut self) -> Result<(), Self::Error> {
643			Ok(())
644		}
645
646		fn reset(&mut self, _code: u32) {}
647
648		fn poll_closed(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
649			Poll::Ready(Ok(()))
650		}
651	}
652
653	struct FakeRecvStream {
654		data: VecDeque<u8>,
655	}
656
657	impl web_transport_trait::poll::RecvStream for FakeRecvStream {
658		type Error = FakeError;
659
660		fn poll_read(&mut self, _cx: &mut Context<'_>, dst: &mut [u8]) -> Poll<Result<Option<usize>, Self::Error>> {
661			if self.data.is_empty() {
662				return Poll::Ready(Ok(None));
663			}
664
665			let size = dst.len().min(self.data.len());
666			for slot in dst.iter_mut().take(size) {
667				*slot = self.data.pop_front().unwrap();
668			}
669			Poll::Ready(Ok(Some(size)))
670		}
671
672		fn stop(&mut self, _code: u32) {}
673
674		fn poll_closed(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
675			Poll::Ready(Ok(()))
676		}
677	}
678
679	fn mock_server_setup(negotiated: Version) -> Vec<u8> {
680		let mut encoded = Vec::new();
681		let server = setup::Server {
682			version: negotiated.into(),
683			parameters: Bytes::new(),
684		};
685		server
686			.encode(&mut encoded, Version::Ietf(ietf::Version::Draft14))
687			.unwrap();
688
689		// Add a setup-stream SessionInfo frame using the negotiated Lite version.
690		let info = lite::SessionInfo { bitrate: Some(1) };
691		let lite_v = lite::Version::try_from(negotiated).unwrap();
692		info.encode(&mut encoded, lite_v).unwrap();
693
694		encoded
695	}
696
697	async fn run_alpn_lite_fallback_case(protocol: Option<&'static str>) {
698		let fake = FakeSession::new(protocol, mock_server_setup(Version::Lite(lite::Version::Lite01)));
699		let client = Client::new().with_versions(
700			[
701				Version::Lite(lite::Version::Lite03),
702				Version::Lite(lite::Version::Lite02),
703				Version::Lite(lite::Version::Lite01),
704				Version::Ietf(ietf::Version::Draft14),
705			]
706			.into(),
707		);
708
709		// Start the returned driver after the handshake completes.
710		let (_session, driver) = client
711			.connect(tokio::time::Instant::now().into_std(), fake.clone())
712			.await
713			.unwrap();
714		tokio::spawn(crate::time::run(driver));
715
716		// Verify the client setup was encoded using Draft14 framing (ALPN_LITE fallback path).
717		let mut setup_bytes = Bytes::from(fake.control_writes());
718		let setup = setup::Client::decode(&mut setup_bytes, Version::Ietf(ietf::Version::Draft14)).unwrap();
719		let advertised: Vec<Version> = setup.versions.iter().map(|v| Version::try_from(*v).unwrap()).collect();
720		assert_eq!(
721			advertised,
722			vec![
723				Version::Lite(lite::Version::Lite02),
724				Version::Lite(lite::Version::Lite01),
725				Version::Ietf(ietf::Version::Draft14),
726			]
727		);
728
729		// The first close comes from the lite connection driver.
730		// Any non-Version error here means SessionInfo decoded successfully
731		// after set_version(). This test cares about the SETUP framing
732		// fallback, not the specific close code. Cancel is what we'd see
733		// with no origin; a protocol violation (or similar) is what an
734		// auto-created origin's first interaction with a Lite01 peer trips.
735		let (code, _) = fake.wait_for_first_close().await;
736		// Session closes encode through the session registry, so compare against that one.
737		assert_ne!(code, SessionError::Version.to_code(), "SessionInfo failed to decode");
738	}
739
740	/// `connect` must not depend on the peer answering. A peer that opens the announce
741	/// stream and then says nothing (or promises a count it never delivers) used to hold
742	/// `connect` for the life of the session, since it waited for the initial announce
743	/// set. Resolving a path you need is `routed`'s job, which waits for
744	/// that path rather than for the peer to finish talking.
745	#[tokio::test(start_paused = true)]
746	async fn connect_does_not_wait_for_the_peer_to_announce() {
747		// Serves bidi streams, so the announce stream opens, and never answers on them.
748		let gate = kio::Producer::new(true);
749		let transport = crate::lite::test_transport::SinkSession::gated_bi(gate.consume())
750			.with_protocol(crate::version::ALPN_LITE_05);
751
752		// A subscribe origin is what makes the client open an announce stream at all.
753		let origin = crate::origin::Config::new(crate::Hop::new(1).unwrap()).produce();
754		let client = Client::new()
755			.with_versions([Version::Lite(lite::Version::Lite05)].into())
756			.with_subscriber(origin);
757
758		// Paused time auto-advances while every task is idle, so a `connect` that waits
759		// on the silent peer trips this rather than hanging the suite.
760		let (_session, _driver) = tokio::time::timeout(
761			std::time::Duration::from_secs(30),
762			client.connect(tokio::time::Instant::now().into_std(), transport),
763		)
764		.await
765		.expect("connect waited on a peer that never announced")
766		.expect("connect failed");
767	}
768
769	/// The client SETUP on the bidi control stream (the pre-draft-17 framing) carries the
770	/// AUTHORITY next to the PATH.
771	#[tokio::test(start_paused = true)]
772	async fn draft14_setup_carries_the_authority() {
773		let fake = FakeSession::new(Some(ALPN_LITE), mock_server_setup(Version::Lite(lite::Version::Lite01)));
774		let client = Client::new()
775			.with_versions(
776				[
777					Version::Lite(lite::Version::Lite01),
778					Version::Ietf(ietf::Version::Draft14),
779				]
780				.into(),
781			)
782			.with_path("/anon")
783			.with_authority("relay.example.com:4443");
784
785		let (_session, driver) = client
786			.connect(tokio::time::Instant::now().into_std(), fake.clone())
787			.await
788			.unwrap();
789		tokio::spawn(crate::time::run(driver));
790
791		let mut setup_bytes = Bytes::from(fake.control_writes());
792		let setup = setup::Client::decode(&mut setup_bytes, Version::Ietf(ietf::Version::Draft14)).unwrap();
793		let mut parameters = setup.parameters;
794		let parameters = ietf::Parameters::decode(&mut parameters, ietf::Version::Draft14).unwrap();
795		assert_eq!(
796			parameters.get_bytes(ietf::ParameterBytes::Authority),
797			Some(b"relay.example.com:4443".as_ref())
798		);
799		assert_eq!(
800			parameters.get_bytes(ietf::ParameterBytes::Path),
801			Some(b"/anon".as_ref())
802		);
803	}
804
805	#[tokio::test(start_paused = true)]
806	async fn alpn_lite_falls_back_to_draft14_and_switches_version_post_setup() {
807		run_alpn_lite_fallback_case(Some(ALPN_LITE)).await;
808	}
809
810	#[tokio::test(start_paused = true)]
811	async fn no_alpn_falls_back_to_draft14_and_switches_version_post_setup() {
812		run_alpn_lite_fallback_case(None).await;
813	}
814
815	// No executor is running: only explicitly polling the driver may process
816	// a session close, and the driver must not retain a session handle.
817	#[test]
818	fn driver_is_caller_polled_and_holds_no_session() {
819		let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
820		let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
821
822		let runtime = crate::runtime::Test::new();
823		let (session, mut driver) = futures::executor::block_on(client.connect(runtime.now(), fake.clone())).unwrap();
824		assert_eq!(session.version(), Version::Lite(lite::Version::Lite04));
825
826		// Construction leaves the driver idle until the caller polls it.
827		assert!(driver.poll(runtime.now(), &kio::Waiter::noop()).is_ok());
828
829		// The caller drops their only session clone; the machine observes the
830		// last handle going away and closes the transport.
831		drop(session);
832		assert!(fake.state.close_events.lock().unwrap().is_empty());
833		let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
834		assert_eq!(
835			fake.state.close_events.lock().unwrap()[0].0,
836			SessionError::Cancel.to_code()
837		);
838	}
839
840	// Clones share the connection: the transport closes on the LAST drop, and
841	// abort() closes it explicitly (first close wins). Both are relayed through
842	// the machine, so each takes a tick to land.
843	#[test]
844	fn session_clones_share_the_close() {
845		let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
846		let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
847
848		let runtime = crate::runtime::Test::new();
849		let (session, mut driver) = futures::executor::block_on(client.connect(runtime.now(), fake.clone())).unwrap();
850		let clone = session.clone();
851
852		// One clone dropping does nothing while another is alive.
853		drop(session);
854		assert!(fake.state.close_events.lock().unwrap().is_empty());
855		let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
856		assert!(fake.state.close_events.lock().unwrap().is_empty());
857
858		clone.abort(Error::Cancel);
859		let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
860		assert_eq!(
861			fake.state.close_events.lock().unwrap()[0].0,
862			SessionError::Cancel.to_code()
863		);
864
865		// And the machine publishes the transport's terminal error, which is
866		// what `closed()` reports.
867		let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
868		futures::executor::block_on(clone.closed());
869
870		// The final drop requests no second close: the handle-side close is once.
871		let closes = fake.state.close_events.lock().unwrap().len();
872		drop(clone);
873		let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
874		assert_eq!(fake.state.close_events.lock().unwrap().len(), closes);
875	}
876
877	// Dropping the driver instead of running it tears the session
878	// down: the machine was the only transport holder, and `closed()` resolves
879	// rather than parking forever on a machine nobody polls.
880	#[test]
881	fn dropped_driver_resolves_closed() {
882		let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
883		let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
884
885		let runtime = crate::runtime::Test::new();
886		let (session, driver) = futures::executor::block_on(client.connect(runtime.now(), fake.clone())).unwrap();
887
888		drop(driver);
889		assert!(matches!(futures::executor::block_on(session.closed()), Error::Cancel));
890	}
891
892	/// A transport made deliberately `!Send` by an `Rc` marker on the session and
893	/// both stream types: compiling at all is the point, proving the lite path
894	/// never demands thread mobility of any transport piece.
895	#[derive(Clone)]
896	struct LocalSession {
897		inner: FakeSession,
898		_local: std::rc::Rc<()>,
899	}
900
901	struct LocalSend {
902		inner: FakeSendStream,
903		_local: std::rc::Rc<()>,
904	}
905
906	struct LocalRecv {
907		inner: FakeRecvStream,
908		_local: std::rc::Rc<()>,
909	}
910
911	impl web_transport_trait::poll::Session for LocalSession {
912		type SendStream = LocalSend;
913		type RecvStream = LocalRecv;
914		type Error = FakeError;
915
916		fn poll_accept_uni(&mut self, cx: &mut Context<'_>) -> Poll<Result<Self::RecvStream, Self::Error>> {
917			self.inner.poll_accept_uni(cx).map_ok(|stream| LocalRecv {
918				inner: stream,
919				_local: self._local.clone(),
920			})
921		}
922
923		fn poll_accept_bi(
924			&mut self,
925			cx: &mut Context<'_>,
926		) -> Poll<Result<(Self::SendStream, Self::RecvStream), Self::Error>> {
927			self.inner.poll_accept_bi(cx).map_ok(|(send, recv)| {
928				(
929					LocalSend {
930						inner: send,
931						_local: self._local.clone(),
932					},
933					LocalRecv {
934						inner: recv,
935						_local: self._local.clone(),
936					},
937				)
938			})
939		}
940
941		fn poll_open_bi(
942			&mut self,
943			cx: &mut Context<'_>,
944		) -> Poll<Result<(Self::SendStream, Self::RecvStream), Self::Error>> {
945			self.inner.poll_open_bi(cx).map_ok(|(send, recv)| {
946				(
947					LocalSend {
948						inner: send,
949						_local: self._local.clone(),
950					},
951					LocalRecv {
952						inner: recv,
953						_local: self._local.clone(),
954					},
955				)
956			})
957		}
958
959		fn poll_open_uni(&mut self, cx: &mut Context<'_>) -> Poll<Result<Self::SendStream, Self::Error>> {
960			self.inner.poll_open_uni(cx).map_ok(|stream| LocalSend {
961				inner: stream,
962				_local: self._local.clone(),
963			})
964		}
965
966		fn poll_send_datagram(&mut self, cx: &mut Context<'_>, payload: &[u8]) -> Poll<Result<(), Self::Error>> {
967			self.inner.poll_send_datagram(cx, payload)
968		}
969
970		fn poll_recv_datagram(&mut self, cx: &mut Context<'_>) -> Poll<Result<Bytes, Self::Error>> {
971			self.inner.poll_recv_datagram(cx)
972		}
973
974		fn max_datagram_size(&self) -> usize {
975			self.inner.max_datagram_size()
976		}
977
978		fn protocol(&self) -> Option<&str> {
979			self.inner.protocol()
980		}
981
982		fn close(&mut self, code: u32, reason: &str) {
983			self.inner.close(code, reason);
984		}
985
986		fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Self::Error> {
987			self.inner.poll_closed(cx)
988		}
989
990		fn stats(&self) -> impl web_transport_trait::Stats {
991			self.inner.stats()
992		}
993	}
994
995	impl web_transport_trait::poll::SendStream for LocalSend {
996		type Error = FakeError;
997
998		fn poll_write(&mut self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize, Self::Error>> {
999			self.inner.poll_write(cx, buf)
1000		}
1001
1002		fn set_priority(&mut self, order: u8) {
1003			self.inner.set_priority(order);
1004		}
1005
1006		fn finish(&mut self) -> Result<(), Self::Error> {
1007			self.inner.finish()
1008		}
1009
1010		fn reset(&mut self, code: u32) {
1011			self.inner.reset(code);
1012		}
1013
1014		fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
1015			web_transport_trait::poll::SendStream::poll_closed(&mut self.inner, cx)
1016		}
1017	}
1018
1019	impl web_transport_trait::poll::RecvStream for LocalRecv {
1020		type Error = FakeError;
1021
1022		fn poll_read(&mut self, cx: &mut Context<'_>, dst: &mut [u8]) -> Poll<Result<Option<usize>, Self::Error>> {
1023			self.inner.poll_read(cx, dst)
1024		}
1025
1026		fn stop(&mut self, code: u32) {
1027			self.inner.stop(code);
1028		}
1029
1030		fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
1031			web_transport_trait::poll::RecvStream::poll_closed(&mut self.inner, cx)
1032		}
1033	}
1034
1035	// The point of the lite-only entry: a `!Send` transport yields a `!Send`
1036	// driver polled by its caller, while the severed Session handle stays
1037	// Send + Sync. Compiling is most of the assertion; the rest checks the
1038	// machine still relays the close.
1039	#[test]
1040	fn connect_lite_over_a_send_less_transport() {
1041		let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1042		let local = LocalSession {
1043			inner: fake.clone(),
1044			_local: std::rc::Rc::new(()),
1045		};
1046		let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1047
1048		let runtime = crate::runtime::Test::new();
1049		let (session, mut driver) = futures::executor::block_on(client.connect_lite(runtime.now(), local)).unwrap();
1050		assert!(driver.poll(runtime.now(), &kio::Waiter::noop()).is_ok());
1051
1052		fn assert_send_sync<T: Send + Sync>(_: &T) {}
1053		assert_send_sync(&session);
1054
1055		session.abort(Error::Cancel);
1056		let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
1057		assert_eq!(
1058			fake.state.close_events.lock().unwrap()[0].0,
1059			SessionError::Cancel.to_code()
1060		);
1061	}
1062
1063	// The server-side twin: a `!Send` transport accepts a lite session whose
1064	// driver the caller polls directly.
1065	#[test]
1066	fn accept_lite_over_a_send_less_transport() {
1067		let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1068		let local = LocalSession {
1069			inner: fake.clone(),
1070			_local: std::rc::Rc::new(()),
1071		};
1072		let server = crate::Server::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1073
1074		let runtime = crate::runtime::Test::new();
1075		let (session, mut driver) = futures::executor::block_on(server.accept_lite(runtime.now(), local)).unwrap();
1076		assert_eq!(session.version(), Version::Lite(lite::Version::Lite04));
1077		assert!(driver.poll(runtime.now(), &kio::Waiter::noop()).is_ok());
1078
1079		drop(session);
1080		assert!(fake.state.close_events.lock().unwrap().is_empty());
1081		let _ = driver.poll(runtime.now(), &kio::Waiter::noop());
1082		assert_eq!(
1083			fake.state.close_events.lock().unwrap()[0].0,
1084			SessionError::Cancel.to_code()
1085		);
1086	}
1087
1088	// The lite-only entry refuses everything that still needs the boxed ietf
1089	// driver, instead of silently negotiating it.
1090	#[test]
1091	fn connect_lite_refuses_ietf_alpns() {
1092		let fake = FakeSession::new(Some(ALPN_19), Vec::new());
1093		let local = LocalSession {
1094			inner: fake,
1095			_local: std::rc::Rc::new(()),
1096		};
1097		let client = Client::new();
1098		let runtime = crate::runtime::Test::new();
1099		let result = futures::executor::block_on(client.connect_lite(runtime.now(), local));
1100		assert!(matches!(result, Err(Error::Version)));
1101	}
1102
1103	// `stats()` reads the machine's latest sample and primes the sampler, so a
1104	// periodic poller observes fresh counters without consuming the bandwidth
1105	// channel.
1106	#[tokio::test(start_paused = true)]
1107	async fn stats_reads_prime_the_sampler() {
1108		let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1109		fake.set_send_rate(Some(1_000_000));
1110
1111		let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1112		let (session, driver) = client
1113			.connect(tokio::time::Instant::now().into_std(), fake.clone())
1114			.await
1115			.unwrap();
1116		tokio::spawn(crate::time::run(driver));
1117
1118		// The construction-time snapshot, before the machine sampled anything.
1119		assert_eq!(
1120			session.stats().estimated_send_rate,
1121			Some(crate::bandwidth::Rate::from_bps(1_000_000))
1122		);
1123
1124		// That read was demand: the machine keeps sampling while stats are read,
1125		// so the new rate shows up within an interval (paused time auto-advances).
1126		fake.set_send_rate(Some(2_000_000));
1127		while session.stats().estimated_send_rate != Some(crate::bandwidth::Rate::from_bps(2_000_000)) {
1128			tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1129		}
1130	}
1131
1132	// Sampling stops when the supervisor ends, but `stats()` keeps serving its
1133	// cell, so the last thing the supervisor does is take a final snapshot.
1134	// Without one, "what did that session move?" asked at teardown answers with
1135	// the construction-time snapshot: this backend reports no send rate, so
1136	// there is no bandwidth consumer keeping the sampler ticking, and the test
1137	// never reads stats while the session is live.
1138	#[tokio::test(start_paused = true)]
1139	async fn stats_capture_the_final_counters() {
1140		let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1141		fake.set_send_rate(None);
1142		fake.set_bytes_sent(Some(0));
1143
1144		let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1145		let (session, driver) = client
1146			.connect(tokio::time::Instant::now().into_std(), fake.clone())
1147			.await
1148			.unwrap();
1149		tokio::spawn(crate::time::run(driver));
1150		assert!(
1151			session.send_bandwidth().is_none(),
1152			"no send-rate estimate, so nothing samples on its own"
1153		);
1154
1155		fake.set_bytes_sent(Some(4242));
1156
1157		session.abort(Error::Cancel);
1158		session.closed().await;
1159
1160		assert_eq!(
1161			session.stats().bytes_sent,
1162			Some(4242),
1163			"the closing snapshot must carry the session's final counters"
1164		);
1165	}
1166
1167	// The send-bandwidth sampler lives inside the driver: it samples as soon as a
1168	// consumer exists and keeps sampling on its interval. Paused tokio time makes
1169	// the interval fire deterministically.
1170	#[tokio::test(start_paused = true)]
1171	async fn send_bandwidth_samples_while_the_driver_runs() {
1172		let fake = FakeSession::new(Some(ALPN_LITE_04), Vec::new());
1173		fake.set_send_rate(Some(1_000_000));
1174
1175		let client = Client::new().with_versions(Version::Lite(lite::Version::Lite04).into());
1176		let (session, driver) = client
1177			.connect(tokio::time::Instant::now().into_std(), fake.clone())
1178			.await
1179			.unwrap();
1180		tokio::spawn(crate::time::run(driver));
1181
1182		let mut bandwidth = session.send_bandwidth().expect("backend reports an estimate");
1183		assert_eq!(
1184			bandwidth.changed().await.unwrap(),
1185			Some(crate::bandwidth::Rate::from_bps(1_000_000))
1186		);
1187
1188		// A later change is picked up by the next interval tick.
1189		fake.set_send_rate(Some(2_000_000));
1190		assert_eq!(
1191			bandwidth.changed().await.unwrap(),
1192			Some(crate::bandwidth::Rate::from_bps(2_000_000))
1193		);
1194	}
1195}