moq_hls/server/

mod.rs

//! HTTP server: serves HLS and DASH for MoQ broadcasts, fetching media on demand.
//!
//! Routes are path-based, so one server can expose many broadcasts:
//!
//! ```text
//! GET /{broadcast}/master.m3u8
//! GET /{broadcast}/manifest.mpd
//! GET /{broadcast}/{kind}/{rendition}/media.m3u8
//! GET /{broadcast}/{kind}/{rendition}/init.{hash}.mp4
//! GET /{broadcast}/{kind}/{rendition}/seg/{segment}.m4s
//! GET /{broadcast}/{kind}/{rendition}/seg/t{pts}.m4s
//! ```
//!
//! `{kind}` is `video` or `audio`, so a video and an audio rendition that share a
//! name address distinct resources. The two `seg/` forms name the same bytes: HLS
//! playlists address a segment by its aligned number, the DASH manifest by its
//! timeline `pts` (`$Time$`). `{hash}` is a hash of the init segment's bytes, so a
//! reconfigured rendition gets a new init URL. A broadcaster given a generation
//! ([`Broadcaster::set_generation`](crate::export::Broadcaster::set_generation))
//! prefixes every segment file with it (`seg/{generation}.{segment}.m4s`), and
//! refuses segment URLs carrying any other.
//!
//! Every request is served. To gate access, wrap [`Server::router`] in your own
//! [`axum`] middleware. It runs before routing, so a rejected request never reaches
//! the origin, but it also sees the raw request URI rather than the path segments
//! the server decodes. A broadcast occupies every segment before the HLS resource
//! suffix, still percent-encoded: `/project/li%76e/master.m3u8` serves the broadcast
//! `project/live`. Decode each segment before matching it against a policy, or a
//! name can be encoded past the check.
//!
//! An embedder that authorizes each request itself (a token scoping the broadcast, a
//! stats-metered origin per tenant) skips the router instead: parse the path with
//! [`Route::parse`], check [`Route::broadcast`], rewrite it relative to the scope, and
//! answer with [`Server::respond`] on a `Server` built from that scope's origin. The
//! decoded broadcast is what a policy should check.
//!
//! ```no_run
//! use axum::http::StatusCode;
//! use axum::middleware::{self, Next};
//! use axum::extract::Request;
//! use axum::response::Response;
//!
//! async fn gate(req: Request, next: Next) -> Result<Response, StatusCode> {
//!     match req.headers().get("authorization") {
//!         Some(token) if token == "secret" => Ok(next.run(req).await),
//!         _ => Err(StatusCode::UNAUTHORIZED),
//!     }
//! }
//!
//! # fn build(server: moq_hls::Server) -> axum::Router {
//! server.router().layer(middleware::from_fn(gate))
//! # }
//! ```

mod routes;

pub use routes::{Resource, Route};

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use axum::Router;

use crate::export::{Broadcaster, Config};

/// How long to wait for a requested broadcast to be announced by the relay.
const RESOLVE_TIMEOUT: Duration = Duration::from_secs(5);

/// HLS export HTTP server. Cheap to clone (shared inner).
#[derive(Clone)]
pub struct Server {
	inner: Arc<Inner>,
}

struct Inner {
	origin: moq_net::origin::Consumer,
	config: Config,
	broadcasters: Mutex<HashMap<String, Arc<Broadcaster>>>,
}

impl Server {
	/// Build a server reading broadcasts from `origin`. Every request is served;
	/// gate access by layering middleware onto [`router`](Self::router).
	pub fn new(origin: moq_net::origin::Consumer, config: Config) -> Self {
		Self {
			inner: Arc::new(Inner {
				origin,
				config,
				broadcasters: Mutex::new(HashMap::new()),
			}),
		}
	}

	/// The axum router for the HLS endpoints, ready to nest or wrap in middleware.
	pub fn router(&self) -> Router {
		routes::router(self.clone())
	}

	/// Get or create the [`Broadcaster`] for `name`, resolving the broadcast from
	/// the relay (waiting briefly for its announcement). Returns `None` if the
	/// broadcast never shows up.
	pub async fn broadcaster(&self, name: &str) -> Option<Arc<Broadcaster>> {
		{
			let mut broadcasters = self.inner.broadcasters.lock().unwrap();
			if let Some(existing) = broadcasters.get(name) {
				if !existing.is_closed() {
					return Some(existing.clone());
				}
				broadcasters.remove(name);
			}
		}

		// Confirm a route covers the broadcast (and it is in scope) before building a
		// broadcaster; `Broadcaster::new` re-resolves it through the origin, which also
		// lets a rendition's catalog `broadcast` field reference a sibling broadcast.
		tokio::time::timeout(RESOLVE_TIMEOUT, self.inner.origin.routed(name))
			.await
			.ok()
			.flatten()?;

		let source = moq_mux::Source::new(self.inner.origin.consume(), name);
		let broadcaster = Broadcaster::new(source, self.inner.config.clone())
			.await
			.map_err(|err| tracing::warn!(%err, %name, "failed to resolve broadcast catalog"))
			.ok()?;

		let mut broadcasters = self.inner.broadcasters.lock().unwrap();
		if let Some(existing) = broadcasters.get(name) {
			if !existing.is_closed() {
				return Some(existing.clone());
			}
			broadcasters.remove(name);
		}

		let name = name.to_string();
		broadcasters.insert(name.clone(), broadcaster.clone());
		tokio::spawn(evict_closed(self.inner.clone(), name, broadcaster.clone()));
		Some(broadcaster)
	}
}

async fn evict_closed(inner: Arc<Inner>, name: String, broadcaster: Arc<Broadcaster>) {
	broadcaster.closed().await;

	let mut broadcasters = inner.broadcasters.lock().unwrap();
	if broadcasters
		.get(&name)
		.is_some_and(|current| Arc::ptr_eq(current, &broadcaster))
	{
		broadcasters.remove(&name);
	}
}

#[cfg(test)]
mod tests {
	/// Build an origin producer, spawning its driver on the ambient runtime.
	fn produce_origin() -> moq_net::origin::Producer {
		let (producer, driver) = moq_net::origin::Producer::new(moq_net::origin::Config::default());
		if tokio::runtime::Handle::try_current().is_ok() {
			tokio::spawn(moq_net::time::run(driver));
		} else {
			// A sync test: nothing polls the driver, and dropping it would tear
			// the origin down, so leak it and rely on the synchronous half.
			std::mem::forget(driver);
		}
		producer
	}

	use super::*;

	/// Let the origin's spawned attach/detach tasks run: a created broadcast
	/// becomes routable (and a finished one unroutable) asynchronously.
	async fn settle() {
		for _ in 0..10 {
			tokio::task::yield_now().await;
		}
	}

	/// Mirrors the middleware example in this module's docs, which `doctest = false`
	/// keeps out of the compiler's reach. Paused time skips the allowed request's
	/// RESOLVE_TIMEOUT wait for a broadcast that never arrives.
	#[tokio::test(start_paused = true)]
	async fn router_middleware_gates_requests() {
		use axum::body::Body;
		use axum::extract::Request;
		use axum::http::{StatusCode, header};
		use axum::middleware::{self, Next};
		use axum::response::Response;
		use tower::ServiceExt;

		async fn gate(req: Request, next: Next) -> Result<Response, StatusCode> {
			match req.headers().get(header::AUTHORIZATION) {
				Some(token) if token == "secret" => Ok(next.run(req).await),
				_ => Err(StatusCode::UNAUTHORIZED),
			}
		}

		// An origin with no broadcasts: a request that reaches the handlers 404s after
		// RESOLVE_TIMEOUT, so a 401 proves the middleware rejected it first.
		let origin = produce_origin();
		let server = Server::new(origin.consume(), Config::default());
		let app = server.router().layer(middleware::from_fn(gate));

		let denied = app
			.clone()
			.oneshot(Request::builder().uri("/live/master.m3u8").body(Body::empty()).unwrap())
			.await
			.unwrap();
		assert_eq!(denied.status(), StatusCode::UNAUTHORIZED);

		let allowed = app
			.oneshot(
				Request::builder()
					.uri("/live/master.m3u8")
					.header(header::AUTHORIZATION, "secret")
					.body(Body::empty())
					.unwrap(),
			)
			.await
			.unwrap();
		assert_eq!(allowed.status(), StatusCode::NOT_FOUND);
		// Pin the 404 to our handler: an unmatched route would also 404, but without
		// the no-store that not_found() sets, so a broken URI can't fake this pass.
		assert_eq!(allowed.headers()[header::CACHE_CONTROL], "no-store");
	}

	#[tokio::test(start_paused = true)]
	async fn router_matches_multisegment_broadcast_when_nested() {
		use axum::body::Body;
		use axum::extract::Request;
		use axum::http::{StatusCode, header};
		use tower::ServiceExt;

		let origin = produce_origin();
		let server = Server::new(origin.consume(), Config::default());
		let app = Router::new().nest("/hls", server.router());

		let response = app
			.oneshot(
				Request::builder()
					.uri("/hls/project/live/master.m3u8")
					.body(Body::empty())
					.unwrap(),
			)
			.await
			.unwrap();

		assert_eq!(response.status(), StatusCode::NOT_FOUND);
		// The handler adds no-store. Axum's unmatched-route 404 does not.
		assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
	}

	async fn closed_broadcaster() -> Arc<Broadcaster> {
		let origin = produce_origin();
		let producer = origin.create_broadcast("gone").expect("publish allowed");
		producer.announce(Default::default()).expect("publish allowed");
		settle().await;
		let source = moq_mux::Source::new(origin.consume(), "gone");
		let broadcaster = Broadcaster::new(source, Config::default())
			.await
			.expect("catalog broadcast resolves while announced");
		// Close the publisher so the resolved broadcast (and the broadcaster) reports closed.
		producer.close();
		settle().await;
		broadcaster
	}

	#[tokio::test]
	async fn broadcaster_replaces_finished_cached_instance() {
		let origin = produce_origin();
		let server = Server::new(origin.consume(), Config::default());
		let stale = closed_broadcaster().await;

		server
			.inner
			.broadcasters
			.lock()
			.unwrap()
			.insert("live".to_string(), stale.clone());
		let _producer = origin.create_broadcast("live").expect("publish allowed");
		_producer.announce(Default::default()).expect("publish allowed");
		settle().await;

		let fresh = server.broadcaster("live").await.expect("broadcast announced");

		assert!(!Arc::ptr_eq(&stale, &fresh));
		assert!(server.inner.broadcasters.lock().unwrap().contains_key("live"));
	}

	#[tokio::test]
	async fn eviction_keeps_newer_cached_instance() {
		let origin = produce_origin();
		let server = Server::new(origin.consume(), Config::default());
		let old = closed_broadcaster().await;
		let new_producer = origin.create_broadcast("live").expect("publish allowed");
		new_producer.announce(Default::default()).expect("publish allowed");
		settle().await;
		let new = Broadcaster::new(moq_mux::Source::new(origin.consume(), "live"), Config::default())
			.await
			.expect("catalog broadcast resolves while announced");

		server
			.inner
			.broadcasters
			.lock()
			.unwrap()
			.insert("live".to_string(), new.clone());

		evict_closed(server.inner.clone(), "live".to_string(), old).await;

		let cached = server.inner.broadcasters.lock().unwrap().get("live").cloned();
		assert!(cached.is_some_and(|cached| Arc::ptr_eq(&cached, &new)));
		new_producer.close();
	}
}