moq_net/lite/

setup.rs

//! The lite-05+ SETUP message: each endpoint advertises its capabilities once, as
//! the sole message on a unidirectional Setup Stream, then closes it.

use crate::coding::*;

use super::{Message, Parameters, Version};

/// Setup Parameter id for the Probe capability level.
const PARAM_PROBE: u64 = 0x1;
/// Setup Parameter id for the request Path (client-only, URI-less transports).
const PARAM_PATH: u64 = 0x2;
/// Setup Parameter id for the client's intended [`Role`] (client-only).
const PARAM_ROLE: u64 = 0x3;
/// Setup Parameter id for the link cost the dialer assigns to this connection.
const PARAM_COST: u64 = 0x4;
/// Setup Parameter id for the endpoint's Hop ID.
const PARAM_HOP: u64 = 0x5;

/// The cost of crossing a link that neither end priced.
///
/// One, so a mesh that configures no costs accumulates a route cost equal to the
/// hop count and ranks routes exactly as pre-lite-06 shortest-path routing did. Pricing
/// a link at 0 makes it free (a sibling in the same datacenter); pricing it higher
/// makes it a last resort (a metered backbone).
pub const DEFAULT_COST: u64 = 1;

/// The probe capability an endpoint advertises in SETUP.
///
/// Monotonic: a higher level implies every lower one. An unknown (future) value
/// decodes as the highest level we understand, so a peer that gains a new level is
/// treated as at least [`Increase`](Self::Increase).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)]
pub enum ProbeLevel {
	/// No probing. Equivalent to omitting the parameter.
	#[default]
	None,
	/// The publisher can measure and periodically report at least one of the PROBE
	/// metrics: its estimated bitrate, its round-trip time, or both. Either may be
	/// unknown in any given report, since the two are independent on the wire.
	Report,
	/// The publisher can additionally pad the connection (or send redundant data).
	Increase,
}

impl ProbeLevel {
	/// The level to advertise for `session`, from what its transport actually exposes.
	///
	/// [`Report`](Self::Report) claims the publisher can measure and periodically
	/// report. A transport that exposes neither a send-rate estimate nor an RTT can
	/// honour neither, and the draft requires such a publisher to reset any Probe
	/// Stream a subscriber opens. Advertising [`None`](Self::None) instead stops the
	/// subscriber opening one at all.
	///
	/// Both metrics are sampled rather than declared, so this only works for a
	/// transport whose figures exist by the time the session starts. QUIC and TCP
	/// both qualify: their RTT comes from the handshake, which has already happened.
	pub fn detect<S: crate::transport::poll::Session>(session: &S) -> Self {
		use web_transport_trait::Stats as _;
		let stats = session.stats();
		match stats.estimated_send_rate().is_some() || stats.rtt().is_some() {
			true => Self::Report,
			false => Self::None,
		}
	}

	/// Map the wire value to a level, saturating unknown values to [`Increase`](Self::Increase).
	fn from_code(code: u64) -> Self {
		match code {
			0 => Self::None,
			1 => Self::Report,
			_ => Self::Increase,
		}
	}

	/// The wire value for this level.
	fn to_code(self) -> u64 {
		match self {
			Self::None => 0,
			Self::Report => 1,
			Self::Increase => 2,
		}
	}
}

/// The single direction a client intends to use the session for.
///
/// A client advertises this in its SETUP so the server can reject a token that lacks
/// the matching scope during the handshake, instead of accepting a connection that
/// then silently carries no media (a subscribe-only token used to publish, or vice
/// versa). It only ever narrows what the server grants, so it is not a security
/// boundary: the server still enforces the token's scope regardless.
///
/// A session is bidirectional by default, which the wire says by omitting the
/// parameter. `Option<Role>` mirrors that: `None` is the default, and it's also what
/// a client that predates the parameter decodes to.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum Role {
	/// The client will publish tracks (ingest); the server must consume.
	Publisher,
	/// The client will subscribe to tracks (egress); the server must publish.
	Subscriber,
}

impl Role {
	/// Map the wire value to a role. `0` and any unrecognized future value are `None`
	/// (bidirectional): the draft requires a receiver that does not recognize the value
	/// to treat it as both directions, so a newer client can't break an older server (it
	/// just loses the early reject and defers fully to the token's scope).
	fn from_code(code: u64) -> Option<Self> {
		match code {
			1 => Some(Role::Publisher),
			2 => Some(Role::Subscriber),
			_ => None,
		}
	}

	/// The wire value for this role.
	fn to_code(self) -> u64 {
		match self {
			Role::Publisher => 1,
			Role::Subscriber => 2,
		}
	}

	/// Derive the advertised role from which origins a client wired up: publish-only is
	/// a [`Publisher`](Role::Publisher), consume-only a [`Subscriber`](Role::Subscriber),
	/// and both (or neither) advertises nothing. This keeps the advertised role from
	/// drifting away from what the session actually does.
	pub(crate) fn from_origins(publishes: bool, consumes: bool) -> Option<Self> {
		match (publishes, consumes) {
			(true, false) => Some(Role::Publisher),
			(false, true) => Some(Role::Subscriber),
			_ => None,
		}
	}

	/// Lowercase label for this role (`"publisher"` / `"subscriber"`).
	pub fn as_str(self) -> &'static str {
		match self {
			Role::Publisher => "publisher",
			Role::Subscriber => "subscriber",
		}
	}
}

impl std::fmt::Display for Role {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.write_str(self.as_str())
	}
}

/// The SETUP message, sent once per endpoint on the unidirectional Setup Stream.
///
/// lite-05+ only. The two endpoints' SETUP messages are independent: neither side
/// blocks on the peer's before opening other streams, but a stream whose encoding
/// depends on a negotiated capability (e.g. PROBE) must wait for it.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Setup {
	/// The probe capability this endpoint supports. [`ProbeLevel::None`] when absent.
	pub probe: ProbeLevel,
	/// The request path, for transports that carry no request URI (native QUIC,
	/// qmux over TCP/TLS, unix sockets), with `?` and the URI query appended when
	/// there is one. Sent only by the client; a server never sends one and a relay
	/// never forwards it. `None` on URI-carrying bindings, where it would be a
	/// protocol violation. An empty path means the same thing as `None`; both are
	/// on the wire so a client need not special-case the root.
	pub path: Option<String>,
	/// The single direction the client intends to use, or `None` for a bidirectional
	/// session. `None` is sent as the absence of the parameter, which is also how a
	/// client that predates the parameter decodes.
	pub role: Option<Role>,
	/// What subscribing from this endpoint costs (lite-06+), added by the peer to the
	/// route cost of every announcement we forward it.
	///
	/// Directional: it prices the sender's own egress, so both ends declare their own
	/// and the two need not match. `None` means the default cost of 1.
	pub cost: Option<u64>,
	/// This endpoint's Hop ID, the identity it stamps onto forwarded
	/// announcements. The peer uses it to serve this endpoint's subscriptions from
	/// a route that does not flow through it (the same split horizon the announce
	/// filter applies). `None` when the endpoint has no meaningful identity (a
	/// leaf that never forwards); a wire value of 0 decodes as `None`.
	pub hop: Option<crate::Hop>,
}

impl Message for Setup {
	const CONTROL: Option<crate::stats::Control> = Some(crate::stats::Control::Session);
	const MAX_SIZE: usize = crate::setup::MAX_SETUP_SIZE;

	fn decode_msg<R: bytes::Buf>(r: &mut R, version: Version) -> Result<Self, DecodeError> {
		if !version.has_setup_stream() {
			return Err(DecodeError::Version);
		}

		let params = Parameters::decode(r, version)?;
		let probe = params
			.get_varint(PARAM_PROBE, version)?
			.map(ProbeLevel::from_code)
			.unwrap_or_default();
		let path = match params.get_bytes(PARAM_PATH) {
			Some(bytes) => Some(
				std::str::from_utf8(bytes)
					.map_err(|_| DecodeError::InvalidValue)?
					.to_string(),
			),
			None => None,
		};
		let role = params.get_varint(PARAM_ROLE, version)?.and_then(Role::from_code);
		let cost = params.get_varint(PARAM_COST, version)?;
		// 0 is legal on the wire but carries no identity (it can't be excluded),
		// so it decodes as "not declared" rather than an error.
		let hop = params
			.get_varint(PARAM_HOP, version)?
			.and_then(|id| crate::Hop::new(id).ok());

		Ok(Self {
			probe,
			path,
			role,
			cost,
			hop,
		})
	}

	fn encode_msg<W: bytes::BufMut>(&self, w: &mut W, version: Version) -> Result<(), EncodeError> {
		if !version.has_setup_stream() {
			return Err(EncodeError::Version);
		}

		let mut params = Parameters::default();
		// None is the wire default, so omit it to keep the message empty when nothing is set.
		if self.probe != ProbeLevel::None {
			params.set_varint(PARAM_PROBE, self.probe.to_code(), version)?;
		}
		if let Some(path) = &self.path {
			params.set_bytes(PARAM_PATH, path.as_bytes().to_vec());
		}
		// Bidirectional is the wire default (absence of the parameter), so only a
		// directional role is encoded.
		if let Some(role) = self.role {
			params.set_varint(PARAM_ROLE, role.to_code(), version)?;
		}
		if let Some(cost) = self.cost {
			params.set_varint(PARAM_COST, cost, version)?;
		}
		if let Some(hop) = self.hop {
			params.set_varint(PARAM_HOP, hop.id(), version)?;
		}

		params.encode(w, version)
	}
}

/// Shared slot for the peer's SETUP, written once when its Setup stream is read.
///
/// Streams whose encoding depends on a negotiated capability (e.g. the PROBE
/// stream) wait on this before deciding what to do. Cheap to clone: every handle
/// shares the same slot.
#[derive(Clone, Default)]
pub(crate) struct PeerSetup(kio::Shared<PeerSetupState>);

#[derive(Default)]
struct PeerSetupState {
	seen: bool,
	setup: Option<Setup>,
}

impl PeerSetup {
	/// Record the peer's SETUP.
	pub fn set(&self, setup: Setup) {
		let mut state = self.0.lock();
		state.seen = true;
		state.setup = Some(setup);
	}

	/// Claim the Setup Stream before decoding its body, refusing any second stream.
	pub fn claim(&self) -> Result<(), crate::Error> {
		let mut state = self.0.lock();
		if state.seen {
			return Err(crate::Error::ProtocolViolation);
		}
		state.seen = true;
		Ok(())
	}

	/// Poll for the peer's advertised probe level, waiting until its SETUP arrives.
	pub fn poll_probe_level(&self, waiter: &kio::Waiter) -> std::task::Poll<ProbeLevel> {
		self.poll_get(waiter, |setup| setup.probe)
	}

	/// Poll for the link cost the peer (the dialing side) declared in its SETUP.
	/// `None` when it declared none, meaning the default cost of 1.
	pub fn poll_cost(&self, waiter: &kio::Waiter) -> std::task::Poll<Option<u64>> {
		self.poll_get(waiter, |setup| setup.cost)
	}

	/// Poll for the [`Hop`](crate::Hop) id the peer declared in its SETUP `Hop`
	/// parameter. `None` when it declared none: a leaf with no identity worth excluding.
	pub fn poll_hop(&self, waiter: &kio::Waiter) -> std::task::Poll<Option<crate::Hop>> {
		self.poll_get(waiter, |setup| setup.hop)
	}

	/// Poll for a field of the peer's SETUP.
	///
	/// The peer MUST send exactly one SETUP, so this resolves once that stream is read.
	/// Pends forever if it never does; the caller is session state, dropped with the
	/// driver.
	fn poll_get<T>(&self, waiter: &kio::Waiter, f: impl FnOnce(&Setup) -> T) -> std::task::Poll<T> {
		let slot = std::task::ready!(self.0.poll(waiter, |setup| {
			if setup.setup.is_some() {
				std::task::Poll::Ready(())
			} else {
				std::task::Poll::Pending
			}
		}));
		std::task::Poll::Ready(f(slot.setup.as_ref().expect("waited for Some")))
	}
}

#[cfg(test)]
mod tests {
	use super::*;

	fn round_trip(msg: &Setup) -> Setup {
		let mut buf = bytes::BytesMut::new();
		msg.encode(&mut buf, Version::Lite05).unwrap();
		let mut slice = &buf[..];
		let got = Setup::decode(&mut slice, Version::Lite05).unwrap();
		assert!(bytes::Buf::remaining(&slice) == 0, "trailing bytes after decode");
		got
	}

	#[test]
	fn empty_round_trip() {
		let msg = Setup::default();
		assert_eq!(round_trip(&msg), msg);
	}

	/// A transport exposing neither metric can't honour a `Report` claim, and the
	/// draft makes such a publisher reset any Probe Stream a subscriber opens. It
	/// must advertise `None` so the subscriber never opens one.
	#[test]
	fn detect_reports_nothing_without_stats() {
		use crate::lite::test_transport::{SinkSession, SinkStats};
		let session = SinkSession::new(Default::default()).with_stats(SinkStats::default());
		assert_eq!(ProbeLevel::detect(&session), ProbeLevel::None);
	}

	/// Either metric alone is enough to report, since the two PROBE fields are
	/// independent on the wire.
	#[test]
	fn detect_reports_with_either_metric() {
		use crate::lite::test_transport::{SinkSession, SinkStats};

		let rtt_only = SinkStats::default().with_rtt(std::time::Duration::from_millis(40));
		let session = SinkSession::new(Default::default()).with_stats(rtt_only);
		assert_eq!(ProbeLevel::detect(&session), ProbeLevel::Report);

		let rate_only = SinkStats::default().with_send_rate(1_000_000);
		let session = SinkSession::new(Default::default()).with_stats(rate_only);
		assert_eq!(ProbeLevel::detect(&session), ProbeLevel::Report);
	}

	#[test]
	fn probe_levels_round_trip() {
		for probe in [ProbeLevel::None, ProbeLevel::Report, ProbeLevel::Increase] {
			let msg = Setup {
				probe,
				..Default::default()
			};
			assert_eq!(round_trip(&msg), msg);
		}
	}

	#[test]
	fn cost_round_trip() {
		// Zero is a meaningful price (a free same-datacenter link), so it must survive
		// the round trip as `Some(0)` rather than collapsing into "unpriced".
		for cost in [None, Some(0), Some(1), Some(7)] {
			let msg = Setup {
				cost,
				..Default::default()
			};
			assert_eq!(round_trip(&msg), msg);
		}
	}

	/// Parameter values follow the session's varint codec, not a fixed one: a value that
	/// takes the two-byte QUIC form on lite-06 fits one leading-ones byte on lite-07.
	#[test]
	fn parameter_values_use_the_version_codec() {
		let msg = Setup {
			cost: Some(100),
			..Default::default()
		};
		for (version, wire) in [
			(Version::Lite06, &[0x05, 0x01, 0x04, 0x02, 0x40, 0x64][..]),
			(Version::Lite07, &[0x04, 0x01, 0x04, 0x01, 0x64][..]),
		] {
			let mut buf = Vec::new();
			msg.encode(&mut buf, version).unwrap();
			assert_eq!(buf, wire, "{version}");
			assert_eq!(Setup::decode(&mut &buf[..], version).unwrap(), msg, "{version}");
		}
	}

	/// Never emit a SETUP our own receiver would refuse.
	#[test]
	fn encode_enforces_the_setup_limit() {
		// Count, id, and a 4-byte length varint precede the path.
		let at_limit = crate::setup::MAX_SETUP_SIZE - 6;
		let msg = Setup {
			path: Some("a".repeat(at_limit)),
			..Default::default()
		};
		assert_eq!(round_trip(&msg), msg);

		let msg = Setup {
			path: Some("a".repeat(at_limit + 1)),
			..Default::default()
		};
		let mut buf = bytes::BytesMut::new();
		assert!(matches!(
			msg.encode(&mut buf, Version::Lite05),
			Err(EncodeError::TooLarge)
		));
		assert!(buf.is_empty());
	}

	#[test]
	fn path_round_trip() {
		let msg = Setup {
			probe: ProbeLevel::Report,
			path: Some("/room/123".to_string()),
			..Default::default()
		};
		assert_eq!(round_trip(&msg), msg);
	}

	#[test]
	fn hop_round_trip() {
		let msg = Setup {
			hop: Some(crate::Hop::new(42).unwrap()),
			..Default::default()
		};
		assert_eq!(round_trip(&msg), msg);
	}

	// A declared id of 0 carries no identity (it cannot be excluded), so it
	// decodes as absent rather than erroring.
	#[test]
	fn hop_zero_decodes_as_none() {
		use crate::coding::Encode;

		let version = Version::Lite05;
		let mut params = Parameters::default();
		params.set_varint(super::PARAM_HOP, 0, version).unwrap();
		let mut body = bytes::BytesMut::new();
		params.encode(&mut body, version).unwrap();
		// Frame the body with the Message Length prefix `Setup::decode` expects.
		let mut buf = bytes::BytesMut::new();
		(body.len() as u64).encode(&mut buf, version).unwrap();
		buf.extend_from_slice(&body);
		let mut slice = &buf[..];
		let got = Setup::decode(&mut slice, version).unwrap();
		assert_eq!(got.hop, None);
	}

	#[test]
	fn empty_path_round_trips() {
		// An empty path is valid and distinct from absent only on the wire; both mean
		// the root, so a client doesn't have to special-case it.
		let msg = Setup {
			path: Some(String::new()),
			..Default::default()
		};
		assert_eq!(round_trip(&msg), msg);
	}

	#[test]
	fn roles_round_trip() {
		for role in [Some(Role::Publisher), Some(Role::Subscriber), None] {
			let msg = Setup {
				path: Some("/room/123".to_string()),
				role,
				..Default::default()
			};
			assert_eq!(round_trip(&msg), msg);
		}
	}

	#[test]
	fn unknown_probe_level_saturates_to_increase() {
		// Frame a SETUP message carrying an unknown probe level (99) by hand: the
		// parameters body, prefixed with its length (the lite Message size prefix).
		let mut params = Parameters::default();
		params.set_varint(PARAM_PROBE, 99, Version::Lite05).unwrap();
		let mut body = Vec::new();
		params.encode(&mut body, Version::Lite05).unwrap();

		let mut buf = bytes::BytesMut::new();
		body.len().encode(&mut buf, Version::Lite05).unwrap();
		buf.extend_from_slice(&body);

		let mut slice = &buf[..];
		let got = Setup::decode(&mut slice, Version::Lite05).unwrap();
		assert_eq!(got.probe, ProbeLevel::Increase);
	}

	#[test]
	fn role_wire_codes() {
		// The draft pins Publisher=1 / Subscriber=2. A swap here would still round-trip
		// against our own decoder, but break every other implementation.
		for (role, code) in [(Role::Publisher, 1u64), (Role::Subscriber, 2)] {
			assert_eq!(role.to_code(), code);
			assert_eq!(Role::from_code(code), Some(role));
		}
	}

	#[test]
	fn unknown_role_decodes_as_bidirectional() {
		// A role value the receiver doesn't recognize (a future extension, or an explicit
		// 0) decodes to `None` rather than failing, so a newer client can't break an older
		// server. The draft mandates this fallback.
		for code in [0u64, 9, 250] {
			let mut params = Parameters::default();
			params.set_varint(PARAM_ROLE, code, Version::Lite05).unwrap();
			let mut body = Vec::new();
			params.encode(&mut body, Version::Lite05).unwrap();

			let mut buf = bytes::BytesMut::new();
			body.len().encode(&mut buf, Version::Lite05).unwrap();
			buf.extend_from_slice(&body);

			let mut slice = &buf[..];
			let got = Setup::decode(&mut slice, Version::Lite05).unwrap();
			assert_eq!(got.role, None, "role code {code} should decode as bidirectional");
		}
	}

	#[test]
	fn rejects_before_lite05() {
		let msg = Setup::default();
		let mut buf = bytes::BytesMut::new();
		assert!(matches!(
			msg.encode(&mut buf, Version::Lite04),
			Err(EncodeError::Version)
		));
	}

	#[test]
	fn ignores_unknown_parameters() {
		// Frame a SETUP carrying an unknown parameter ID alongside the path.
		let mut params = Parameters::default();
		params.set_bytes(PARAM_PATH, b"/foo".to_vec());
		params.set_bytes(0xbeef, b"whatever".to_vec());

		let mut body = Vec::new();
		params.encode(&mut body, Version::Lite05).unwrap();

		// Wrap with the message size prefix the Message impl expects.
		let mut buf = bytes::BytesMut::new();
		body.len().encode(&mut buf, Version::Lite05).unwrap();
		buf.extend_from_slice(&body);

		let mut slice = &buf[..];
		let got = Setup::decode(&mut slice, Version::Lite05).unwrap();
		assert_eq!(got.path.as_deref(), Some("/foo"));
	}
}