reinhardt_core/

ws.rs

//! WebSocket routing primitives shared across reinhardt crates.
//!
//! This module provides foundational WebSocket types used by both
//! `reinhardt-websockets` (connection handling) and `reinhardt-urls`
//! (`UnifiedRouter::websocket()` builder). Placing them here avoids a
//! circular dependency between those two crates.

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

// ── Endpoint metadata ─────────────────────────────────────────────────────

/// Compile-time WebSocket endpoint metadata.
///
/// Implemented on the consumer struct generated by `#[websocket]`.
/// Parallel to `EndpointInfo` for HTTP views.
pub trait WebSocketEndpointInfo {
	/// Returns the URL path pattern for this WebSocket endpoint.
	fn path() -> &'static str;
	/// Returns the optional route name for this WebSocket endpoint.
	fn name() -> Option<&'static str>;
}

/// Inventory metadata submitted by `#[websocket]` at compile time.
///
/// Used by `UrlReverser` to resolve WebSocket route names.
pub struct WebSocketEndpointMetadata {
	/// URL path pattern (e.g. `"/ws/chat/{room_id}/"`).
	pub path: &'static str,
	/// Route name used for URL reversal.
	pub name: &'static str,
	/// Handler function name for diagnostics.
	pub fn_name: &'static str,
	/// Rust module path of the handler for diagnostics.
	pub module_path: &'static str,
}

inventory::collect!(WebSocketEndpointMetadata);

/// Substitute path parameters in a WebSocket URL pattern.
///
/// `"/ws/chat/{room_id}/"` + `[("room_id", "42")]` → `"/ws/chat/42/"`
pub fn substitute_ws_params(path: &str, params: &[(&str, &str)]) -> String {
	let mut result = path.to_string();
	for (name, value) in params {
		result = result.replace(&format!("{{{}}}", name), value);
	}
	result
}

// ── Routing types ─────────────────────────────────────────────────────────

/// Routing result type
pub type RouteResult = Result<(), RouteError>;

/// Routing errors for WebSocket routes
#[derive(Debug, thiserror::Error)]
pub enum RouteError {
	/// No route registered for the given path.
	#[error("Route not found: {0}")]
	NotFound(String),
	/// A route with the given path is already registered.
	#[error("Route already exists: {0}")]
	AlreadyExists(String),
	/// The provided route pattern is syntactically invalid.
	#[error("Invalid route pattern: {0}")]
	InvalidPattern(String),
}

/// A registered WebSocket route (path + optional name + metadata).
#[derive(Debug, Clone)]
pub struct WebSocketRoute {
	path: String,
	name: Option<String>,
	metadata: HashMap<String, String>,
}

impl WebSocketRoute {
	/// Creates a new route with the given path and optional name.
	pub fn new(path: String, name: Option<String>) -> Self {
		Self {
			path,
			name,
			metadata: HashMap::new(),
		}
	}

	/// Returns the URL path pattern for this route.
	pub fn path(&self) -> &str {
		&self.path
	}

	/// Returns the optional name for this route.
	pub fn name(&self) -> Option<&str> {
		self.name.as_deref()
	}

	/// Attaches a key-value metadata entry to this route.
	pub fn with_metadata(mut self, key: String, value: String) -> Self {
		self.metadata.insert(key, value);
		self
	}

	/// Returns the metadata value for the given key, if present.
	pub fn get_metadata(&self, key: &str) -> Option<&String> {
		self.metadata.get(key)
	}
}

// ── WebSocketRouter ───────────────────────────────────────────────────────

/// WebSocket router: build-time registration + runtime lookup.
///
/// The build-time API (`consumer()`, `reverse()`, `find_pending()`) is used
/// by `UnifiedRouter::websocket()` and WebSocket route declarations.
/// The async API (`register_route()`, `find_route()`, etc.) is used at
/// connection-handling time in `reinhardt-websockets`.
#[derive(Clone)]
pub struct WebSocketRouter {
	routes: Arc<RwLock<HashMap<String, WebSocketRoute>>>,
	names: Arc<RwLock<HashMap<String, String>>>,
	/// Build-time consumer registrations (added by `consumer()` builder).
	pending_consumers: Vec<WebSocketRoute>,
	/// Optional namespace (app label) for this router.
	namespace: Option<String>,
}

impl WebSocketRouter {
	/// Creates a new empty router.
	pub fn new() -> Self {
		Self {
			routes: Arc::new(RwLock::new(HashMap::new())),
			names: Arc::new(RwLock::new(HashMap::new())),
			pending_consumers: Vec::new(),
			namespace: None,
		}
	}

	/// Set the namespace for this router.
	///
	/// Parallel to `ServerRouter::with_namespace`.
	/// WebSocket route paths are absolute today and are not rewritten
	/// with this namespace; the value is stored for parity with other
	/// routers and future use. See reinhardt-web#3829.
	pub fn with_namespace(mut self, namespace: impl Into<String>) -> Self {
		self.namespace = Some(namespace.into());
		self
	}

	/// Returns the namespace set via [`with_namespace`], if any.
	///
	/// [`with_namespace`]: Self::with_namespace
	pub fn namespace(&self) -> Option<&str> {
		self.namespace.as_deref()
	}

	/// Register a WebSocket consumer by its factory function.
	///
	/// Parallel to `ServerRouter::endpoint()`. Path and name are derived
	/// from `C`'s `WebSocketEndpointInfo` impl at compile time.
	pub fn consumer<C, F>(mut self, _f: F) -> Self
	where
		F: Fn() -> C,
		C: WebSocketEndpointInfo + 'static,
	{
		self.pending_consumers.push(WebSocketRoute::new(
			C::path().to_string(),
			C::name().map(|s| s.to_string()),
		));
		self
	}

	/// Find a pending consumer route by name.
	pub fn find_pending(&self, name: &str) -> Option<&WebSocketRoute> {
		self.pending_consumers
			.iter()
			.find(|r| r.name() == Some(name))
	}

	/// Resolve a WebSocket URL by route name, substituting path parameters.
	pub fn reverse(&self, name: &str, params: &[(&str, &str)]) -> Option<String> {
		self.pending_consumers
			.iter()
			.find(|r| r.name() == Some(name))
			.map(|r| substitute_ws_params(r.path(), params))
	}

	/// Register a route at runtime (async, used by the connection handler).
	pub async fn register_route(&mut self, route: WebSocketRoute) -> RouteResult {
		let mut routes = self.routes.write().await;
		if routes.contains_key(&route.path) {
			return Err(RouteError::AlreadyExists(route.path.clone()));
		}
		if let Some(name) = &route.name {
			let mut names = self.names.write().await;
			names.insert(name.clone(), route.path.clone());
		}
		routes.insert(route.path.clone(), route);
		Ok(())
	}

	/// Looks up a registered route by its exact path.
	pub async fn find_route(&self, path: &str) -> Option<WebSocketRoute> {
		let routes = self.routes.read().await;
		routes.get(path).cloned()
	}

	/// Looks up a registered route by its name.
	pub async fn find_route_by_name(&self, name: &str) -> Option<WebSocketRoute> {
		let names = self.names.read().await;
		if let Some(path) = names.get(name) {
			let routes = self.routes.read().await;
			routes.get(path).cloned()
		} else {
			None
		}
	}

	/// Removes the registered route for the given path.
	pub async fn remove_route(&mut self, path: &str) -> RouteResult {
		let mut routes = self.routes.write().await;
		let route = routes
			.remove(path)
			.ok_or_else(|| RouteError::NotFound(path.to_string()))?;
		if let Some(name) = &route.name {
			let mut names = self.names.write().await;
			names.remove(name);
		}
		Ok(())
	}

	/// Returns all currently registered routes.
	pub async fn all_routes(&self) -> Vec<WebSocketRoute> {
		let routes = self.routes.read().await;
		routes.values().cloned().collect()
	}

	/// Returns `true` if a route is registered for the given path.
	pub async fn has_route(&self, path: &str) -> bool {
		self.routes.read().await.contains_key(path)
	}

	/// Returns the number of currently registered routes.
	pub async fn route_count(&self) -> usize {
		self.routes.read().await.len()
	}

	/// Removes all registered routes and name mappings.
	pub async fn clear(&mut self) {
		self.routes.write().await.clear();
		self.names.write().await.clear();
	}
}

impl Default for WebSocketRouter {
	fn default() -> Self {
		Self::new()
	}
}

// ── Global registry ───────────────────────────────────────────────────────

static GLOBAL_ROUTER: once_cell::sync::Lazy<Arc<RwLock<Option<WebSocketRouter>>>> =
	once_cell::sync::Lazy::new(|| Arc::new(RwLock::new(None)));

/// Installs `router` as the process-wide WebSocket router.
pub async fn register_websocket_router(router: WebSocketRouter) {
	*GLOBAL_ROUTER.write().await = Some(router);
}

/// Returns a clone of the current process-wide WebSocket router, if set.
pub async fn get_websocket_router() -> Option<WebSocketRouter> {
	GLOBAL_ROUTER.read().await.clone()
}

/// Clears the process-wide WebSocket router (primarily for tests).
pub async fn clear_websocket_router() {
	*GLOBAL_ROUTER.write().await = None;
}

/// Resolves a registered or pending WebSocket URL by route name.
pub async fn reverse_websocket_url(router: &WebSocketRouter, name: &str) -> Option<String> {
	let names = router.names.read().await;
	if let Some(path) = names.get(name) {
		let routes = router.routes.read().await;
		routes.get(path).map(|r| r.path().to_string())
	} else {
		router.find_pending(name).map(|r| r.path().to_string())
	}
}

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

	struct TestConsumer;
	impl WebSocketEndpointInfo for TestConsumer {
		fn path() -> &'static str {
			"/ws/chat/{room_id}/"
		}
		fn name() -> Option<&'static str> {
			Some("chat_ws")
		}
	}

	#[rstest]
	fn test_substitute_no_params() {
		assert_eq!(substitute_ws_params("/ws/notif/", &[]), "/ws/notif/");
	}

	#[rstest]
	fn test_substitute_one_param() {
		assert_eq!(
			substitute_ws_params("/ws/chat/{room_id}/", &[("room_id", "42")]),
			"/ws/chat/42/"
		);
	}

	#[rstest]
	fn test_consumer_builder() {
		let router = WebSocketRouter::new().consumer(|| TestConsumer);
		let route = router.find_pending("chat_ws");
		assert!(route.is_some());
		assert_eq!(route.unwrap().path(), "/ws/chat/{room_id}/");
	}

	#[rstest]
	fn test_with_namespace_stores_value_without_rewriting_paths() {
		let router = WebSocketRouter::new()
			.with_namespace("auth")
			.consumer(|| TestConsumer);
		assert_eq!(router.namespace(), Some("auth"));
		assert_eq!(
			router.find_pending("chat_ws").unwrap().path(),
			"/ws/chat/{room_id}/"
		);
	}

	#[rstest]
	fn test_reverse() {
		let router = WebSocketRouter::new().consumer(|| TestConsumer);
		assert_eq!(
			router.reverse("chat_ws", &[("room_id", "99")]),
			Some("/ws/chat/99/".to_string())
		);
		assert_eq!(router.reverse("unknown", &[]), None);
	}
}