mosaik 0.4.2

use {
	super::{
		CallContext,
		FunctionId,
		FunctionNotFound,
		Functions,
		HandlerFailure,
		NoCapacity,
		NotAllowed,
		handler::{Registry, registry::DispatchError},
	},
	crate::{
		NetworkId,
		discovery::{Discovery, PeerInfo},
		network::{
			LocalNode,
			UnknownPeer,
			error::{DifferentNetwork, ProtocolViolation},
			link::{Link, Protocol},
		},
		primitives::{Bytes, Short, ShortFmtExt},
	},
	core::fmt,
	iroh::{
		endpoint::Connection,
		protocol::{AcceptError, ProtocolHandler},
	},
	n0_error::Meta,
	serde::{Deserialize, Serialize},
	std::sync::Arc,
};

/// Functions protocol acceptor
///
/// This type is responsible for accepting incoming call connections from
/// remote callers and dispatching them to the registered function handler
/// for the requested [`FunctionId`].
///
/// The call protocol works as follows:
///
/// - A remote caller connects using the functions protocol ALPN identifier and
///   sends a [`CallRequest`] containing the desired function id and the encoded
///   call payload.
///
/// - The acceptor looks up the registered handler for the requested function
///   id, authorizes the caller against the handler's peer requirements and
///   ticket validators, runs the function, and sends the [`CallReply`] back
///   over the same link before closing it.
///
/// - If the connecting caller peer is not known in the discovery catalog, the
///   call is rejected. When the caller gets an `UnknownPeer` close reason, it
///   should re-sync its catalog and retry the call.
pub(super) struct Acceptor {
	local: LocalNode,
	discovery: Discovery,
	registry: Arc<Registry>,
}

impl Acceptor {
	/// Creates a new [`Acceptor`] instance for the [`Functions`] subsystem.
	pub(super) fn new(functions: &Functions) -> Self {
		Self {
			local: functions.local.clone(),
			discovery: functions.discovery.clone(),
			registry: Arc::clone(&functions.registry),
		}
	}
}

impl fmt::Debug for Acceptor {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		// Safety: ALPN is valid UTF-8 hardcoded at compile time
		unsafe {
			write!(
				f,
				"Functions({})",
				str::from_utf8_unchecked(Functions::ALPN)
			)
		}
	}
}

impl ProtocolHandler for Acceptor {
	#[allow(clippy::too_many_lines)]
	async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
		let cancel = self.local.termination().clone();
		let mut link =
			Link::<Functions>::accept_with_cancel(connection, cancel).await?;
		let remote_peer_id = link.remote_id();
		let catalog = self.discovery.catalog();
		let network_label =
			[("network", self.local.network_id().short().to_string())];

		let Some(peer) = catalog.get(&remote_peer_id) else {
			tracing::trace!(
				peer_id = %Short(&remote_peer_id),
				"rejecting unidentified caller",
			);

			metrics::counter!("mosaik.functions.calls.rejected", &network_label)
				.increment(1);

			// Close the link with a reason before returning error
			link
				.close(UnknownPeer)
				.await
				.map_err(AcceptError::from_err)?;

			return Err(AcceptError::NotAllowed {
				meta: Meta::default(),
			});
		};

		// Receive the caller's call request
		let request: CallRequest = link
			.recv()
			.await
			.inspect_err(|e| {
				tracing::debug!(
					caller_id = %Short(peer.id()),
					error = %e,
					"Failed to receive call request",
				);
			})
			.map_err(AcceptError::from_err)?;

		// ensure that the caller is connecting to the correct network
		if request.network_id != self.local.network_id() {
			tracing::debug!(
				caller_id = %Short(peer.id()),
				function_id = %Short(request.function_id),
				expected_network = %Short(self.local.network_id()),
				received_network = %Short(request.network_id),
				"Caller connected to wrong network",
			);

			metrics::counter!("mosaik.functions.calls.rejected", &network_label)
				.increment(1);

			link
				.close(DifferentNetwork)
				.await
				.map_err(AcceptError::from_err)?;

			return Err(AcceptError::NotAllowed {
				meta: Meta::default(),
			});
		}

		// Lookup the registered handler for the requested function id
		let Some(entry) = self.registry.open(request.function_id) else {
			tracing::debug!(
				caller_id = %Short(peer.id()),
				function_id = %request.function_id,
				"Caller requesting unavailable function",
			);

			metrics::counter!("mosaik.functions.calls.rejected", &network_label)
				.increment(1);

			link
				.close(FunctionNotFound)
				.await
				.map_err(AcceptError::from_err)?;

			return Err(AcceptError::NotAllowed {
				meta: Meta::default(),
			});
		};

		// Authorize the caller against the handler's peer requirements and
		// ticket validators.
		let info = PeerInfo::from_tracker(peer, self.discovery.rtt_tracker());
		if !(entry.require)(&info)
			|| peer.validate_tickets(&entry.caller_auth).is_err()
		{
			tracing::debug!(
				caller_id = %Short(peer.id()),
				function_id = %Short(request.function_id),
				"Caller not allowed to invoke function",
			);

			metrics::counter!("mosaik.functions.calls.rejected", &network_label)
				.increment(1);

			link
				.close(NotAllowed)
				.await
				.map_err(AcceptError::from_err)?;

			return Err(AcceptError::NotAllowed {
				meta: Meta::default(),
			});
		}

		// Bound handler concurrency; the permit is held for the duration of
		// the dispatch and released when this scope ends.
		let Ok(_permit) = Arc::clone(&entry.permits).try_acquire_owned() else {
			tracing::debug!(
				caller_id = %Short(peer.id()),
				function_id = %Short(request.function_id),
				"Function has no capacity for new calls",
			);

			metrics::counter!("mosaik.functions.calls.rejected", &network_label)
				.increment(1);

			link
				.close(NoCapacity)
				.await
				.map_err(AcceptError::from_err)?;

			return Err(AcceptError::NotAllowed {
				meta: Meta::default(),
			});
		};

		metrics::counter!("mosaik.functions.calls.accepted", &network_label)
			.increment(1);
		metrics::gauge!("mosaik.functions.calls.inflight", &network_label)
			.increment(1.0);
		entry.stats.record_started();

		// Run the function, racing it against handler deregistration and the
		// caller abandoning the call, so long-running dispatches don't outlive
		// either side.
		let context = CallContext::new(peer.clone());
		let dispatch = (entry.dispatch)(request.payload, context);
		let caller_gone = link.closed();
		let result = tokio::select! {
			result = dispatch => Some(result),
			() = entry.cancel.cancelled() => None,
			_ = caller_gone => {
				// the caller closed the connection before the reply was ready
				metrics::gauge!("mosaik.functions.calls.inflight", &network_label)
					.decrement(1.0);
				entry.stats.record_failed();
				return Ok(());
			}
		};

		metrics::gauge!("mosaik.functions.calls.inflight", &network_label)
			.decrement(1.0);

		let Some(result) = result else {
			// the handler was dropped while the call was in flight; for the
			// caller this is equivalent to the function not being registered.
			entry.stats.record_failed();
			link
				.close(FunctionNotFound)
				.await
				.map_err(AcceptError::from_err)?;

			return Err(AcceptError::NotAllowed {
				meta: Meta::default(),
			});
		};

		match result {
			Ok(reply) => {
				entry.stats.record_served();
				metrics::counter!("mosaik.functions.calls.completed", &network_label)
					.increment(1);

				link.send(&reply).await.map_err(AcceptError::from_err)?;

				// The caller closes the link once it has received the reply.
				// Closing from this side instead would race the reply frame with
				// the connection close and could drop it before delivery.
				if let Err(e) = link.closed().await {
					tracing::trace!(
						caller_id = %Short(peer.id()),
						function_id = %Short(request.function_id),
						error = %e,
						"call link closed with error",
					);
				}

				Ok(())
			}
			Err(DispatchError::DecodeRequest) => {
				entry.stats.record_failed();
				metrics::counter!("mosaik.functions.calls.failed", &network_label)
					.increment(1);

				link
					.close(ProtocolViolation)
					.await
					.map_err(AcceptError::from_err)?;

				Err(AcceptError::NotAllowed {
					meta: Meta::default(),
				})
			}
			Err(DispatchError::EncodeReply) => {
				entry.stats.record_failed();
				metrics::counter!("mosaik.functions.calls.failed", &network_label)
					.increment(1);

				link
					.close(HandlerFailure)
					.await
					.map_err(AcceptError::from_err)?;

				Err(AcceptError::NotAllowed {
					meta: Meta::default(),
				})
			}
		}
	}
}

/// A function invocation request sent by a caller when connecting to a
/// handler.
#[derive(Debug, Serialize, Deserialize)]
pub(super) struct CallRequest {
	/// The network id the caller is connecting to.
	pub network_id: NetworkId,

	/// The effective wire id of the function to invoke.
	pub function_id: FunctionId,

	/// The encoded request payload (the function's `Req` type encoded via
	/// [`Datum::encode`](crate::primitives::Datum::encode)).
	pub payload: Bytes,
}

/// The reply sent by a handler after running the function.
///
/// Both variants are successful protocol exchanges — `Err` carries the
/// function's typed application-level error, which never triggers caller
/// failover. Handler-side faults are communicated through close reasons
/// instead.
#[derive(Debug, Serialize, Deserialize)]
pub(super) enum CallReply {
	/// The function completed; carries the encoded `Res` value.
	Ok(Bytes),

	/// The function returned an application error; carries the encoded `E`
	/// value.
	Err(Bytes),
}