use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use bytes::Bytes;
use futures_util::{stream, StreamExt};
use n0_future::time::Instant;
use serde_json::Value;
use unb_core::{Envelope, ErrorCode, Kind, Resolution, TargetPath, DEFAULT_HOPS};
use unb_runtime::WireBody;
use crate::handler::HandlerError;
use crate::layer::Origin;
use crate::node::{Node, NodeSnapshot};
pub(crate) const CALL_TIMEOUT: Duration = Duration::from_secs(30);
pub trait IntoBody: Send {
fn into_body(self) -> Bytes;
}
impl IntoBody for Bytes {
fn into_body(self) -> Bytes {
self
}
}
impl IntoBody for Vec<u8> {
fn into_body(self) -> Bytes {
self.into()
}
}
impl IntoBody for String {
fn into_body(self) -> Bytes {
self.into()
}
}
impl IntoBody for &str {
fn into_body(self) -> Bytes {
Bytes::copy_from_slice(self.as_bytes())
}
}
impl IntoBody for Value {
fn into_body(self) -> Bytes {
Envelope::encode_payload(&self)
}
}
impl IntoBody for () {
fn into_body(self) -> Bytes {
Bytes::new()
}
}
pub trait Destination: Send + Sync {
fn send(
&self,
request: http::Request<Bytes>,
) -> impl Future<Output = Result<http::Response<Bytes>, HandlerError>> + Send;
}
impl Destination for Arc<Node> {
async fn send(
&self,
request: http::Request<Bytes>,
) -> Result<http::Response<Bytes>, HandlerError> {
let response = self.fetch(request).await?;
let (parts, body) = response.into_parts();
match body {
crate::layer::ServiceBody::Unary(payload) => {
Ok(http::Response::from_parts(parts, payload))
}
crate::layer::ServiceBody::Stream(_) => Err(HandlerError::new(
ErrorCode::InvalidInput,
"subscribe is not available over send; use Node::subscribe",
)),
}
}
}
impl<D: Destination + ?Sized> Destination for &D {
async fn send(
&self,
request: http::Request<Bytes>,
) -> Result<http::Response<Bytes>, HandlerError> {
D::send(self, request).await
}
}
#[cfg(feature = "hosting")]
impl Destination for &str {
async fn send(
&self,
request: http::Request<Bytes>,
) -> Result<http::Response<Bytes>, HandlerError> {
n0_future::time::timeout(CALL_TIMEOUT, one_shot_http(self, request))
.await
.map_err(|_| {
HandlerError::new(
ErrorCode::PeerUnreachable,
format!("{self:?} did not answer within the call timeout"),
)
})?
}
}
#[cfg(feature = "hosting")]
impl Destination for String {
async fn send(
&self,
request: http::Request<Bytes>,
) -> Result<http::Response<Bytes>, HandlerError> {
self.as_str().send(request).await
}
}
pub trait SendExt<T> {
fn send<D: Destination>(
self,
destination: D,
) -> impl Future<Output = Result<http::Response<Bytes>, HandlerError>> + Send;
}
impl<T: IntoBody> SendExt<T> for http::Request<T> {
async fn send<D: Destination>(
self,
destination: D,
) -> Result<http::Response<Bytes>, HandlerError> {
let (mut parts, body) = self.into_parts();
if parts.method == http::Method::GET {
parts.method = http::Method::POST;
}
destination
.send(http::Request::from_parts(parts, body.into_body()))
.await
}
}
impl<T: IntoBody> SendExt<T> for Result<http::Request<T>, http::Error> {
async fn send<D: Destination>(
self,
destination: D,
) -> Result<http::Response<Bytes>, HandlerError> {
match self {
Ok(request) => request.send(destination).await,
Err(error) => Err(HandlerError::new(
ErrorCode::InvalidInput,
error.to_string(),
)),
}
}
}
#[cfg(feature = "hosting")]
async fn one_shot_http(
address: &str,
request: http::Request<Bytes>,
) -> Result<http::Response<Bytes>, HandlerError> {
let (tls, remainder) = if let Some(rest) = address.strip_prefix("https://") {
(true, rest)
} else if let Some(rest) = address.strip_prefix("wss://") {
(true, rest)
} else if let Some(rest) = address.strip_prefix("http://") {
(false, rest)
} else if let Some(rest) = address.strip_prefix("ws://") {
(false, rest)
} else {
(false, address)
};
let authority = remainder
.split('/')
.next()
.filter(|authority| !authority.is_empty())
.ok_or_else(|| {
HandlerError::new(
ErrorCode::InvalidInput,
format!("{address:?} names no host to send to"),
)
})?;
let has_port = match authority.rfind(']') {
Some(bracket) => authority[bracket + 1..].contains(':'),
None => authority.contains(':'),
};
let authority = if has_port {
authority.to_string()
} else {
format!("{authority}:{}", if tls { 443 } else { 80 })
};
let unreachable = |error: String| HandlerError::new(ErrorCode::PeerUnreachable, error);
let stream = tokio::net::TcpStream::connect(&authority)
.await
.map_err(|error| unreachable(error.to_string()))?;
if tls {
let host = authority
.rsplit_once(':')
.map(|(host, _)| host)
.unwrap_or(&authority)
.trim_start_matches('[')
.trim_end_matches(']');
let server_name = rustls_pki_types::ServerName::try_from(host.to_string())
.map_err(|error| HandlerError::new(ErrorCode::InvalidInput, error.to_string()))?;
let config = unb_transport::ws::tls_client_config()
.map_err(|error| unreachable(error.to_string()))?;
let stream = tokio_rustls::TlsConnector::from(config)
.connect(server_name, stream)
.await
.map_err(|error| unreachable(error.to_string()))?;
exchange_http1(stream, &authority, request).await
} else {
exchange_http1(stream, &authority, request).await
}
}
#[cfg(feature = "hosting")]
async fn exchange_http1<T>(
stream: T,
authority: &str,
request: http::Request<Bytes>,
) -> Result<http::Response<Bytes>, HandlerError>
where
T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
{
let target = request
.uri()
.path_and_query()
.map(|target| target.as_str())
.filter(|target| !target.is_empty())
.unwrap_or("/");
let unreachable = |error: String| HandlerError::new(ErrorCode::PeerUnreachable, error);
let (mut sender, connection) =
hyper::client::conn::http1::handshake(hyper_util::rt::TokioIo::new(stream))
.await
.map_err(|error| unreachable(error.to_string()))?;
tokio::spawn(async move {
let _ = connection.await;
});
let mut outbound = http::Request::builder()
.method(http::Method::POST)
.uri(target)
.header(http::header::HOST, authority);
for (name, value) in request.headers() {
if matches!(
*name,
http::header::HOST
| http::header::CONTENT_LENGTH
| http::header::TRANSFER_ENCODING
| http::header::CONNECTION
) {
continue;
}
outbound = outbound.header(name, value);
}
let outbound = outbound
.body(http_body_util::Full::new(request.into_body()))
.map_err(|error| HandlerError::new(ErrorCode::InvalidInput, error.to_string()))?;
let response = sender
.send_request(outbound)
.await
.map_err(|error| unreachable(error.to_string()))?;
let (parts, body) = response.into_parts();
let body = http_body_util::Limited::new(body, unb_transport::DEFAULT_MAX_FRAME_SIZE);
let body = http_body_util::BodyExt::collect(body)
.await
.map_err(|error| HandlerError::new(ErrorCode::Protocol, error.to_string()))?
.to_bytes();
let mut projected = http::Response::builder().status(parts.status);
for (name, value) in &parts.headers {
if matches!(
*name,
http::header::CONNECTION
| http::header::CONTENT_LENGTH
| http::header::TRANSFER_ENCODING
| http::header::DATE
) {
continue;
}
projected = projected.header(name, value);
}
projected
.body(body)
.map_err(|error| HandlerError::new(ErrorCode::Protocol, error.to_string()))
}
impl Node {
pub async fn fetch_body(
self: &Arc<Self>,
request: http::Request<WireBody>,
) -> Result<http::Response<WireBody>, HandlerError> {
self.fetch_body_until(request, Instant::now() + CALL_TIMEOUT)
.await
}
pub async fn fetch_body_with_timeout(
self: &Arc<Self>,
request: http::Request<WireBody>,
timeout: Duration,
) -> Result<http::Response<WireBody>, HandlerError> {
self.fetch_body_until(request, Instant::now() + timeout)
.await
}
pub(crate) async fn fetch_body_until(
self: &Arc<Self>,
request: http::Request<WireBody>,
deadline: Instant,
) -> Result<http::Response<WireBody>, HandlerError> {
if let Some(name) = request
.headers()
.keys()
.find(|name| name.as_str().starts_with("unb-"))
{
return Err(HandlerError::new(
ErrorCode::InvalidInput,
format!("{name}: unb-* headers are reserved for framing metadata"),
));
}
let (parts, body) = request.into_parts();
if parts.uri.query().is_some() {
return Err(HandlerError::new(
ErrorCode::InvalidInput,
"application request targets must not carry a query string",
));
}
let target_path = TargetPath::parse_application(parts.uri.path())
.map_err(|error| HandlerError::new(ErrorCode::InvalidInput, error.to_string()))?;
let target = target_path.target().to_owned();
let subject = target_path.subject().to_owned();
let (snapshot, resolution) = self.resolve_unary_until(&target, deadline).await?;
match resolution {
Resolution::Local => {
let (payload, streaming_body) = match body {
WireBody::Bytes(payload) => (payload, None),
WireBody::Stream(body) => (Bytes::new(), Some(body)),
};
let envelope = Envelope::from_request(http::Request::from_parts(parts, payload))
.map_err(|error| {
HandlerError::new(ErrorCode::InvalidInput, error.to_string())
})?;
let mut request = Self::inbound_request(&envelope)?;
if let Some(body) = streaming_body {
request
.extensions_mut()
.insert(crate::service::StreamingBody(Arc::new(
std::sync::Mutex::new(Some(body)),
)));
}
match self
.run_service(snapshot.clone(), request, Origin::Local)
.await
{
Some(outcome) => {
let response = outcome?;
let (parts, body) = response.into_parts();
let body = match body {
crate::layer::ServiceBody::Unary(payload) => WireBody::Bytes(payload),
crate::layer::ServiceBody::Stream(body) => {
WireBody::Stream(Box::pin(body.map(|item| {
item.map_err(|error| {
unb_core::CoreError::Malformed(error.to_string())
})
})))
}
};
Ok(http::Response::from_parts(parts, body))
}
None => Err(Self::teach_unknown_subject(&snapshot, &subject)),
}
}
Resolution::Route(peer_name) => {
let link = self.route_link(&peer_name).await?;
let remaining = Self::remaining_unary_time(deadline)?;
let mut response = link
.wire
.client_session()
.fetch_body(http::Request::from_parts(parts, body), remaining)
.await
.map_err(Self::client_error)?;
let reserved = response
.headers()
.keys()
.filter(|name| {
name.as_str().starts_with("unb-") && name.as_str() != unb_core::UNB_CODE
})
.cloned()
.collect::<Vec<_>>();
for name in reserved {
response.headers_mut().remove(name);
}
Ok(response)
}
Resolution::Conflicted { owners } => Err(HandlerError::new(
ErrorCode::PeerUnreachable,
format!(
"destination node {target:?} has multiple live incarnations: {}",
owners.join(", ")
),
)),
Resolution::Unknown => Err(Self::teach_unknown_target(&snapshot, &target)),
}
}
pub async fn fetch(
self: &Arc<Self>,
request: http::Request<bytes::Bytes>,
) -> Result<http::Response<crate::layer::ServiceBody>, HandlerError> {
self.fetch_until(request, Instant::now() + CALL_TIMEOUT)
.await
}
pub(crate) async fn fetch_until(
self: &Arc<Self>,
request: http::Request<bytes::Bytes>,
deadline: Instant,
) -> Result<http::Response<crate::layer::ServiceBody>, HandlerError> {
let (parts, body) = request.into_parts();
let response = self
.fetch_body_until(
http::Request::from_parts(parts, WireBody::Bytes(body)),
deadline,
)
.await?;
let (parts, body) = response.into_parts();
let body = match body {
WireBody::Bytes(payload) => crate::layer::ServiceBody::Unary(payload),
WireBody::Stream(body) => {
crate::layer::ServiceBody::Stream(Box::pin(body.map(|item| {
item.map_err(|error| HandlerError::new(ErrorCode::Protocol, error.to_string()))
})))
}
};
Ok(http::Response::from_parts(parts, body))
}
pub async fn subscribe(
self: &Arc<Self>,
target_path: &str,
payload: Value,
) -> Result<crate::EventStream, HandlerError> {
self.subscribe_with(target_path, payload, serde_json::Map::new())
.await
}
pub async fn subscribe_with(
self: &Arc<Self>,
target_path: &str,
payload: Value,
headers: serde_json::Map<String, Value>,
) -> Result<crate::EventStream, HandlerError> {
self.subscribe_bytes(target_path, Envelope::encode_payload(&payload), headers)
.await
}
pub async fn subscribe_bytes(
self: &Arc<Self>,
target_path: &str,
payload: Bytes,
headers: serde_json::Map<String, Value>,
) -> Result<crate::EventStream, HandlerError> {
let target_path = TargetPath::parse_application(target_path)
.map_err(|error| HandlerError::new(ErrorCode::InvalidInput, error.to_string()))?;
let target = target_path.target().to_owned();
let subject = target_path.subject().to_owned();
let target_path = target_path.to_string();
let (snapshot, resolution) = self
.resolve_unary_until(&target, Instant::now() + CALL_TIMEOUT)
.await?;
match resolution {
Resolution::Local => {
let request =
self.local_request(Kind::Subscribe, &subject, payload, headers.clone())?;
match self
.run_service(snapshot.clone(), request, Origin::Local)
.await
{
Some(Ok(response)) => match response.into_body() {
crate::layer::ServiceBody::Stream(stream) => Ok(stream),
crate::layer::ServiceBody::Unary(_) => Err(HandlerError::new(
ErrorCode::Internal,
"a streaming operation produced a unary response",
)),
},
Some(Err(error)) => Err(error),
None => Err(Self::teach_unknown_subject(&snapshot, &subject)),
}
}
Resolution::Route(peer_name) => {
let link = self.peer(&peer_name).await.ok_or_else(|| {
HandlerError::new(
ErrorCode::PeerUnreachable,
format!("no live connection to peer {peer_name:?}"),
)
})?;
let stream = link
.wire
.client_session()
.start(
&target_path,
Kind::Subscribe,
payload,
Some(DEFAULT_HOPS),
headers,
)
.await
.map_err(|error| {
HandlerError::new(ErrorCode::PeerUnreachable, error.to_string())
})?;
Ok(Box::pin(stream::unfold(stream, |mut stream| async move {
let item = match stream.next().await {
Ok(Some(envelope)) if envelope.kind == Kind::Event => {
Some(Ok(envelope.payload))
}
Ok(Some(envelope)) if envelope.kind == Kind::Response => None,
Ok(Some(_)) => Some(Err(HandlerError::new(
ErrorCode::Protocol,
"unexpected frame in subscription",
))),
Ok(None) => None,
Err(error) => Some(Err(Node::client_error(error))),
};
item.map(|item| (item, stream))
})))
}
Resolution::Conflicted { owners } => Err(HandlerError::new(
ErrorCode::PeerUnreachable,
format!(
"destination node {target:?} has multiple live incarnations: {}",
owners.join(", ")
),
)),
Resolution::Unknown => Err(Self::teach_unknown_target(&snapshot, &target)),
}
}
pub(crate) async fn call_nested(
self: &Arc<Self>,
target_path: &str,
payload: Value,
headers: serde_json::Map<String, Value>,
) -> Result<Value, HandlerError> {
self.call_with_origin(target_path, payload, headers, Origin::Nested)
.await
}
async fn call_with_origin(
self: &Arc<Self>,
target_path: &str,
payload: Value,
headers: serde_json::Map<String, Value>,
origin: Origin,
) -> Result<Value, HandlerError> {
let target_path = TargetPath::parse_application(target_path)
.map_err(|error| HandlerError::new(ErrorCode::InvalidInput, error.to_string()))?;
let target = target_path.target().to_owned();
let subject = target_path.subject().to_owned();
let target_path = target_path.to_string();
let deadline = Instant::now() + CALL_TIMEOUT;
let (snapshot, resolution) = self.resolve_unary_until(&target, deadline).await?;
match resolution {
Resolution::Local => {
let request = self.local_request(
Kind::Request,
&subject,
Envelope::encode_payload(&payload),
headers,
)?;
let outcome = self.run_service(snapshot.clone(), request, origin).await;
match outcome {
Some(outcome) => match outcome?.into_body() {
crate::layer::ServiceBody::Unary(payload) => {
Self::json_profile_payload(&payload)
}
crate::layer::ServiceBody::Stream(_) => Err(HandlerError::new(
ErrorCode::Internal,
"a unary operation produced a stream",
)),
},
None => Err(Self::teach_unknown_subject(&snapshot, &subject)),
}
}
Resolution::Route(peer_name) => {
let link = self.route_link(&peer_name).await?;
self.call_peer(link, &target_path, payload, headers, DEFAULT_HOPS, deadline)
.await
}
Resolution::Conflicted { owners } => Err(HandlerError::new(
ErrorCode::PeerUnreachable,
format!(
"destination node {target:?} has multiple live incarnations: {}",
owners.join(", ")
),
)),
Resolution::Unknown => Err(Self::teach_unknown_target(&snapshot, &target)),
}
}
pub(crate) async fn route_link(
&self,
peer_name: &str,
) -> Result<crate::node::PeerLink, HandlerError> {
self.peer(peer_name).await.ok_or_else(|| {
HandlerError::new(
ErrorCode::PeerUnreachable,
format!("no live connection to peer {peer_name:?}"),
)
})
}
async fn call_peer(
&self,
link: crate::node::PeerLink,
target_path: &str,
payload: Value,
headers: serde_json::Map<String, Value>,
hops: u8,
deadline: Instant,
) -> Result<Value, HandlerError> {
let reply = self
.call_peer_envelope(
link,
target_path,
Envelope::encode_payload(&payload),
headers,
hops,
deadline,
)
.await?;
Self::json_profile_payload(&reply.payload)
}
async fn call_peer_envelope(
&self,
link: crate::node::PeerLink,
target_path: &str,
payload: bytes::Bytes,
headers: serde_json::Map<String, Value>,
hops: u8,
deadline: Instant,
) -> Result<Envelope, HandlerError> {
let remaining = Self::remaining_unary_time(deadline)?;
let operation = async move {
let mut stream = link
.wire
.client_session()
.start(target_path, Kind::Request, payload, Some(hops), headers)
.await
.map_err(|error| {
HandlerError::new(ErrorCode::PeerUnreachable, error.to_string())
})?;
match stream.next().await {
Ok(Some(envelope)) if envelope.kind == Kind::Response => Ok(envelope),
Err(error) => Err(Self::client_error(error)),
Ok(Some(_)) => Err(HandlerError::new(
ErrorCode::Protocol,
"downstream call returned an unexpected frame",
)),
Ok(None) => Err(HandlerError::new(
ErrorCode::PeerUnreachable,
"downstream call did not complete",
)),
}
};
match n0_future::time::timeout(remaining, operation).await {
Ok(result) => result,
Err(_) => Err(HandlerError::new(
ErrorCode::PeerUnreachable,
"downstream call did not complete before its deadline",
)),
}
}
pub(crate) async fn resolve_unary_until(
&self,
target: &str,
deadline: Instant,
) -> Result<(Arc<NodeSnapshot>, Resolution), HandlerError> {
let mut route_changes = self.route_changes();
let readiness_waits = self.readiness_waits_for_destination(target);
let snapshot = self.snapshot.load_full();
let resolution = snapshot.node_core.resolve(target);
if !matches!(resolution, Resolution::Unknown) {
return Ok((snapshot, resolution));
}
if readiness_waits.is_empty() {
return Ok((snapshot, Resolution::Unknown));
}
let mut waiting = stream::FuturesUnordered::new();
for readiness in readiness_waits {
waiting.push(readiness.wait());
}
let mut restored = false;
while !waiting.is_empty() {
let remaining = Self::remaining_unary_time(deadline)?;
let wake = n0_future::time::timeout(remaining, async {
tokio::select! {
result = waiting.next() => (result, false),
changed = route_changes.changed() => {
let _ = changed;
(None, true)
}
}
})
.await
.map_err(|_| {
HandlerError::new(
ErrorCode::PeerUnreachable,
format!(
"recovery did not restore target {target:?} before the request deadline"
),
)
})?;
if !wake.1 && wake.0.is_some_and(|result| result.is_ok()) {
restored = true;
}
let snapshot = self.snapshot.load_full();
let resolution = snapshot.node_core.resolve(target);
if !matches!(resolution, Resolution::Unknown) {
return Ok((snapshot, resolution));
}
if waiting.is_empty() {
if restored {
return Ok((snapshot, Resolution::Unknown));
}
return Err(HandlerError::new(
ErrorCode::PeerUnreachable,
format!("recovery did not restore target {target:?}"),
));
}
}
Ok((snapshot, Resolution::Unknown))
}
pub(crate) async fn await_target_readiness(
&self,
target: &str,
cancellation: &unb_runtime::CancellationToken,
) -> Result<(), String> {
let mut route_changes = self.route_changes();
let readiness_waits = self.readiness_waits_for_destination(target);
if !matches!(
self.snapshot.load().node_core.resolve(target),
Resolution::Unknown
) {
return Ok(());
}
if readiness_waits.is_empty() {
return Err(format!(
"no recovering connection carried target {target:?}"
));
}
let mut waiting = stream::FuturesUnordered::new();
for readiness in readiness_waits {
waiting.push(readiness.wait());
}
loop {
if !matches!(
self.snapshot.load().node_core.resolve(target),
Resolution::Unknown
) {
return Ok(());
}
if waiting.is_empty() {
return Err(format!("recovery did not restore target {target:?}"));
}
tokio::select! {
biased;
() = cancellation.cancelled() => {
return Err(format!("source stream ended while waiting for target {target:?}"));
}
changed = route_changes.changed() => {
if changed.is_err() {
return Err("route readiness notifications closed".into());
}
}
result = waiting.next() => {
if result.is_some_and(|result| result.is_err()) && waiting.is_empty() {
return Err(format!("recovery did not restore target {target:?}"));
}
}
}
}
}
fn remaining_unary_time(deadline: Instant) -> Result<Duration, HandlerError> {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
Err(HandlerError::new(
ErrorCode::PeerUnreachable,
"unary request deadline elapsed",
))
} else {
Ok(remaining)
}
}
fn json_profile_payload(payload: &Bytes) -> Result<Value, HandlerError> {
if payload.is_empty() {
return Ok(Value::Null);
}
serde_json::from_slice(payload)
.map_err(|error| HandlerError::new(ErrorCode::Protocol, error.to_string()))
}
fn client_error(error: unb_runtime::ClientError) -> HandlerError {
match error {
unb_runtime::ClientError::Protocol { code, message, .. } => {
HandlerError::new(code, message)
}
unb_runtime::ClientError::Cancelled(_) => {
HandlerError::new(ErrorCode::Cancelled, error.to_string())
}
unb_runtime::ClientError::Invalid(message) => {
HandlerError::new(ErrorCode::InvalidInput, message)
}
_ => HandlerError::new(ErrorCode::PeerUnreachable, error.to_string()),
}
}
}