use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;
use futures::{FutureExt, StreamExt};
use tracing::{error, info, warn};
use crate::nats::HasNatsClient;
use crate::{
FlatbedRouteError, FlatbedWorkerError, HeaderMap, HeaderName, HeaderValue, Response,
ToFlatBuffer,
};
pub const REQUEST_ID_HEADER: &str = "x-request-id";
pub const ERROR_CODE_HEADER: &str = "x-error-code";
pub const ERROR_MESSAGE_HEADER: &str = "x-error-message";
pub const ERROR_STATUS_HEADER: &str = "x-error-status";
pub(crate) const CONTENT_TYPE_HEADER: &str = "content-type";
const JSON_CONTENT_TYPE: &str = "application/json";
const FLATBUFFER_CONTENT_TYPE: &str = "application/x-flatbuffers";
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NatsEncoding {
Json,
FlatBuffer,
}
impl NatsEncoding {
#[must_use]
pub fn from_content_type(content_type: Option<&str>) -> Self {
match content_type {
Some(ct) if ct.contains(JSON_CONTENT_TYPE) => Self::Json,
_ => Self::FlatBuffer,
}
}
#[must_use]
pub fn content_type(self) -> &'static str {
match self {
Self::Json => JSON_CONTENT_TYPE,
Self::FlatBuffer => FLATBUFFER_CONTENT_TYPE,
}
}
}
#[derive(Debug)]
pub struct NatsRequestParts {
pub subject: String,
pub headers: HeaderMap,
pub params: HashMap<String, String>,
pub request_id: String,
pub encoding: NatsEncoding,
}
#[derive(Debug)]
pub struct NatsReply {
pub payload: Vec<u8>,
pub headers: HeaderMap,
}
pub type NatsHandlerFn =
fn(
NatsRequestParts,
Vec<u8>,
Arc<dyn Any + Send + Sync>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = NatsReply> + Send>>;
#[derive(Clone, Copy)]
pub struct NatsRouteInfo {
pub subject: &'static str,
pub wire_subject: &'static str,
pub params: &'static [(&'static str, usize)],
pub queue: Option<&'static str>,
pub request_type: &'static str,
pub response_type: &'static str,
pub handler: NatsHandlerFn,
}
impl std::fmt::Debug for NatsRouteInfo {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("NatsRouteInfo")
.field("subject", &self.subject)
.field("wire_subject", &self.wire_subject)
.field("params", &self.params)
.field("queue", &self.queue)
.field("request_type", &self.request_type)
.field("response_type", &self.response_type)
.field("handler", &"<async_function>")
.finish()
}
}
inventory::collect!(NatsRouteInfo);
#[must_use]
pub fn get_nats_routes() -> Vec<&'static NatsRouteInfo> {
inventory::iter::<NatsRouteInfo>.into_iter().collect()
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NatsRouteConflict {
pub first_subject: String,
pub second_subject: String,
}
impl std::fmt::Display for NatsRouteConflict {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"subjects '{}' and '{}' overlap: a message matching both is answered twice",
self.first_subject, self.second_subject
)
}
}
impl std::error::Error for NatsRouteConflict {}
pub fn validate_nats_routes() -> Result<(), NatsRouteConflict> {
match find_conflict(get_nats_routes()) {
Some(conflict) => Err(conflict),
None => Ok(()),
}
}
fn subjects_overlap(first: &str, second: &str) -> bool {
let first: Vec<&str> = first.split('.').collect();
let second: Vec<&str> = second.split('.').collect();
first.len() == second.len()
&& first
.iter()
.zip(&second)
.all(|(a, b)| *a == "*" || *b == "*" || a == b)
}
fn find_conflict<'a>(
routes: impl IntoIterator<Item = &'a NatsRouteInfo>,
) -> Option<NatsRouteConflict> {
let routes: Vec<&NatsRouteInfo> = routes.into_iter().collect();
routes
.iter()
.enumerate()
.flat_map(|(i, first)| routes[i + 1..].iter().map(move |second| (*first, *second)))
.find(|(first, second)| subjects_overlap(first.wire_subject, second.wire_subject))
.map(|(first, second)| NatsRouteConflict {
first_subject: first.subject.to_string(),
second_subject: second.subject.to_string(),
})
}
#[must_use]
pub fn reply_ok<T: ToFlatBuffer>(
encoding: NatsEncoding,
request_id: &str,
mut response: Response<T>,
) -> NatsReply {
let (payload, content_type) = match response.take_raw() {
Some((bytes, raw_content_type)) => (bytes, raw_content_type.into_owned()),
None if encoding == NatsEncoding::Json => match serde_json::to_vec(&response.body) {
Ok(bytes) => (bytes, encoding.content_type().to_string()),
Err(e) => {
return reply_err(
encoding,
request_id,
&FlatbedRouteError::internal(format!("JSON serialization error: {e}"))
.code("SERIALIZATION_ERROR"),
)
}
},
None => (
response.body.to_flatbuffer(),
encoding.content_type().to_string(),
),
};
let mut headers = copy_headers(&response.headers);
for name in [ERROR_CODE_HEADER, ERROR_MESSAGE_HEADER, ERROR_STATUS_HEADER] {
headers.remove(name);
}
set_header(&mut headers, CONTENT_TYPE_HEADER, &content_type);
set_header(&mut headers, REQUEST_ID_HEADER, request_id);
NatsReply { payload, headers }
}
#[derive(serde::Serialize)]
struct ErrorBody<'a, D: serde::Serialize> {
code: &'a str,
message: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
details: &'a Option<D>,
}
#[must_use]
pub fn reply_err<D: ToFlatBuffer>(
encoding: NatsEncoding,
request_id: &str,
err: &FlatbedRouteError<D>,
) -> NatsReply {
let code = if err.code.is_empty() {
"ERROR"
} else {
&err.code
};
let payload = match encoding {
NatsEncoding::Json => serde_json::to_vec(&ErrorBody {
code,
message: &err.message,
details: &err.details,
})
.unwrap_or_default(),
NatsEncoding::FlatBuffer => err
.details
.as_ref()
.map(ToFlatBuffer::to_flatbuffer)
.unwrap_or_default(),
};
let mut headers = copy_headers(&err.headers);
set_header(&mut headers, CONTENT_TYPE_HEADER, encoding.content_type());
set_header(&mut headers, REQUEST_ID_HEADER, request_id);
set_header(&mut headers, ERROR_CODE_HEADER, code);
set_header(&mut headers, ERROR_MESSAGE_HEADER, &err.message);
set_header(&mut headers, ERROR_STATUS_HEADER, err.status.as_str());
NatsReply { payload, headers }
}
fn copy_headers(source: &HeaderMap) -> HeaderMap {
let mut headers = HeaderMap::new();
for (name, value) in source.iter() {
headers.append(name.clone(), value.clone());
}
headers
}
pub(crate) fn set_header(headers: &mut HeaderMap, name: &str, value: &str) {
let flattened: String = value
.chars()
.map(|c| if c.is_control() { ' ' } else { c })
.collect();
let (Ok(name), Ok(value)) = (
HeaderName::try_from(name),
HeaderValue::try_from(flattened.as_str()),
) else {
return;
};
headers.insert(name, value);
}
pub(crate) fn from_nats_headers(source: Option<&async_nats::HeaderMap>) -> HeaderMap {
let mut headers = HeaderMap::new();
let Some(source) = source else {
return headers;
};
for (name, values) in source.iter() {
let Ok(name) = HeaderName::try_from(AsRef::<str>::as_ref(name)) else {
continue;
};
for value in values {
let Ok(value) = HeaderValue::try_from(value.as_str()) else {
continue;
};
headers.append(name.clone(), value);
}
}
headers
}
pub(crate) fn to_nats_headers(source: &HeaderMap) -> async_nats::HeaderMap {
let mut headers = async_nats::HeaderMap::new();
for (name, value) in source.iter() {
let Ok(value) = value.to_str() else {
continue;
};
headers.append(name.as_str(), value);
}
headers
}
pub(crate) fn header_value<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
headers.get(name).and_then(|value| value.to_str().ok())
}
fn extract_params(params: &[(&'static str, usize)], subject: &str) -> HashMap<String, String> {
if params.is_empty() {
return HashMap::new();
}
let tokens: Vec<&str> = subject.split('.').collect();
params
.iter()
.filter_map(|(name, index)| {
tokens
.get(*index)
.map(|token| ((*name).to_string(), (*token).to_string()))
})
.collect()
}
async fn dispatch(
route: NatsRouteInfo,
message: &async_nats::Message,
ctx: Arc<dyn Any + Send + Sync>,
) -> NatsReply {
let headers = from_nats_headers(message.headers.as_ref());
let encoding = NatsEncoding::from_content_type(header_value(&headers, CONTENT_TYPE_HEADER));
let request_id = header_value(&headers, REQUEST_ID_HEADER)
.map(str::to_string)
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let parts = NatsRequestParts {
subject: message.subject.to_string(),
params: extract_params(route.params, &message.subject),
headers,
request_id: request_id.clone(),
encoding,
};
let handled = (route.handler)(parts, message.payload.to_vec(), ctx);
match std::panic::AssertUnwindSafe(handled).catch_unwind().await {
Ok(reply) => reply,
Err(_) => {
error!(
subject = %message.subject,
"nats_route handler panicked; answering with an error reply"
);
reply_err(
encoding,
&request_id,
&FlatbedRouteError::internal("handler panicked").code("PANIC"),
)
}
}
}
pub async fn run_nats_route<C>(
ctx: Arc<dyn Any + Send + Sync>,
route: NatsRouteInfo,
) -> Result<(), FlatbedWorkerError>
where
C: HasNatsClient + Send + Sync + 'static,
{
let ctx: Arc<C> = ctx.downcast::<C>().map_err(|_| {
FlatbedWorkerError::new(
"nats_route_context",
format!(
"[{}] context type mismatch: the service booted with a different context",
route.subject
),
)
})?;
let client = HasNatsClient::nats_client(&*ctx).clone();
let subscribe = match route.queue {
Some(queue) => {
client
.queue_subscribe(route.wire_subject, queue.to_string())
.await
}
None => client.subscribe(route.wire_subject).await,
};
let mut subscription = subscribe.map_err(|e| {
FlatbedWorkerError::new(
"nats_route_subscribe",
format!("[{}] subscribe failed: {}", route.subject, e),
)
})?;
info!(
subject = route.subject,
wire_subject = route.wire_subject,
queue = route.queue.unwrap_or("<none>"),
"nats route subscribed"
);
while let Some(message) = subscription.next().await {
let Some(reply_subject) = message.reply.clone() else {
warn!(
subject = %message.subject,
"message has no reply subject; nothing to answer"
);
continue;
};
let ctx: Arc<dyn Any + Send + Sync> = ctx.clone();
let client = client.clone();
tokio::spawn(async move {
let reply = dispatch(route, &message, ctx).await;
let headers = to_nats_headers(&reply.headers);
if let Err(e) = client
.publish_with_headers(reply_subject, headers, reply.payload.into())
.await
{
error!(subject = route.subject, error = %e, "reply publish failed");
}
});
}
Err(FlatbedWorkerError::new(
"nats_route_subscription_ended",
format!(
"[{}] subscription ended; the subject is no longer answered",
route.subject
),
))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::StatusCode;
#[test]
fn encoding_defaults_to_flatbuffers_for_absent_or_unknown_content_type() {
assert_eq!(
NatsEncoding::from_content_type(None),
NatsEncoding::FlatBuffer
);
assert_eq!(
NatsEncoding::from_content_type(Some("text/plain")),
NatsEncoding::FlatBuffer
);
assert_eq!(
NatsEncoding::from_content_type(Some("application/x-flatbuffers")),
NatsEncoding::FlatBuffer
);
}
#[test]
fn encoding_selects_json_from_a_parameterized_content_type() {
assert_eq!(
NatsEncoding::from_content_type(Some("application/json; charset=utf-8")),
NatsEncoding::Json
);
assert_eq!(NatsEncoding::Json.content_type(), "application/json");
assert_eq!(
NatsEncoding::FlatBuffer.content_type(),
"application/x-flatbuffers"
);
}
fn unreachable_handler(
_: NatsRequestParts,
_: Vec<u8>,
_: Arc<dyn Any + Send + Sync>,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = NatsReply> + Send>> {
unreachable!("the conflict scan never dispatches")
}
fn route(subject: &'static str, wire_subject: &'static str) -> NatsRouteInfo {
NatsRouteInfo {
subject,
wire_subject,
params: &[],
queue: None,
request_type: "Req",
response_type: "Res",
handler: unreachable_handler,
}
}
#[test]
fn distinct_wire_subjects_do_not_conflict() {
let routes = [
route("a.{id}.status", "a.*.status"),
route("a.{id}.config", "a.*.config"),
route("b.report", "b.report"),
];
assert!(find_conflict(&routes).is_none());
}
#[test]
fn patterns_differing_only_in_capture_names_conflict() {
let routes = [
route("a.report", "a.report"),
route("plonk.sat.{id}.status", "plonk.sat.*.status"),
route("plonk.sat.{name}.status", "plonk.sat.*.status"),
];
let conflict = find_conflict(&routes).expect("the colliding pair must be reported");
assert_eq!(conflict.first_subject, "plonk.sat.{id}.status");
assert_eq!(conflict.second_subject, "plonk.sat.{name}.status");
assert_eq!(
conflict.to_string(),
"subjects 'plonk.sat.{id}.status' and 'plonk.sat.{name}.status' \
overlap: a message matching both is answered twice",
);
}
#[test]
fn a_wildcard_conflicts_with_a_literal_in_the_same_position() {
let routes = [
route("plonk.sat.{id}.status", "plonk.sat.*.status"),
route("plonk.sat.alpha.status", "plonk.sat.alpha.status"),
];
let conflict = find_conflict(&routes).expect("the overlapping pair must be reported");
assert_eq!(conflict.first_subject, "plonk.sat.{id}.status");
assert_eq!(conflict.second_subject, "plonk.sat.alpha.status");
}
#[test]
fn a_wildcard_does_not_conflict_across_different_token_counts() {
let routes = [
route("plonk.{id}", "plonk.*"),
route("plonk.{id}.status", "plonk.*.status"),
];
assert!(find_conflict(&routes).is_none());
}
#[test]
fn params_are_read_off_the_concrete_subject_by_token_index() {
let params = extract_params(&[("id", 2), ("verb", 4)], "plonk.satellite.x07.call.status");
assert_eq!(params.get("id").map(String::as_str), Some("x07"));
assert_eq!(params.get("verb").map(String::as_str), Some("status"));
}
#[test]
fn params_are_empty_for_a_pattern_without_tokens() {
assert!(extract_params(&[], "plonk.ground.report").is_empty());
}
#[test]
fn error_reply_carries_code_message_and_status_headers() {
let err = FlatbedRouteError::not_found("no such satellite").code("NOT_FOUND");
let reply = reply_err(NatsEncoding::FlatBuffer, "req-1", &err);
assert_eq!(
header_value(&reply.headers, ERROR_CODE_HEADER),
Some("NOT_FOUND")
);
assert_eq!(
header_value(&reply.headers, ERROR_MESSAGE_HEADER),
Some("no such satellite")
);
assert_eq!(
header_value(&reply.headers, ERROR_STATUS_HEADER),
Some("404")
);
assert_eq!(
header_value(&reply.headers, REQUEST_ID_HEADER),
Some("req-1")
);
assert!(
reply.payload.is_empty(),
"a FlatBuffers error without details has an empty payload"
);
}
#[test]
fn error_reply_without_a_code_falls_back_to_a_generic_one() {
let err = FlatbedRouteError::with_status(StatusCode::IM_A_TEAPOT, "brew failed");
let reply = reply_err(NatsEncoding::Json, "req-2", &err);
assert_eq!(
header_value(&reply.headers, ERROR_CODE_HEADER),
Some("ERROR")
);
assert_eq!(
header_value(&reply.headers, ERROR_STATUS_HEADER),
Some("418")
);
let body: serde_json::Value = serde_json::from_slice(&reply.payload).unwrap();
assert_eq!(body["code"], "ERROR");
assert_eq!(body["message"], "brew failed");
assert!(
body.get("details").is_none(),
"an error without details omits the field"
);
}
#[test]
fn handler_headers_never_displace_the_frameworks_own() {
let err = FlatbedRouteError::not_found("gone")
.code("NOT_FOUND")
.header("x-error-code", "OK")
.header("content-type", "text/plain")
.header("x-trace", "t-9");
let reply = reply_err(NatsEncoding::Json, "req-4", &err);
assert_eq!(
header_value(&reply.headers, ERROR_CODE_HEADER),
Some("NOT_FOUND")
);
assert_eq!(
header_value(&reply.headers, CONTENT_TYPE_HEADER),
Some("application/json")
);
assert_eq!(header_value(&reply.headers, "x-trace"), Some("t-9"));
}
#[test]
fn a_multi_line_error_message_is_flattened_onto_the_header() {
let err = FlatbedRouteError::bad_request("line one\nline two\r\tend").code("BAD");
let reply = reply_err(NatsEncoding::FlatBuffer, "req-5", &err);
assert_eq!(
header_value(&reply.headers, ERROR_MESSAGE_HEADER),
Some("line one line two end")
);
}
#[test]
fn a_success_reply_never_carries_the_error_headers() {
let reply = reply_ok(
NatsEncoding::Json,
"req-6",
Response::ok(()).header("x-error-code", "NOT_FOUND"),
);
assert_eq!(header_value(&reply.headers, ERROR_CODE_HEADER), None);
}
#[test]
fn raw_response_bytes_are_published_verbatim_under_their_own_content_type() {
let reply = reply_ok(
NatsEncoding::Json,
"req-3",
Response::raw(b"id,count\n1,2\n".to_vec(), "text/csv"),
);
assert_eq!(reply.payload, b"id,count\n1,2\n");
assert_eq!(
header_value(&reply.headers, CONTENT_TYPE_HEADER),
Some("text/csv")
);
}
#[test]
fn nats_headers_round_trip_through_the_http_header_map() {
let mut source = async_nats::HeaderMap::new();
source.append("Content-Type", "application/json");
source.append("X-Request-Id", "abc");
let converted = from_nats_headers(Some(&source));
assert_eq!(
header_value(&converted, CONTENT_TYPE_HEADER),
Some("application/json"),
"NATS header names are matched case-insensitively after conversion"
);
assert_eq!(header_value(&converted, REQUEST_ID_HEADER), Some("abc"));
let back = to_nats_headers(&converted);
assert_eq!(
back.get("content-type")
.map(async_nats::HeaderValue::as_str),
Some("application/json")
);
}
}