use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use crate::core::auth::Identity;
use crate::core::types::Capabilities;
use crate::protocol::wire::{CallError, ResponseEnvelope};
use crate::registry::context::{AbortPolicy, OperationContext, ScopedPeerEnv};
use crate::registry::env::LocalOperationEnv;
use crate::registry::registration::OperationRegistry;
use crate::registry::spec::{AccessResult, Visibility};
use futures::stream::BoxStream;
use serde_json::Value;
const SERVICES_SCHEMA: &str = "services/schema";
pub const DEFAULT_DEADLINE: Duration = Duration::from_secs(30);
pub struct GatewayDispatch {
registry: Arc<OperationRegistry>,
deadline: Option<Duration>,
invoke_count: AtomicUsize,
}
impl GatewayDispatch {
pub fn new(registry: Arc<OperationRegistry>) -> Self {
Self {
registry,
deadline: Some(DEFAULT_DEADLINE),
invoke_count: AtomicUsize::new(0),
}
}
pub fn with_deadline(mut self, deadline: Option<Duration>) -> Self {
self.deadline = deadline;
self
}
pub fn registry(&self) -> &Arc<OperationRegistry> {
&self.registry
}
pub fn invoke_count(&self) -> usize {
self.invoke_count.load(Ordering::Relaxed)
}
pub async fn invoke(
&self,
identity: Option<Identity>,
op: &str,
input: Value,
) -> ResponseEnvelope {
self.invoke_count.fetch_add(1, Ordering::Relaxed);
let operation_name = strip_leading_slash(op).to_string();
if let Some(error) =
schema_via_call_denial(&self.registry, &operation_name, &input, identity.as_ref())
{
return ResponseEnvelope::error(uuid::Uuid::new_v4().to_string(), error);
}
let request_id = uuid::Uuid::new_v4().to_string();
let context = self.build_root_context(&request_id, &operation_name, identity);
let fut = self.registry.invoke(&operation_name, input, context);
match self.deadline {
Some(deadline) => match tokio::time::timeout(deadline, fut).await {
Ok(envelope) => envelope,
Err(_elapsed) => ResponseEnvelope::error(
request_id,
CallError::timeout(format!(
"operation did not complete within the {deadline:?} dispatch deadline"
)),
),
},
None => fut.await,
}
}
pub fn invoke_streaming(
&self,
identity: Option<Identity>,
op: &str,
input: Value,
) -> BoxStream<'static, ResponseEnvelope> {
let operation_name = strip_leading_slash(op).to_string();
if let Some(error) =
schema_via_call_denial(&self.registry, &operation_name, &input, identity.as_ref())
{
let request_id = uuid::Uuid::new_v4().to_string();
return Box::pin(futures::stream::once(async move {
ResponseEnvelope::error(request_id, error)
}));
}
let request_id = uuid::Uuid::new_v4().to_string();
let context = self.build_root_context_streaming(&request_id, &operation_name, identity);
self.registry
.invoke_streaming(&operation_name, input, context)
}
pub async fn invoke_sink(
&self,
identity: Option<Identity>,
op: &str,
input: Value,
publish_stream: crate::registry::registration::PublishStream,
) -> ResponseEnvelope {
let operation_name = strip_leading_slash(op).to_string();
let request_id = uuid::Uuid::new_v4().to_string();
let context = self.build_root_context_sink(&request_id, &operation_name, identity);
let fut = self
.registry
.invoke_sink(&operation_name, input, publish_stream, context);
match self.deadline {
Some(deadline) => match tokio::time::timeout(deadline, fut).await {
Ok(envelope) => envelope,
Err(_elapsed) => ResponseEnvelope::error(
request_id,
CallError::timeout(format!(
"operation did not complete within the {deadline:?} dispatch deadline"
)),
),
},
None => fut.await,
}
}
fn build_root_context_sink(
&self,
request_id: &str,
operation_name: &str,
identity: Option<Identity>,
) -> OperationContext {
self.build_root_context_inner(request_id, operation_name, identity, false)
}
fn build_root_context(
&self,
request_id: &str,
operation_name: &str,
identity: Option<Identity>,
) -> OperationContext {
self.build_root_context_inner(request_id, operation_name, identity, true)
}
fn build_root_context_streaming(
&self,
request_id: &str,
operation_name: &str,
identity: Option<Identity>,
) -> OperationContext {
self.build_root_context_inner(request_id, operation_name, identity, false)
}
fn build_root_context_inner(
&self,
request_id: &str,
operation_name: &str,
identity: Option<Identity>,
bounded: bool,
) -> OperationContext {
let registration = self.registry.registration(operation_name);
let (composition_authority, capabilities, scoped_env) = match registration {
Some(r) => (
r.composition_authority.clone(),
r.capabilities.clone(),
r.scoped_env.clone().unwrap_or_else(ScopedPeerEnv::empty),
),
None => (None, Capabilities::new(), ScopedPeerEnv::empty()),
};
let env: Arc<dyn crate::registry::env::OperationEnv + Send + Sync> =
Arc::new(LocalOperationEnv::new(Arc::clone(&self.registry)));
OperationContext {
request_id: request_id.to_string(),
parent_request_id: None,
identity,
handler_identity: composition_authority,
forwarded_for: None,
capabilities,
metadata: HashMap::new(),
deadline: bounded
.then_some(self.deadline)
.flatten()
.map(|deadline| Instant::now() + deadline),
scoped_env,
env,
abort_policy: AbortPolicy::default(),
internal: false,
ownership: None,
}
}
}
fn strip_leading_slash(operation_id: &str) -> &str {
operation_id.strip_prefix('/').unwrap_or(operation_id)
}
fn schema_via_call_denial(
registry: &OperationRegistry,
operation: &str,
input: &Value,
identity: Option<&Identity>,
) -> Option<CallError> {
let name = input.get("name").and_then(Value::as_str)?;
if !matches!(
registry.registration(operation),
Some(registration) if registration.spec.name == SERVICES_SCHEMA
) {
return None;
}
schema_disclosure_denial(registry, name, identity)
}
pub fn schema_disclosure_denial(
registry: &OperationRegistry,
operation: &str,
identity: Option<&Identity>,
) -> Option<CallError> {
let name = strip_leading_slash(operation);
let registration = registry.registration(name)?;
if registration.spec.visibility == Visibility::Internal {
return Some(CallError::not_found(operation));
}
if let AccessResult::Forbidden(message) =
registration.spec.access_control.check(identity, None, None)
{
return Some(CallError::forbidden(message));
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use crate::registry::registration::{
make_handler, make_sink_handler, make_streaming_handler, HandlerKind, HandlerRegistration,
OperationProvenance,
};
use crate::registry::spec::{AccessControl, OperationSpec, OperationType, Visibility};
use futures::StreamExt;
fn spec(name: &str, visibility: Visibility, op_type: OperationType) -> OperationSpec {
OperationSpec::new(
name,
op_type,
visibility,
serde_json::json!({}),
serde_json::json!({}),
vec![],
AccessControl::default(),
None,
)
}
fn echo_registry() -> Arc<OperationRegistry> {
let registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
spec("echo/run", Visibility::External, OperationType::Query),
HandlerKind::Once(make_handler(|input, ctx| async move {
ResponseEnvelope::ok(ctx.request_id, input)
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
Arc::new(registry)
}
#[tokio::test]
async fn invoke_external_op_round_trips() {
let dispatch = GatewayDispatch::new(echo_registry());
let envelope = dispatch
.invoke(None, "/echo/run", serde_json::json!({ "x": 1 }))
.await;
assert!(envelope.result.is_ok(), "expected ok, got {envelope:?}");
}
#[tokio::test]
async fn invoke_unknown_op_returns_not_found() {
let dispatch = GatewayDispatch::new(echo_registry());
let envelope = dispatch
.invoke(None, "/missing/op", serde_json::json!({}))
.await;
match envelope.result {
Err(error) => assert_eq!(error.code, "NOT_FOUND"),
Ok(v) => panic!("expected error, got {v:?}"),
}
}
#[tokio::test]
async fn invoke_internal_op_returns_not_found() {
let registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
spec("internal/op", Visibility::Internal, OperationType::Query),
HandlerKind::Once(make_handler(|_input, ctx| async move {
ResponseEnvelope::ok(ctx.request_id, serde_json::json!({}))
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
let dispatch = GatewayDispatch::new(Arc::new(registry));
let envelope = dispatch
.invoke(None, "/internal/op", serde_json::json!({}))
.await;
match envelope.result {
Err(error) => assert_eq!(error.code, "NOT_FOUND"),
Ok(v) => panic!("expected error, got {v:?}"),
}
}
#[tokio::test]
async fn invoke_enforces_the_configured_deadline_on_a_hung_handler() {
let registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
spec("hung/op", Visibility::External, OperationType::Query),
HandlerKind::Once(make_handler(|_input, ctx| async move {
tokio::time::sleep(Duration::from_secs(120)).await;
ResponseEnvelope::ok(ctx.request_id, serde_json::json!({}))
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
let dispatch =
GatewayDispatch::new(Arc::new(registry)).with_deadline(Some(Duration::from_millis(50)));
let started = std::time::Instant::now();
let envelope = dispatch
.invoke(None, "/hung/op", serde_json::json!({}))
.await;
assert!(
started.elapsed() < Duration::from_secs(10),
"the 50 ms deadline must fire well before the 120 s handler sleep"
);
match envelope.result {
Err(error) => {
assert_eq!(error.code, "TIMEOUT");
assert!(error.retryable, "the deadline error is retryable");
}
Ok(v) => panic!("expected a TIMEOUT error, got {v:?}"),
}
}
#[tokio::test]
async fn invoke_completes_within_the_deadline_for_a_fast_handler() {
let dispatch = GatewayDispatch::new(echo_registry());
let envelope = dispatch
.invoke(None, "/echo/run", serde_json::json!({}))
.await;
assert!(envelope.result.is_ok(), "a fast handler must not time out");
}
#[tokio::test]
async fn invoke_with_no_deadline_completes_a_slow_handler() {
let registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
spec("slow/op", Visibility::External, OperationType::Query),
HandlerKind::Once(make_handler(|_input, ctx| async move {
tokio::time::sleep(Duration::from_millis(80)).await;
ResponseEnvelope::ok(ctx.request_id, serde_json::json!({}))
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
let dispatch = GatewayDispatch::new(Arc::new(registry)).with_deadline(None);
let envelope = dispatch
.invoke(None, "/slow/op", serde_json::json!({}))
.await;
assert!(
envelope.result.is_ok(),
"a no-deadline spine must not time out, got {envelope:?}"
);
}
#[tokio::test]
async fn invoke_sink_enforces_the_configured_deadline_on_a_hung_sink_handler() {
let registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
spec("hung/sink", Visibility::External, OperationType::Pub),
HandlerKind::Sink(make_sink_handler(|_input, ctx, _chunks| async move {
tokio::time::sleep(Duration::from_secs(120)).await;
ResponseEnvelope::ok(ctx.request_id, serde_json::json!({}))
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
let dispatch =
GatewayDispatch::new(Arc::new(registry)).with_deadline(Some(Duration::from_millis(50)));
let chunks: crate::registry::registration::PublishStream =
Box::pin(futures::stream::empty());
let started = std::time::Instant::now();
let envelope = dispatch
.invoke_sink(None, "/hung/sink", serde_json::json!({}), chunks)
.await;
assert!(
started.elapsed() < Duration::from_secs(10),
"the 50 ms deadline must fire well before the 120 s handler sleep"
);
match envelope.result {
Err(error) => {
assert_eq!(error.code, "TIMEOUT");
assert!(error.retryable, "the deadline error is retryable");
}
Ok(v) => panic!("expected a TIMEOUT error, got {v:?}"),
}
}
#[tokio::test]
async fn invoke_sink_completes_within_the_deadline_for_a_fast_sink_handler() {
let registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
spec("fast/sink", Visibility::External, OperationType::Pub),
HandlerKind::Sink(make_sink_handler(|_input, ctx, mut chunks| async move {
let mut count = 0u64;
while let Some(chunk) = chunks.next().await {
if chunk.is_ok() {
count += 1;
}
}
ResponseEnvelope::ok(ctx.request_id, serde_json::json!({ "count": count }))
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
let dispatch = GatewayDispatch::new(Arc::new(registry));
let chunks: crate::registry::registration::PublishStream =
Box::pin(futures::stream::iter(vec![
Ok(serde_json::json!({ "n": 1 })),
Ok(serde_json::json!({ "n": 2 })),
]));
let envelope = dispatch
.invoke_sink(None, "/fast/sink", serde_json::json!({}), chunks)
.await;
assert!(
envelope.result.is_ok(),
"a fast sink handler must not time out, got {envelope:?}"
);
}
#[tokio::test]
async fn invoke_sink_with_no_deadline_completes_a_slow_sink_handler() {
let registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
spec("slow/sink", Visibility::External, OperationType::Pub),
HandlerKind::Sink(make_sink_handler(|_input, ctx, mut chunks| async move {
while let Some(chunk) = chunks.next().await {
let _ = chunk;
}
tokio::time::sleep(Duration::from_millis(80)).await;
ResponseEnvelope::ok(ctx.request_id, serde_json::json!({}))
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
let dispatch = GatewayDispatch::new(Arc::new(registry)).with_deadline(None);
let chunks: crate::registry::registration::PublishStream = Box::pin(futures::stream::iter(
vec![Ok(serde_json::json!({ "n": 1 }))],
));
let envelope = dispatch
.invoke_sink(None, "/slow/sink", serde_json::json!({}), chunks)
.await;
assert!(
envelope.result.is_ok(),
"a no-deadline spine must not time out a slow sink, got {envelope:?}"
);
}
#[tokio::test]
async fn streaming_sub_op_streams_envelopes() {
let registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
spec("tick/stream", Visibility::External, OperationType::Sub),
HandlerKind::Stream(make_streaming_handler(|input, ctx| {
let count = input.get("count").and_then(|v| v.as_u64()).unwrap_or(2);
let request_id = ctx.request_id.clone();
futures::stream::iter(0..count).map(move |i| {
ResponseEnvelope::ok(request_id.clone(), serde_json::json!({ "tick": i }))
})
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
let dispatch = GatewayDispatch::new(Arc::new(registry));
let mut stream =
dispatch.invoke_streaming(None, "/tick/stream", serde_json::json!({ "count": 3 }));
let mut ticks = Vec::new();
while let Some(envelope) = stream.next().await {
ticks.push(envelope);
}
assert_eq!(ticks.len(), 3);
}
#[tokio::test]
async fn invoke_count_counts_invokes_only() {
let dispatch = GatewayDispatch::new(echo_registry());
let _ = dispatch
.invoke(None, "/echo/run", serde_json::json!({}))
.await;
let _ = dispatch.invoke_streaming(None, "/echo/run", serde_json::json!({}));
assert_eq!(dispatch.invoke_count(), 1);
}
fn registry_with_services_schema_over(inner_ops: Vec<OperationSpec>) -> Arc<OperationRegistry> {
use crate::registry::discovery::{services_schema_handler, services_schema_spec};
let inner = Arc::new({
let registry = OperationRegistry::new();
for op_spec in &inner_ops {
registry
.register(HandlerRegistration::new(
op_spec.clone(),
HandlerKind::Once(make_handler(|input, ctx| async move {
ResponseEnvelope::ok(ctx.request_id, input)
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
}
registry
});
let registry = OperationRegistry::new();
for op_spec in &inner_ops {
registry
.register(HandlerRegistration::new(
op_spec.clone(),
HandlerKind::Once(make_handler(|input, ctx| async move {
ResponseEnvelope::ok(ctx.request_id, input)
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
}
registry
.register(HandlerRegistration::new(
services_schema_spec(),
HandlerKind::Once(services_schema_handler(Arc::clone(&inner))),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
Arc::new(registry)
}
#[tokio::test]
async fn invoke_of_services_schema_with_internal_inner_name_is_blocked_pre_dispatch() {
let registry = registry_with_services_schema_over(vec![spec(
"secret/op",
Visibility::Internal,
OperationType::Query,
)]);
let dispatch = GatewayDispatch::new(registry);
let envelope = dispatch
.invoke(
None,
"services/schema",
serde_json::json!({ "name": "secret/op" }),
)
.await;
match envelope.result {
Err(error) => {
assert_eq!(error.code, "NOT_FOUND");
assert!(error.message.contains("secret/op"));
}
Ok(v) => panic!("the internal op spec must not be returned, got {v:?}"),
}
}
#[tokio::test]
async fn invoke_of_services_schema_with_authorized_inner_name_still_projects() {
let registry = registry_with_services_schema_over(vec![spec(
"public/op",
Visibility::External,
OperationType::Query,
)]);
let dispatch = GatewayDispatch::new(registry);
let envelope = dispatch
.invoke(
None,
"services/schema",
serde_json::json!({ "name": "public/op" }),
)
.await;
assert!(
envelope.result.is_ok(),
"an allowed inner name must still project, got {envelope:?}"
);
assert_eq!(
envelope
.result
.as_ref()
.ok()
.and_then(|v| v.get("name"))
.and_then(Value::as_str),
Some("public/op")
);
}
#[tokio::test]
async fn invoke_of_services_schema_with_forbidden_inner_name_is_denied() {
let restricted = OperationSpec::new(
"admin/op",
OperationType::Query,
Visibility::External,
serde_json::json!({}),
serde_json::json!({}),
vec![],
AccessControl {
required_scopes: vec!["admin".to_string()],
..Default::default()
},
None,
);
let registry = registry_with_services_schema_over(vec![restricted]);
let dispatch = GatewayDispatch::new(registry);
let envelope = dispatch
.invoke(
Some(Identity {
id: "user".to_string(),
scopes: vec!["user".to_string()],
resources: HashMap::new(),
}),
"services/schema",
serde_json::json!({ "name": "admin/op" }),
)
.await;
match envelope.result {
Err(error) => assert_eq!(error.code, "FORBIDDEN"),
Ok(v) => panic!("the ACL-denied spec must not be returned, got {v:?}"),
}
}
#[tokio::test]
async fn invoke_streaming_of_services_schema_with_internal_inner_name_is_blocked() {
let registry = registry_with_services_schema_over(vec![spec(
"secret/op",
Visibility::Internal,
OperationType::Query,
)]);
let dispatch = GatewayDispatch::new(registry);
let mut stream = dispatch.invoke_streaming(
None,
"services/schema",
serde_json::json!({ "name": "secret/op" }),
);
let envelopes: Vec<ResponseEnvelope> = stream.by_ref().collect().await;
assert_eq!(envelopes.len(), 1, "the guard error is the only item");
match &envelopes[0].result {
Err(error) => assert_eq!(error.code, "NOT_FOUND"),
Ok(v) => panic!("the internal op spec must not be returned, got {v:?}"),
}
}
#[test]
fn schema_disclosure_denial_hides_internal_ops() {
let registry = OperationRegistry::new();
registry
.register(HandlerRegistration::new(
spec("secret/op", Visibility::Internal, OperationType::Query),
HandlerKind::Once(make_handler(|_input, ctx| async move {
ResponseEnvelope::ok(ctx.request_id, serde_json::json!({}))
})),
OperationProvenance::Local,
None,
None,
Capabilities::new(),
))
.unwrap();
let error = schema_disclosure_denial(®istry, "secret/op", None).unwrap();
assert_eq!(error.code, "NOT_FOUND");
assert!(schema_disclosure_denial(®istry, "missing/op", None).is_none());
}
#[test]
fn schema_disclosure_denial_allows_unrestricted_ops_without_identity() {
let error = schema_disclosure_denial(&echo_registry(), "echo/run", None);
assert!(
error.is_none(),
"default-ACL ops are fetchable unauthenticated"
);
}
}