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