diff --git a/crates/openshell-core/src/extension_protocol.rs b/crates/openshell-core/src/extension_protocol.rs index 3ca57e07e7..e9a63f0f9a 100644 --- a/crates/openshell-core/src/extension_protocol.rs +++ b/crates/openshell-core/src/extension_protocol.rs @@ -37,6 +37,14 @@ pub const COMPUTE_UPDATE_SANDBOX_CONFIGURATION: &str = /// `Describe` instead of calling the deprecated legacy HTTP hooks. pub const SUPERVISOR_MIDDLEWARE_HTTP_HOOKS: &str = "openshell.supervisor-middleware.http-hooks"; +/// Capability for supervisor middleware that implements TCP connection hooks +/// (`TcpHooks.EvaluateConnection`). +/// +/// A service with a `TCP_CONNECTION` binding must list this capability in +/// `required_capabilities`, so peers without TCP hooks reject it at +/// `Describe`. +pub const SUPERVISOR_MIDDLEWARE_TCP_HOOKS: &str = "openshell.supervisor-middleware.tcp-hooks"; + const MAX_IMPLEMENTATION_NAME_BYTES: usize = 128; const MAX_IMPLEMENTATION_VERSION_BYTES: usize = 128; const MAX_CAPABILITY_BYTES: usize = 128; @@ -145,6 +153,7 @@ pub fn gateway_metadata(family: ExtensionFamily) -> PeerMetadata { } ExtensionFamily::SupervisorMiddleware => { supported_capabilities.push(SUPERVISOR_MIDDLEWARE_HTTP_HOOKS.to_string()); + supported_capabilities.push(SUPERVISOR_MIDDLEWARE_TCP_HOOKS.to_string()); } ExtensionFamily::Credentials | ExtensionFamily::GatewayInterceptor => {} } @@ -211,6 +220,28 @@ pub fn http_hooks_middleware_metadata( metadata } +/// Extension metadata for supervisor middleware that implements TCP +/// connection hooks. +/// +/// It supports and requires [`SUPERVISOR_MIDDLEWARE_TCP_HOOKS`], so peers +/// without TCP hooks reject the service at `Describe`. +#[must_use] +pub fn tcp_hooks_middleware_metadata( + implementation_name: impl Into, + implementation_version: impl Into, +) -> PeerMetadata { + let mut metadata = extension_metadata( + ExtensionFamily::SupervisorMiddleware, + implementation_name, + implementation_version, + [SUPERVISOR_MIDDLEWARE_TCP_HOOKS.to_string()], + ); + metadata + .required_capabilities + .push(SUPERVISOR_MIDDLEWARE_TCP_HOOKS.to_string()); + metadata +} + pub fn negotiate( family: ExtensionFamily, configured_name: impl Into, diff --git a/crates/openshell-core/src/middleware.rs b/crates/openshell-core/src/middleware.rs index f6ba747bfc..566a763699 100644 --- a/crates/openshell-core/src/middleware.rs +++ b/crates/openshell-core/src/middleware.rs @@ -14,10 +14,15 @@ use crate::proto::{ HttpHeader, HttpRequestEvaluation, HttpRequestHookEvent, HttpRequestHookResult, HttpRequestResult, HttpRequestTarget, HttpResponseEvent, HttpResponseEventResult, HttpResponseHookEvent, HttpResponseHookResult, MiddlewareDescribeRequest, MiddlewareManifest, - RequestContext, SupervisorMiddlewarePhase, ValidateConfigRequest, ValidateConfigResponse, - WebSocketSessionEvent, WebSocketSessionEventResult, + RequestContext, SupervisorMiddlewarePhase, TcpHookEvent, TcpHookResult, ValidateConfigRequest, + ValidateConfigResponse, WebSocketSessionEvent, WebSocketSessionEventResult, }; +/// Transport-neutral result stream for one `TcpHooks.EvaluateConnection` +/// exchange. +pub type TcpHookResultStream = + Pin> + Send + 'static>>; + /// Transport-neutral result stream for one `HttpHooks.EvaluateRequest` /// exchange. pub type HttpRequestHookResultStream = Pin< @@ -100,6 +105,16 @@ pub trait SupervisorMiddlewareEndpoint: Send + Sync { "middleware does not implement HTTP hook responses", )) } + + /// Open one TCP connection hook exchange (`TcpHooks.EvaluateConnection`). + async fn open_tcp_connection_hook( + &self, + _requests: mpsc::Receiver, + ) -> Result { + Err(Status::unimplemented( + "middleware does not implement TCP connection hooks", + )) + } } /// Borrowed request state exposed to one in-process middleware invocation. @@ -340,6 +355,16 @@ pub trait InProcessMiddleware: Send + Sync { "middleware does not implement HTTP hook responses", )) } + + /// Open one TCP connection hook exchange (`TcpHooks.EvaluateConnection`). + async fn open_tcp_connection_hook( + &self, + _requests: mpsc::Receiver, + ) -> std::result::Result { + Err(Status::unimplemented( + "middleware does not implement TCP connection hooks", + )) + } } /// Default timeout for one supervisor middleware RPC. diff --git a/crates/openshell-policy/src/lib.rs b/crates/openshell-policy/src/lib.rs index 056cf39460..f5c4b90502 100644 --- a/crates/openshell-policy/src/lib.rs +++ b/crates/openshell-policy/src/lib.rs @@ -3379,6 +3379,7 @@ network_policies: } #[test] + #[ignore = "spike 02b: tls: skip conflict moved to gateway validation"] fn validate_rejects_middleware_selector_matching_tls_skip_endpoint() { let mut policy = restrictive_default_policy(); add_middleware( @@ -3436,6 +3437,7 @@ network_policies: } #[test] + #[ignore = "spike 02b: tls: skip conflict moved to gateway validation"] fn validate_rejects_explicit_fail_closed_middleware_on_tls_skip_endpoint() { let mut policy = restrictive_default_policy(); let mut middleware = middleware_config("openshell/regex"); @@ -3464,6 +3466,7 @@ network_policies: } #[test] + #[ignore = "spike 02b: tls: skip conflict moved to gateway validation"] fn validate_rejects_concrete_selector_overlapping_tls_skip_wildcard() { let mut policy = restrictive_default_policy(); let mut middleware = middleware_config("openshell/regex"); diff --git a/crates/openshell-policy/src/middleware.rs b/crates/openshell-policy/src/middleware.rs index 20e9ee05b2..cf0e263a11 100644 --- a/crates/openshell-policy/src/middleware.rs +++ b/crates/openshell-policy/src/middleware.rs @@ -283,7 +283,11 @@ pub fn validate(policy: &SandboxPolicy) -> Vec { None }; - let requires_inspection = matches!(middleware.on_error.as_str(), "" | "fail_closed"); + // Spike 02b: TCP-only middleware may select tls: skip endpoints under + // fail_closed, which this static check cannot tell apart. The + // gateway's registry-aware validation (validate_tcp_hook_rules) + // enforces the conflict for services that inspect content. + let requires_inspection = false; for (key, rule) in &policy.network_policies { let policy_name = if rule.name.is_empty() { key diff --git a/crates/openshell-server/src/middleware.rs b/crates/openshell-server/src/middleware.rs index a86c64f9a7..7be904530b 100644 --- a/crates/openshell-server/src/middleware.rs +++ b/crates/openshell-server/src/middleware.rs @@ -13,7 +13,8 @@ pub async fn validate_policy( ) -> Result<(), Status> { async { registry.validate_policy_configs(policy).await?; - registry.validate_http_hook_rules(policy).await + registry.validate_http_hook_rules(policy).await?; + registry.validate_tcp_hook_rules(policy).await } .await .map_err(|error| { diff --git a/crates/openshell-supervisor-middleware/src/lib.rs b/crates/openshell-supervisor-middleware/src/lib.rs index d8edfb9cff..c294f46a0e 100644 --- a/crates/openshell-supervisor-middleware/src/lib.rs +++ b/crates/openshell-supervisor-middleware/src/lib.rs @@ -9,6 +9,7 @@ pub mod headers; mod http_hooks; mod remote; mod response; +mod tcp_hooks; mod websocket; pub use http_hooks::{ @@ -44,8 +45,8 @@ use miette::{Result, miette}; use prost::Message; use openshell_core::extension_protocol::{ - ExtensionFamily, NegotiatedExtension, SUPERVISOR_MIDDLEWARE_HTTP_HOOKS, gateway_metadata, - negotiate, + ExtensionFamily, NegotiatedExtension, SUPERVISOR_MIDDLEWARE_HTTP_HOOKS, + SUPERVISOR_MIDDLEWARE_TCP_HOOKS, gateway_metadata, negotiate, }; use openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware; use openshell_core::proto::{ @@ -62,7 +63,11 @@ use binding_policy::BindingPolicy; pub use openshell_core::middleware::{ HttpRequestHookResultStream, HttpRequestView, HttpResponseHookResultStream, HttpResponseResultStream, InProcessMiddleware, SupervisorMiddlewareEndpoint, - WebSocketResponseStream, + TcpHookResultStream, WebSocketResponseStream, +}; +pub use tcp_hooks::{ + MAX_TCP_STATS_INTERVAL, MIN_TCP_STATS_INTERVAL, MeteredStream, TcpCloseObserver, + TcpConnectionMeter, TcpHookAdmission, TcpHookDenial, TcpHookRequest, TcpHookSession, }; pub type MiddlewareService = @@ -235,6 +240,13 @@ impl InProcessMiddleware for EndpointInProcessAdapter { ) -> std::result::Result { self.endpoint.open_http_response_hook(requests).await } + + async fn open_tcp_connection_hook( + &self, + requests: tokio::sync::mpsc::Receiver, + ) -> std::result::Result { + self.endpoint.open_tcp_connection_hook(requests).await + } } /// Adapt a transport-neutral endpoint to the in-process registry contract. @@ -490,6 +502,41 @@ fn validate_manifest_http_hooks_capability( Ok(()) } +/// True when the service requires the TCP hooks capability. +fn requires_tcp_hooks(manifest: &MiddlewareManifest) -> bool { + manifest.extension.as_ref().is_some_and(|extension| { + extension + .required_capabilities + .iter() + .any(|capability| capability == SUPERVISOR_MIDDLEWARE_TCP_HOOKS) + }) +} + +fn is_tcp_binding(binding: &MiddlewareBinding) -> bool { + binding.operation == SupervisorMiddlewareOperation::TcpConnection as i32 +} + +/// A service with a TCP connection binding must require the TCP hooks +/// capability, so peers without TCP hooks reject it at Describe. +fn validate_manifest_tcp_hooks_capability( + source: &str, + manifest: &MiddlewareManifest, +) -> Result<()> { + if manifest.bindings.iter().any(is_tcp_binding) && !requires_tcp_hooks(manifest) { + return Err(miette!( + "{source} declares a TCP_CONNECTION binding without requiring {SUPERVISOR_MIDDLEWARE_TCP_HOOKS}" + )); + } + Ok(()) +} + +/// True when every binding of the service is a TCP connection binding. Such a +/// service observes connections without inspecting content, so it never +/// requires traffic to be inspectable. +fn manifest_is_tcp_only(manifest: &MiddlewareManifest) -> bool { + !manifest.bindings.is_empty() && manifest.bindings.iter().all(is_tcp_binding) +} + /// Legacy HTTP hooks and WebSocket bindings carry payloads. An HTTP hook /// binding that advertises no payload limit is offered no body modes, so it /// can only continue or deny at preflight. @@ -497,6 +544,9 @@ fn binding_requires_payload_limit( manifest: &MiddlewareManifest, binding: &MiddlewareBinding, ) -> bool { + if is_tcp_binding(binding) { + return false; + } binding_http_hook_version(manifest, binding) != Some(HttpHookVersion::Current) || binding.max_payload_bytes > 0 } @@ -869,6 +919,16 @@ impl MiddlewareDispatch { Self::Grpc(service) => service.open_http_response_hook(receiver).await, } } + + async fn open_tcp_connection_hook( + &self, + receiver: tokio::sync::mpsc::Receiver, + ) -> std::result::Result { + match self { + Self::InProcess(service) => service.open_tcp_connection_hook(receiver).await, + Self::Grpc(service) => service.open_tcp_connection_hook(receiver).await, + } + } } struct MiddlewareServiceState { @@ -1119,6 +1179,10 @@ fn supported_binding( Some(SupervisorMiddlewareOperation::WebsocketMessage), Some(SupervisorMiddlewarePhase::PreCredentials), ) => Ok(SupervisorMiddlewareBinding::WebsocketMessagePreCredentials), + ( + Some(SupervisorMiddlewareOperation::TcpConnection), + Some(SupervisorMiddlewarePhase::PreConnect), + ) => Ok(SupervisorMiddlewareBinding::TcpConnectionPreConnect), ( Some(SupervisorMiddlewareOperation::WebsocketMessage), Some(SupervisorMiddlewarePhase::PreReturn), @@ -1140,6 +1204,7 @@ fn validate_manifest_bindings( } validate_manifest_http_hooks_capability(source, manifest)?; + validate_manifest_tcp_hooks_capability(source, manifest)?; let mut bindings = BTreeMap::new(); for binding in &manifest.bindings { let selector = supported_binding(source, binding)?; @@ -1541,6 +1606,78 @@ impl MiddlewareRegistry { }) } + /// True unless this registry described the middleware and every one of its + /// bindings is a TCP connection binding. Only such a service may select + /// traffic `OpenShell` cannot inspect (`tls: skip`, raw TCP) under + /// `fail_closed`: it observes connections without inspecting content. + #[must_use] + pub fn requires_inspection(&self, implementation: &str) -> bool { + !self.services.iter().any(|state| { + state.manifest.get().is_some_and(|manifest| { + ChainRunner::attachment_name(state, manifest) == implementation + && manifest_is_tcp_only(manifest) + }) + }) + } + + /// Gateway rules for policies that use TCP hook middleware. + /// + /// - `on_error: fail_open` is rejected on an entry whose service has a TCP + /// connection binding: TCP hooks always fail closed. + /// - A `fail_closed` entry whose service inspects content (any HTTP or + /// WebSocket binding, or an undescribed service) must not select a + /// `tls: skip` endpoint. TCP-only services may. + pub async fn validate_tcp_hook_rules(&self, policy: &SandboxPolicy) -> Result<()> { + use openshell_core::host_pattern::{HostPattern, HostSelector}; + let manifests = ChainRunner::from_registry(self.clone()).manifests().await?; + let manifest_for = |implementation: &str| { + manifests + .iter() + .find(|(state, manifest)| { + ChainRunner::attachment_name(state, manifest) == implementation + }) + .map(|(_, manifest)| manifest) + }; + let mut configs: Vec<_> = policy.network_middlewares.iter().collect(); + configs.sort_unstable_by_key(|(name, _)| name.as_str()); + for (name, config) in configs { + let manifest = manifest_for(&config.middleware); + let on_error = OnError::parse(&config.on_error)?; + if on_error == OnError::FailOpen + && manifest.is_some_and(|manifest| manifest.bindings.iter().any(is_tcp_binding)) + { + return Err(miette!( + "middleware config '{name}' cannot use on_error: fail_open with '{}': TCP hook middleware is always fail-closed; use fail_closed", + config.middleware + )); + } + if on_error == OnError::FailOpen || manifest.is_some_and(manifest_is_tcp_only) { + continue; + } + let Some(selector) = config + .endpoints + .as_ref() + .and_then(|selector| HostSelector::new(&selector.include, &selector.exclude).ok()) + else { + continue; + }; + for (key, rule) in &policy.network_policies { + for endpoint in &rule.endpoints { + if endpoint.tls == openshell_core::proto::NetworkTlsMode::Skip as i32 + && HostPattern::new(&endpoint.host) + .is_ok_and(|pattern| selector.may_match_pattern(&pattern)) + { + return Err(miette!( + "fail-closed middleware config '{name}' selects tls: skip endpoint '{}' in network policy '{key}'; its service inspects content it cannot see there. Exclude the host, or use a TCP-only middleware service", + endpoint.host + )); + } + } + } + } + Ok(()) + } + /// Gateway rules for policies that use HTTP hook middleware. /// /// - `on_error: fail_open` is rejected on an entry whose service runs HTTP @@ -1984,6 +2121,12 @@ impl ChainRunner { }) } + /// See [`MiddlewareRegistry::requires_inspection`]. + #[must_use] + pub fn requires_inspection(&self, implementation: &str) -> bool { + self.registry.requires_inspection(implementation) + } + /// HTTP hook version of a registered middleware, without a network call. See /// [`MiddlewareRegistry::http_hook_version_of`]. #[must_use] diff --git a/crates/openshell-supervisor-middleware/src/remote.rs b/crates/openshell-supervisor-middleware/src/remote.rs index ad354a844b..f06660a294 100644 --- a/crates/openshell-supervisor-middleware/src/remote.rs +++ b/crates/openshell-supervisor-middleware/src/remote.rs @@ -4,15 +4,17 @@ use miette::{IntoDiagnostic, Result, WrapErr}; use openshell_core::middleware::{ HttpRequestHookResultStream, HttpRequestView, HttpResponseHookResultStream, - HttpResponseResultStream, SupervisorMiddlewareEndpoint, WebSocketResponseStream, + HttpResponseResultStream, SupervisorMiddlewareEndpoint, TcpHookResultStream, + WebSocketResponseStream, }; use openshell_core::proto::middleware::v1::http_hooks_client::HttpHooksClient; use openshell_core::proto::middleware::v1::http_response_pre_return_client::HttpResponsePreReturnClient; use openshell_core::proto::middleware::v1::supervisor_middleware_client::SupervisorMiddlewareClient; +use openshell_core::proto::middleware::v1::tcp_hooks_client::TcpHooksClient; use openshell_core::proto::{ HttpRequestEvaluation, HttpRequestHookEvent, HttpRequestResult, HttpResponseEvent, - HttpResponseHookEvent, MiddlewareDescribeRequest, MiddlewareManifest, ValidateConfigRequest, - ValidateConfigResponse, WebSocketSessionEvent, + HttpResponseHookEvent, MiddlewareDescribeRequest, MiddlewareManifest, TcpHookEvent, + ValidateConfigRequest, ValidateConfigResponse, WebSocketSessionEvent, }; use openshell_extension_core::{ BearerTokenInterceptor, BearerTokenSlot, ExtensionChannelConfig, ExtensionServerTrust, @@ -137,6 +139,14 @@ impl GrpcMiddlewareService { ) -> std::result::Result { self.service.open_http_response_hook(receiver).await } + + /// Open a remote TCP connection hook exchange through the gRPC adapter. + pub async fn open_tcp_connection_hook( + &self, + receiver: tokio::sync::mpsc::Receiver, + ) -> std::result::Result { + self.service.open_tcp_connection_hook(receiver).await + } } #[derive(Clone)] @@ -144,6 +154,7 @@ pub struct RemoteMiddlewareService { client: SupervisorMiddlewareClient, response_client: HttpResponsePreReturnClient, http_hooks_client: HttpHooksClient, + tcp_hooks_client: TcpHooksClient, } impl RemoteMiddlewareService { @@ -177,7 +188,10 @@ impl RemoteMiddlewareService { response_client: HttpResponsePreReturnClient::new(channel.clone()) .max_decoding_message_size(MIDDLEWARE_GRPC_MESSAGE_BYTES) .max_encoding_message_size(MIDDLEWARE_GRPC_MESSAGE_BYTES), - http_hooks_client: HttpHooksClient::new(channel) + http_hooks_client: HttpHooksClient::new(channel.clone()) + .max_decoding_message_size(MIDDLEWARE_GRPC_MESSAGE_BYTES) + .max_encoding_message_size(MIDDLEWARE_GRPC_MESSAGE_BYTES), + tcp_hooks_client: TcpHooksClient::new(channel) .max_decoding_message_size(MIDDLEWARE_GRPC_MESSAGE_BYTES) .max_encoding_message_size(MIDDLEWARE_GRPC_MESSAGE_BYTES), }) @@ -269,6 +283,20 @@ impl SupervisorMiddlewareEndpoint for RemoteMiddlewareService { .into_inner(); Ok(Box::pin(responses)) } + + async fn open_tcp_connection_hook( + &self, + receiver: tokio::sync::mpsc::Receiver, + ) -> std::result::Result { + let mut client = self.tcp_hooks_client.clone(); + let responses = client + .evaluate_connection(Request::new(tokio_stream::wrappers::ReceiverStream::new( + receiver, + ))) + .await? + .into_inner(); + Ok(Box::pin(responses)) + } } #[cfg(test)] diff --git a/crates/openshell-supervisor-middleware/src/tcp_hooks.rs b/crates/openshell-supervisor-middleware/src/tcp_hooks.rs new file mode 100644 index 0000000000..59a5cd3654 --- /dev/null +++ b/crates/openshell-supervisor-middleware/src/tcp_hooks.rs @@ -0,0 +1,584 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! TCP connection hooks (`TcpHooks.EvaluateConnection`). +//! +//! One stream per admitted upstream connection and stage. The preflight runs +//! before the upstream is dialed and can deny the connection. A stage that +//! selects STATS receives cumulative upstream byte totals at its interval. Any +//! stage can close the connection at any time. TCP hooks always fail closed. + +use std::io; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::task::{Context, Poll}; +use std::time::{Duration, Instant}; + +use futures::StreamExt; +use futures::task::AtomicWaker; +use miette::Result; +use openshell_core::proto::{ + MiddlewareSessionEndReason, RequestContext, SupervisorMiddlewareOperation, + SupervisorMiddlewarePhase, TcpEnd, TcpHookEvent, TcpHookResult, TcpInspection, TcpLimits, + TcpMode, TcpPreflight, TcpStats, tcp_hook_event, tcp_hook_result, tcp_inspect, + tcp_preflight_result, +}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot}; + +use crate::{ChainEntry, ChainRunner, DescribedChainEntry, MAX_MIDDLEWARE_REASON_CODE_BYTES}; + +/// Shortest stats interval a stage may select. +pub const MIN_TCP_STATS_INTERVAL: Duration = Duration::from_secs(1); +/// Longest stats interval a stage may select. +pub const MAX_TCP_STATS_INTERVAL: Duration = Duration::from_mins(1); +/// Process-wide cap on concurrently observed connections. Spike: one shared +/// cap, not a budget per service. +pub const MAX_TCP_HOOK_CONNECTIONS: usize = 4096; + +static TCP_HOOK_CONNECTIONS: std::sync::LazyLock> = + std::sync::LazyLock::new(|| Arc::new(Semaphore::new(MAX_TCP_HOOK_CONNECTIONS))); + +/// Upstream byte counters and the close switch for one connection. +#[derive(Debug)] +pub struct TcpConnectionMeter { + bytes_out: AtomicU64, + bytes_in: AtomicU64, + closed: AtomicBool, + read_waker: AtomicWaker, + write_waker: AtomicWaker, + started: Instant, +} + +impl Default for TcpConnectionMeter { + fn default() -> Self { + Self { + bytes_out: AtomicU64::new(0), + bytes_in: AtomicU64::new(0), + closed: AtomicBool::new(false), + read_waker: AtomicWaker::new(), + write_waker: AtomicWaker::new(), + started: Instant::now(), + } + } +} + +impl TcpConnectionMeter { + /// Bytes written to and read from the upstream so far. + #[must_use] + pub fn totals(&self) -> (u64, u64) { + ( + self.bytes_out.load(Ordering::Relaxed), + self.bytes_in.load(Ordering::Relaxed), + ) + } + + #[must_use] + pub fn elapsed(&self) -> Duration { + self.started.elapsed() + } + + #[must_use] + pub fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) + } + + /// Abort the connection: pending and later upstream reads and writes fail. + pub fn close(&self) { + self.closed.store(true, Ordering::Release); + self.read_waker.wake(); + self.write_waker.wake(); + } +} + +fn aborted() -> io::Error { + io::Error::new( + io::ErrorKind::ConnectionAborted, + "connection closed by TCP middleware", + ) +} + +/// Upstream stream wrapper that counts bytes and honors the meter's close. +pub struct MeteredStream { + inner: S, + meter: Arc, +} + +impl MeteredStream { + pub fn new(inner: S, meter: Arc) -> Self { + Self { inner, meter } + } +} + +impl AsyncRead for MeteredStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + this.meter.read_waker.register(cx.waker()); + if this.meter.is_closed() { + return Poll::Ready(Err(aborted())); + } + let before = buf.filled().len(); + let result = Pin::new(&mut this.inner).poll_read(cx, buf); + if matches!(&result, Poll::Ready(Ok(()))) { + let read = buf.filled().len().saturating_sub(before); + this.meter + .bytes_in + .fetch_add(read as u64, Ordering::Relaxed); + } + result + } +} + +impl AsyncWrite for MeteredStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let this = self.get_mut(); + this.meter.write_waker.register(cx.waker()); + if this.meter.is_closed() { + return Poll::Ready(Err(aborted())); + } + let result = Pin::new(&mut this.inner).poll_write(cx, buf); + if let Poll::Ready(Ok(written)) = &result { + this.meter + .bytes_out + .fetch_add(*written as u64, Ordering::Relaxed); + } + result + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.meter.is_closed() { + return Poll::Ready(Err(aborted())); + } + Pin::new(&mut this.inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } +} + +/// Called once when a stage closes an open connection: middleware name, +/// reason code, bytes out, bytes in. +pub type TcpCloseObserver = Arc; + +/// One admitted upstream connection to evaluate. +pub struct TcpHookRequest { + pub host: String, + pub port: u16, + pub inspection: TcpInspection, + pub policy_keys: Vec, + pub context: RequestContext, + pub on_close: Option, +} + +/// A stage denied the connection at preflight, or failed. +#[derive(Debug, Clone)] +pub struct TcpHookDenial { + pub middleware: String, + /// Stable reason code from the stage, or an `OpenShell` failure code. + pub reason_code: String, +} + +pub enum TcpHookAdmission { + /// Relay through [`TcpHookSession::wrap`]. + Allow(TcpHookSession), + Deny(TcpHookDenial), +} + +struct StageHandle { + end: Option>, +} + +/// Observed connection. Dropping it ends every stage stream with +/// `NORMAL`; [`TcpHookSession::finish`] supplies another reason. +pub struct TcpHookSession { + meter: Arc, + stages: Vec, + _permit: Option, +} + +impl TcpHookSession { + #[must_use] + pub fn meter(&self) -> &Arc { + &self.meter + } + + /// True when no stage observes the connection. + #[must_use] + pub fn is_empty(&self) -> bool { + self.stages.is_empty() + } + + /// Wrap the upstream stream so its bytes are counted and a stage can + /// close it. + pub fn wrap(&self, upstream: S) -> MeteredStream { + MeteredStream::new(upstream, Arc::clone(&self.meter)) + } + + /// End every stage stream with `reason`. + pub fn finish(mut self, reason: MiddlewareSessionEndReason) { + self.end_all(reason); + } + + fn end_all(&mut self, reason: MiddlewareSessionEndReason) { + for stage in &mut self.stages { + if let Some(end) = stage.end.take() { + let _ = end.send(reason); + } + } + } +} + +impl Drop for TcpHookSession { + fn drop(&mut self) { + let reason = if self.meter.is_closed() { + MiddlewareSessionEndReason::MiddlewareDenial + } else { + MiddlewareSessionEndReason::Normal + }; + self.end_all(reason); + } +} + +fn duration_proto(duration: Duration) -> prost_types::Duration { + prost_types::Duration::try_from(duration).unwrap_or_default() +} + +/// Accept the stage's code when it is a stable reason code, else a generic one. +fn stage_reason_code(code: &str, fallback: &str) -> String { + if !code.is_empty() + && code.len() <= MAX_MIDDLEWARE_REASON_CODE_BYTES + && crate::is_stable_reason_code(code) + { + code.to_string() + } else { + fallback.to_string() + } +} + +enum StagePreflight { + Continue, + Stats(Duration), + Close(String), +} + +impl ChainRunner { + /// Run the TCP connection preflight of every matching stage, in chain + /// order. The first deny or failure denies the connection. + pub async fn open_tcp_hooks( + &self, + entries: &[ChainEntry], + request: TcpHookRequest, + ) -> Result { + if entries.is_empty() { + return Ok(TcpHookAdmission::Allow(TcpHookSession { + meter: Arc::new(TcpConnectionMeter::default()), + stages: Vec::new(), + _permit: None, + })); + } + let described = self + .describe_chain_for( + entries, + SupervisorMiddlewareOperation::TcpConnection, + SupervisorMiddlewarePhase::PreConnect, + ) + .await? + .entries; + let meter = Arc::new(TcpConnectionMeter::default()); + if described.is_empty() { + return Ok(TcpHookAdmission::Allow(TcpHookSession { + meter, + stages: Vec::new(), + _permit: None, + })); + } + let Ok(permit) = Arc::clone(&TCP_HOOK_CONNECTIONS).try_acquire_owned() else { + return Ok(TcpHookAdmission::Deny(TcpHookDenial { + middleware: described[0].entry.name.clone(), + reason_code: "middleware_capacity_exhausted".to_string(), + })); + }; + let mut session = TcpHookSession { + meter, + stages: Vec::with_capacity(described.len()), + _permit: Some(permit), + }; + let request = Arc::new(request); + for entry in &described { + match open_stage(entry, &request, &session.meter).await { + Ok(stage) => session.stages.push(stage), + Err(reason_code) => { + session.end_all(MiddlewareSessionEndReason::MiddlewareDenial); + return Ok(TcpHookAdmission::Deny(TcpHookDenial { + middleware: entry.entry.name.clone(), + reason_code, + })); + } + } + } + Ok(TcpHookAdmission::Allow(session)) + } +} + +async fn open_stage( + described: &DescribedChainEntry, + request: &Arc, + meter: &Arc, +) -> std::result::Result { + let Some(state) = described.service.as_ref() else { + return Err("middleware_unavailable".to_string()); + }; + let (tx, rx) = mpsc::channel::(4); + let preflight = TcpPreflight { + host: request.host.clone(), + port: u32::from(request.port), + sni: String::new(), + inspection: request.inspection as i32, + policy_keys: request.policy_keys.clone(), + endpoint_id: format!("{}:{}", request.host, request.port), + context: Some(request.context.clone()), + middleware_name: described.entry.name.clone(), + config: Some(described.entry.config.clone()), + permitted_modes: vec![TcpMode::Stats as i32], + limits: Some(TcpLimits { + min_stats_interval: Some(duration_proto(MIN_TCP_STATS_INTERVAL)), + max_stats_interval: Some(duration_proto(MAX_TCP_STATS_INTERVAL)), + }), + phase: SupervisorMiddlewarePhase::PreConnect as i32, + }; + tx.send(TcpHookEvent { + event: Some(tcp_hook_event::Event::Preflight(preflight)), + }) + .await + .map_err(|_| "middleware_failure".to_string())?; + + let timeout = described.timeout; + let mut results = tokio::time::timeout(timeout, state.service.open_tcp_connection_hook(rx)) + .await + .map_err(|_| "middleware_timeout".to_string())? + .map_err(|status| { + tracing::warn!( + middleware = %described.entry.name, + error = %status, + "TCP middleware stream failed to open" + ); + "middleware_failure".to_string() + })?; + let first = tokio::time::timeout(timeout, results.next()) + .await + .map_err(|_| "middleware_timeout".to_string())?; + let decision = match first { + Some(Ok(TcpHookResult { + result: Some(tcp_hook_result::Result::PreflightResult(result)), + })) => match result.decision { + Some(tcp_preflight_result::Decision::ContinueWithoutInspection(_)) => { + StagePreflight::Continue + } + Some(tcp_preflight_result::Decision::Inspect(inspect)) => match inspect.mode { + Some(tcp_inspect::Mode::Stats(mode)) => { + let interval = mode + .interval + .as_ref() + .and_then(|value| openshell_core::time::duration_to_std(value).ok()) + .unwrap_or(MIN_TCP_STATS_INTERVAL) + .clamp(MIN_TCP_STATS_INTERVAL, MAX_TCP_STATS_INTERVAL); + StagePreflight::Stats(interval) + } + None => return Err("middleware_protocol_error".to_string()), + }, + None => return Err("middleware_protocol_error".to_string()), + }, + Some(Ok(TcpHookResult { + result: Some(tcp_hook_result::Result::Close(close)), + })) => StagePreflight::Close(stage_reason_code( + &close.diagnostics.unwrap_or_default().reason_code, + "middleware_denied", + )), + Some(Ok(_)) => return Err("middleware_protocol_error".to_string()), + Some(Err(status)) => { + tracing::warn!( + middleware = %described.entry.name, + error = %status, + "TCP middleware preflight failed" + ); + return Err("middleware_failure".to_string()); + } + None => return Err("middleware_failure".to_string()), + }; + let interval = match decision { + StagePreflight::Close(reason_code) => { + let _ = tx.try_send(end_event( + meter, + MiddlewareSessionEndReason::MiddlewareDenial, + )); + return Err(reason_code); + } + StagePreflight::Continue => None, + StagePreflight::Stats(interval) => Some(interval), + }; + let (end_tx, end_rx) = oneshot::channel(); + tokio::spawn(run_stage( + described.entry.name.clone(), + Arc::clone(request), + Arc::clone(meter), + tx, + results, + interval, + end_rx, + )); + Ok(StageHandle { end: Some(end_tx) }) +} + +fn end_event(meter: &TcpConnectionMeter, reason: MiddlewareSessionEndReason) -> TcpHookEvent { + let (bytes_out, bytes_in) = meter.totals(); + TcpHookEvent { + event: Some(tcp_hook_event::Event::End(TcpEnd { + bytes_out, + bytes_in, + duration: Some(duration_proto(meter.elapsed())), + reason: reason as i32, + })), + } +} + +/// Drive one stage after preflight: send stats when totals moved, apply a +/// close, and send the end. A broken stream closes the connection. +async fn run_stage( + middleware: String, + request: Arc, + meter: Arc, + tx: mpsc::Sender, + mut results: openshell_core::middleware::TcpHookResultStream, + interval: Option, + mut end_rx: oneshot::Receiver, +) { + let mut ticker = interval.map(|interval| { + let mut ticker = tokio::time::interval(interval); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + ticker + }); + let mut last = (0_u64, 0_u64); + let mut results_open = true; + let reason = loop { + tokio::select! { + reason = &mut end_rx => { + break reason.unwrap_or(MiddlewareSessionEndReason::Cancellation); + } + () = async { + match ticker.as_mut() { + Some(ticker) => { ticker.tick().await; } + None => std::future::pending::<()>().await, + } + } => { + let totals = meter.totals(); + if totals != last { + last = totals; + // Latest value wins: a full channel skips this update and + // the next tick sends newer cumulative totals. + let _ = tx.try_send(TcpHookEvent { + event: Some(tcp_hook_event::Event::Stats(TcpStats { + bytes_out: totals.0, + bytes_in: totals.1, + elapsed: Some(duration_proto(meter.elapsed())), + })), + }); + } + } + result = results.next(), if results_open => { + match result { + Some(Ok(TcpHookResult { result: Some(tcp_hook_result::Result::Close(close)) })) => { + let reason_code = stage_reason_code( + &close.diagnostics.unwrap_or_default().reason_code, + "middleware_closed", + ); + close_connection(&middleware, &request, &meter, &reason_code); + } + Some(Ok(_)) => { + close_connection(&middleware, &request, &meter, "middleware_protocol_error"); + } + Some(Err(status)) => { + tracing::warn!(middleware = %middleware, error = %status, "TCP middleware stream failed"); + close_connection(&middleware, &request, &meter, "middleware_failure"); + results_open = false; + } + None => { + // The stage ended its response stream. Fail closed. + close_connection(&middleware, &request, &meter, "middleware_failure"); + results_open = false; + } + } + } + } + }; + let _ = tx.send(end_event(&meter, reason)).await; + // Dropping the response stream cancels the RPC, which can discard the + // queued end before it is flushed. Half-close the request side and let + // the stage finish its responses, so short connections that end before + // their first stats tick still report their totals. + drop(tx); + if results_open { + let _ = tokio::time::timeout(TCP_END_DRAIN_TIMEOUT, async { + while let Some(Ok(_)) = results.next().await {} + }) + .await; + } +} + +/// How long a stage may take to finish its response stream after the end. +const TCP_END_DRAIN_TIMEOUT: Duration = Duration::from_secs(5); + +fn close_connection( + middleware: &str, + request: &TcpHookRequest, + meter: &TcpConnectionMeter, + reason_code: &str, +) { + if meter.is_closed() { + return; + } + meter.close(); + let (bytes_out, bytes_in) = meter.totals(); + if let Some(on_close) = request.on_close.as_ref() { + on_close(middleware, reason_code, bytes_out, bytes_in); + } +} + +#[cfg(test)] +mod tests { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + use super::*; + + #[tokio::test] + async fn metered_stream_counts_bytes_and_close_aborts_pending_read() { + let (client, mut server) = tokio::io::duplex(1024); + let meter = Arc::new(TcpConnectionMeter::default()); + let mut upstream = MeteredStream::new(client, Arc::clone(&meter)); + + upstream.write_all(b"hello").await.unwrap(); + server.write_all(b"world!").await.unwrap(); + let mut buf = [0_u8; 6]; + upstream.read_exact(&mut buf).await.unwrap(); + assert_eq!(meter.totals(), (5, 6)); + + let pending = tokio::spawn(async move { + let mut buf = [0_u8; 1]; + upstream.read(&mut buf).await + }); + tokio::task::yield_now().await; + meter.close(); + let error = pending.await.unwrap().unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::ConnectionAborted); + } +} diff --git a/crates/openshell-supervisor-network/src/l7/middleware.rs b/crates/openshell-supervisor-network/src/l7/middleware.rs index d44120b9dd..03db3a7ba5 100644 --- a/crates/openshell-supervisor-network/src/l7/middleware.rs +++ b/crates/openshell-supervisor-network/src/l7/middleware.rs @@ -337,6 +337,19 @@ pub fn uninspectable_traffic_gate( } } +/// Entries whose service inspects content. TCP-only services observe +/// connections without inspecting them, so they never make traffic +/// "required to be inspectable". +pub fn inspecting_entries( + runner: &openshell_supervisor_middleware::ChainRunner, + chain: Vec, +) -> Vec { + chain + .into_iter() + .filter(|entry| runner.requires_inspection(&entry.implementation)) + .collect() +} + /// Sandbox and originating process identity shown to middleware. The process /// is the one the connection's network policy resolved, without its pid. pub fn middleware_request_context( diff --git a/crates/openshell-supervisor-network/src/l7/relay.rs b/crates/openshell-supervisor-network/src/l7/relay.rs index 9be6870910..23a27f118d 100644 --- a/crates/openshell-supervisor-network/src/l7/relay.rs +++ b/crates/openshell-supervisor-network/src/l7/relay.rs @@ -1255,7 +1255,10 @@ where // The SQL relay is not implemented, so a matching middleware // chain can never inspect this stream: gate it like any other // uninspectable protocol. - let chain = engine.query_middleware_chain(&middleware_network_input(ctx))?; + let chain = crate::l7::middleware::inspecting_entries( + engine.middleware_runner(), + engine.query_middleware_chain(&middleware_network_input(ctx))?, + ); match uninspectable_traffic_gate(&chain) { UninspectableTrafficGate::Deny => { emit_middleware_uninspectable(ctx, "sql passthrough", true); diff --git a/crates/openshell-supervisor-network/src/opa.rs b/crates/openshell-supervisor-network/src/opa.rs index 9362f3ff9f..102b9fcaa1 100644 --- a/crates/openshell-supervisor-network/src/opa.rs +++ b/crates/openshell-supervisor-network/src/opa.rs @@ -12420,44 +12420,8 @@ network_middlewares: "#, "invalid middleware configuration", ), - ( - "tls skip selector", - r#" -network_middlewares: - redactor: - middleware: openshell/regex - endpoints: - include: ["api.example.com"] -network_policies: - api: - endpoints: - - host: api.example.com - port: 443 - tls: skip - binaries: - - { path: /usr/bin/curl } -"#, - "middleware conflicts with TLS inspection", - ), - ( - "tls skip wildcard overlap", - r#" -network_middlewares: - redactor: - middleware: openshell/regex - endpoints: - include: ["api.example.com"] -network_policies: - api: - endpoints: - - host: "*.example.com" - port: 443 - tls: skip - binaries: - - { path: /usr/bin/curl } -"#, - "middleware conflicts with TLS inspection", - ), + // Spike 02b: the tls: skip conflict moved to the gateway's + // registry-aware validation (validate_tcp_hook_rules). ]; for (name, data, expected) in cases { diff --git a/crates/openshell-supervisor-network/src/proxy.rs b/crates/openshell-supervisor-network/src/proxy.rs index 7048214be5..d743808823 100644 --- a/crates/openshell-supervisor-network/src/proxy.rs +++ b/crates/openshell-supervisor-network/src/proxy.rs @@ -6,6 +6,7 @@ pub(crate) mod destination; mod egress; mod relay; +mod tcp_hook; #[cfg(test)] pub(crate) use relay::relay_inspected_http_stream_for_test; @@ -1361,10 +1362,24 @@ async fn handle_transparent_tcp_connection( if middleware_gate == crate::l7::middleware::UninspectableTrafficGate::BypassWithFinding { crate::l7::middleware::emit_middleware_uninspectable(&ctx, "transparent tcp", false); } + let tcp_hook_session = match tcp_hook::admit( + &opa_engine, + &ctx, + openshell_core::proto::TcpInspection::TransparentTcp, + ) + .await + { + openshell_supervisor_middleware::TcpHookAdmission::Allow(session) => session, + openshell_supervisor_middleware::TcpHookAdmission::Deny(denial) => { + tcp_hook::emit_deny(&ctx, &denial); + emit_activity(&activity_tx, true, "tcp_middleware_denied"); + return Ok(()); + } + }; let approved_real_ip_candidates = connector.addrs().to_vec(); generation_guard.ensure_current()?; let dial_span = egress::dial_span(&connect_span); - let mut upstream = + let upstream = dial_transparent_upstream(&upstream_proxy, &host, port, &approved_real_ip_candidates) .instrument(dial_span.clone()) .await @@ -1433,6 +1448,7 @@ async fn handle_transparent_tcp_connection( } )); emit_activity(&activity_tx, false, "transparent_tcp"); + let mut upstream = tcp_hook_session.wrap(upstream); relay::relay_tcp(&mut client, &mut upstream, &generation_guard, &ctx).await } @@ -1830,6 +1846,7 @@ fn middleware_uninspectable_gate( ) -> Result { let input = crate::l7::middleware::middleware_network_input(ctx); let (chain, _generation) = opa_engine.query_middleware_chain_with_generation(&input)?; + let chain = crate::l7::middleware::inspecting_entries(&opa_engine.middleware_runner()?, chain); Ok(crate::l7::middleware::uninspectable_traffic_gate(&chain)) } @@ -3117,6 +3134,31 @@ async fn handle_mediated_connection( } } + let tcp_hook_inspection = if effective_tls_skip { + openshell_core::proto::TcpInspection::TlsPassthrough + } else { + openshell_core::proto::TcpInspection::ConnectTunnel + }; + let tcp_hook_session = match tcp_hook::admit(&opa_engine, &ctx, tcp_hook_inspection).await { + openshell_supervisor_middleware::TcpHookAdmission::Allow(session) => session, + openshell_supervisor_middleware::TcpHookAdmission::Deny(denial) => { + tcp_hook::emit_deny(&ctx, &denial); + emit_activity_simple(activity_tx.as_ref(), true, "tcp_middleware_denied"); + respond( + &mut client, + &build_json_error_response_with_reason( + 403, + "Forbidden", + "middleware_denied", + "connection denied by TCP middleware", + &denial.reason_code, + ), + ) + .await?; + return Ok(()); + } + }; + let dial_span = egress::dial_span(&connect_span); let upstream_result = tokio::select! { result = dial_upstream(&upstream_proxy, &host_lc, &raw_host_lc, port, connector.addrs()) @@ -3140,7 +3182,7 @@ async fn handle_mediated_connection( return Ok(()); }; let mut upstream = match upstream_result { - Ok(upstream) => upstream, + Ok(upstream) => tcp_hook_session.wrap(upstream), Err(error) => { egress::mark_error(&dial_span); if let Some(observer) = connect_endpoint_observer.as_ref() { @@ -6742,10 +6784,39 @@ async fn handle_forward_proxy( // would need absolute-form requests rather than a CONNECT tunnel. Dial // only after every local authorization and transformation step so a // rejected WebSocket preflight cannot contact the destination. + let tcp_hook_session = match tcp_hook::admit( + &opa_engine, + &l7_ctx, + openshell_core::proto::TcpInspection::HttpForward, + ) + .await + { + openshell_supervisor_middleware::TcpHookAdmission::Allow(session) => session, + openshell_supervisor_middleware::TcpHookAdmission::Deny(denial) => { + tcp_hook::emit_deny(&l7_ctx, &denial); + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::MiddlewareSessionEndReason::MiddlewareDenial) + .await; + } + respond( + client, + &build_json_error_response_with_reason( + 403, + "Forbidden", + "middleware_denied", + "connection denied by TCP middleware", + &denial.reason_code, + ), + ) + .await?; + return Ok(()); + } + }; let dial_span = egress::dial_span(&connect_span); let dial_result = connector.connect().instrument(dial_span.clone()).await; let mut upstream = match dial_result { - Ok(s) => s, + Ok(s) => tcp_hook_session.wrap(s), Err(e) => { egress::mark_error(&dial_span); let event = HttpActivityBuilder::new(openshell_ocsf::ctx::ctx()) diff --git a/crates/openshell-supervisor-network/src/proxy/tcp_hook.rs b/crates/openshell-supervisor-network/src/proxy/tcp_hook.rs new file mode 100644 index 0000000000..5243ac75a3 --- /dev/null +++ b/crates/openshell-supervisor-network/src/proxy/tcp_hook.rs @@ -0,0 +1,123 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! TCP connection hooks at the proxy's upstream dial sites. + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use miette::Result; +use openshell_core::proto::TcpInspection; +use openshell_ocsf::{ + ActionId, ActivityId, DispositionId, Endpoint, NetworkActivityBuilder, SeverityId, StatusId, + ocsf_emit, +}; +use openshell_supervisor_middleware::{TcpHookAdmission, TcpHookDenial, TcpHookRequest}; + +use crate::l7::middleware::{middleware_network_input, middleware_request_context}; +use crate::l7::relay::L7EvalContext; +use crate::opa::OpaEngine; + +static NEXT_CONNECTION_ID: AtomicU64 = AtomicU64::new(1); + +/// Run the preflight of every TCP stage matching this admitted connection. +/// Any failure denies it: TCP hooks always fail closed. +pub(super) async fn admit( + opa_engine: &OpaEngine, + ctx: &L7EvalContext, + inspection: TcpInspection, +) -> TcpHookAdmission { + match try_admit(opa_engine, ctx, inspection).await { + Ok(admission) => admission, + Err(error) => { + tracing::warn!(host = %ctx.host, port = ctx.port, error = %error, "TCP middleware preflight failed"); + TcpHookAdmission::Deny(TcpHookDenial { + middleware: "-".to_string(), + reason_code: "middleware_failure".to_string(), + }) + } + } +} + +async fn try_admit( + opa_engine: &OpaEngine, + ctx: &L7EvalContext, + inspection: TcpInspection, +) -> Result { + let (chain, _generation) = + opa_engine.query_middleware_chain_with_generation(&middleware_network_input(ctx))?; + let runner = opa_engine.middleware_runner()?; + let connection_id = format!("tcp-{}", NEXT_CONNECTION_ID.fetch_add(1, Ordering::Relaxed)); + let host = ctx.host.clone(); + let port = ctx.port; + let policy = ctx.policy_name.clone(); + let request = TcpHookRequest { + host: ctx.host.clone(), + port: ctx.port, + inspection, + policy_keys: if ctx.policy_name.is_empty() { + Vec::new() + } else { + vec![ctx.policy_name.clone()] + }, + context: middleware_request_context(ctx, &connection_id), + on_close: Some(Arc::new( + move |middleware, reason_code, bytes_out, bytes_in| { + emit_close( + &host, + port, + &policy, + middleware, + reason_code, + bytes_out, + bytes_in, + ); + }, + )), + }; + runner.open_tcp_hooks(&chain, request).await +} + +/// OCSF event for a connection a TCP stage denied at preflight. +pub(super) fn emit_deny(ctx: &L7EvalContext, denial: &TcpHookDenial) { + let event = NetworkActivityBuilder::new(openshell_ocsf::ctx::ctx()) + .activity(ActivityId::Open) + .action(ActionId::Denied) + .disposition(DispositionId::Blocked) + .severity(SeverityId::Medium) + .status(StatusId::Failure) + .dst_endpoint(Endpoint::from_domain(&ctx.host, ctx.port)) + .firewall_rule(&denial.middleware, "tcp-middleware") + .message(format!( + "TCP middleware denied {}:{} [middleware:{} reason:{}]", + ctx.host, ctx.port, denial.middleware, denial.reason_code + )) + .status_detail(&denial.reason_code) + .build(); + ocsf_emit!(event); +} + +fn emit_close( + host: &str, + port: u16, + policy: &str, + middleware: &str, + reason_code: &str, + bytes_out: u64, + bytes_in: u64, +) { + let event = NetworkActivityBuilder::new(openshell_ocsf::ctx::ctx()) + .activity(ActivityId::Close) + .action(ActionId::Denied) + .disposition(DispositionId::Blocked) + .severity(SeverityId::Medium) + .status(StatusId::Failure) + .dst_endpoint(Endpoint::from_domain(host, port)) + .firewall_rule(middleware, "tcp-middleware") + .message(format!( + "TCP middleware closed {host}:{port} [middleware:{middleware} reason:{reason_code} policy:{policy} bytes_out:{bytes_out} bytes_in:{bytes_in}]" + )) + .status_detail(reason_code) + .build(); + ocsf_emit!(event); +} diff --git a/examples/supervisor-middleware-egress-usage/.gitignore b/examples/supervisor-middleware-egress-usage/.gitignore new file mode 100644 index 0000000000..62f12e5090 --- /dev/null +++ b/examples/supervisor-middleware-egress-usage/.gitignore @@ -0,0 +1,2 @@ +# Spike: no tracked lockfile, so the PR diff stays comparable. +Cargo.lock diff --git a/examples/supervisor-middleware-egress-usage/Cargo.toml b/examples/supervisor-middleware-egress-usage/Cargo.toml new file mode 100644 index 0000000000..cf505b54f4 --- /dev/null +++ b/examples/supervisor-middleware-egress-usage/Cargo.toml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[workspace] + +[package] +name = "openshell-supervisor-middleware-egress-usage" +description = "Example OpenShell TCP middleware that cuts egress bursts" +version = "0.0.0" +edition = "2024" +rust-version = "1.94" +license = "Apache-2.0" +publish = false + +[dependencies] +clap = { version = "4.5", features = ["derive"] } +openshell-core = { path = "../../crates/openshell-core", default-features = false } +prost-types = "0.14" +tokio = { version = "1.43", features = ["macros", "rt-multi-thread", "sync", "time"] } +tokio-stream = "0.1" +tonic = { version = "0.14", features = ["transport"] } + +[dev-dependencies] +openshell-policy = { path = "../../crates/openshell-policy" } + +[[bin]] +name = "supervisor-middleware-egress-usage" +path = "src/main.rs" diff --git a/examples/supervisor-middleware-egress-usage/README.md b/examples/supervisor-middleware-egress-usage/README.md new file mode 100644 index 0000000000..b78638101a --- /dev/null +++ b/examples/supervisor-middleware-egress-usage/README.md @@ -0,0 +1,78 @@ + + +# Supervisor Middleware Egress Usage (spike) + +> [!WARNING] +> Spike for TCP connection hooks (`TcpHooks.EvaluateConnection`). The contract is not settled. + +The service binds `TCP_CONNECTION` at `PRE_CONNECT` and requires the `openshell.supervisor-middleware.tcp-hooks` capability. For every admitted upstream connection its selector matches, OpenShell sends a preflight before dialing. The service selects STATS mode and receives cumulative upstream byte totals at its interval. It sums them into a sliding window per sandbox and network policy key. When a window passes `max_bytes`, it closes that key's open connections and denies new ones for `block_for`, with reason code `egress_budget_exceeded`. + +Covers raw TCP, `tls: skip`, terminated TLS, and plain HTTP the same way, because bytes are counted at the upstream socket. + +## Configure limits in the policy + +The limits come from the middleware's `config` in the sandbox policy, so changing the policy changes the budget. + +> [!NOTE] +> The `config` shape below is for demonstration only. OpenShell passes `config` to the middleware service as an opaque structure and does not define its fields: each service defines and validates its own. This example's service reads the fields below. + +```yaml +network_middlewares: + egress-usage: + middleware: egress-usage-example + config: + max_bytes: 20MB # per sandbox and policy key + window: 60s + interval: 1s # stats interval, 1s to 60s + block_for: 120s + direction: both # both, in, or out + policies: # overrides for one network policy key + other_svc: + max_bytes: 50MB + on_error: fail_closed + endpoints: + include: + - host.openshell.internal +``` + +Every field is optional; the service flags supply defaults. Sizes accept `B`, `KB`, `MB`, `GB` (decimal) or a number of bytes; durations accept `ms`, `s`, `m` or a number of seconds. The gateway calls `ValidateConfig` when a policy is set, so an invalid value is rejected by `openshell policy set`. A policy update applies to connections opened after the sandbox loads it; open connections keep the limits they started with. + +## Run the smoke suite + +```shell +./examples/supervisor-middleware-egress-usage/smoke.sh --test-suite +``` + +The script builds the gateway, CLI, and a local supervisor image, starts the host-run upstream (`upstream.py`), the service, and a Docker-driver gateway, creates a sandbox with `policy.yaml`, and checks: + +- 1 MB over HTTP and raw TCP, and on the `other_svc` key, pass. +- A 200 MB download over `tls: skip` at 10 MB/s is cut shortly after 20 MB. +- New `pkg_proxy` connections are then denied; `other_svc` still passes. +- 40 short 1 MB raw TCP connections in a second sandbox are denied after 20 MB. +- `policy set` with `max_bytes: lots` is rejected. +- 15 MB on `other_svc` passes under its 50 MB budget; after `policy set` lowers it to 10 MB, the same transfer is cut. + +Set `EGRESS_USAGE_SUPERVISOR_IMAGE` to reuse a built supervisor image and skip the builds. + +## Run manually + +```shell +cd examples/supervisor-middleware-egress-usage +cargo run -- --bind 0.0.0.0:50061 +``` + +Register it in the gateway TOML. `max_payload_bytes` is required by the config schema but unused by TCP bindings: + +```toml +[[openshell.supervisor.middleware]] +name = "egress-usage-example" +grpc_endpoint = "http://:50061" +allow_insecure_transport = true +max_payload_bytes = 65536 +timeout = "500ms" +``` + +Select it in the sandbox policy with `network_middlewares`, as in `policy.yaml`. A TCP-only service may select `tls: skip` endpoints under `fail_closed`. diff --git a/examples/supervisor-middleware-egress-usage/policy.yaml b/examples/supervisor-middleware-egress-usage/policy.yaml new file mode 100644 index 0000000000..76270e2de0 --- /dev/null +++ b/examples/supervisor-middleware-egress-usage/policy.yaml @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# pkg_proxy stands in for artifactory.internal. Docker sandboxes reach the +# host-run upstream only through host.openshell.internal, so the two policy +# keys are told apart by port. The middleware groups usage by sandbox and +# policy key, and reads its limits from config below: change the policy to +# change the budget. + +version: 1 + +network_middlewares: + egress-usage: + name: Egress usage + middleware: egress-usage-example + order: 10 + config: + # Per sandbox and policy key. + max_bytes: 20MB + window: 60s + # How often OpenShell reports byte totals (1s to 60s). Bytes over the + # budget before a cut are roughly rate x interval. + interval: 1s + block_for: 120s + direction: both + # Overrides for one network policy key. + policies: + other_svc: + max_bytes: 50MB + on_error: fail_closed + endpoints: + include: + - host.openshell.internal + +network_policies: + pkg_proxy: + name: pkg_proxy + endpoints: + # Inspected plain HTTP. + - host: host.openshell.internal + port: 47080 + protocol: rest + enforcement: enforce + access: full + # TLS passthrough: HTTP hooks never see this traffic. + - host: host.openshell.internal + port: 47443 + tls: skip + # Raw TCP. + - host: host.openshell.internal + port: 47900 + protocol: tcp + binaries: + - path: "/**" + other_svc: + name: other_svc + endpoints: + - host: host.openshell.internal + port: 47081 + protocol: rest + enforcement: enforce + access: full + binaries: + - path: "/**" diff --git a/examples/supervisor-middleware-egress-usage/smoke.sh b/examples/supervisor-middleware-egress-usage/smoke.sh new file mode 100755 index 0000000000..e362d326fc --- /dev/null +++ b/examples/supervisor-middleware-egress-usage/smoke.sh @@ -0,0 +1,412 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +EXAMPLE_DIR="$ROOT/examples/supervisor-middleware-egress-usage" +RUN_TEST_SUITE=0 +PRINT_CONFIG=0 + +usage() { + cat <&2; usage >&2; exit 2 ;; + esac +done + +detect_service_host() { + if [[ -n "${EGRESS_USAGE_SMOKE_HOST:-}" ]]; then + printf '%s\n' "$EGRESS_USAGE_SMOKE_HOST" + return + fi + ip route get 1.1.1.1 2>/dev/null | awk '{ for (i = 1; i <= NF; i++) if ($i == "src") { print $(i + 1); exit } }' +} + +SERVICE_HOST="$(detect_service_host)" +[[ -n "$SERVICE_HOST" ]] || { echo "set EGRESS_USAGE_SMOKE_HOST" >&2; exit 1; } + +SMOKE_TMP_DIR="$(mktemp -d)" +LOG_DIR="$SMOKE_TMP_DIR/logs" +JWT_DIR="$SMOKE_TMP_DIR/jwt" +GATEWAY_CONFIG="$SMOKE_TMP_DIR/gateway.toml" +SETUP_LOG="$LOG_DIR/setup.log" +GATEWAY_LOG="$LOG_DIR/gateway.log" +MIDDLEWARE_LOG="$LOG_DIR/middleware.log" +UPSTREAM_LOG="$LOG_DIR/upstream.log" +SANDBOX_LOG="$LOG_DIR/sandbox.log" +RUN_ID="egress-usage-smoke-$$-$RANDOM" +SUPERVISOR_IMAGE="${EGRESS_USAGE_SUPERVISOR_IMAGE:-localhost/openshell-egress-usage/supervisor:$RUN_ID}" +SANDBOX_NAME="eu-$$-$RANDOM" +SANDBOX_CREATED=0 +# The scenarios need curl and python3 in the sandbox. +SANDBOX_IMAGE="${EGRESS_USAGE_SANDBOX_IMAGE:-ghcr.io/nvidia/openshell-community/sandboxes/base:latest}" + +mkdir -p "$LOG_DIR" + +cleanup() { + local status=$? + trap - EXIT + if [[ "$SANDBOX_CREATED" -eq 1 && -n "${CLI+x}" ]]; then + "${CLI[@]}" sandbox delete "$SANDBOX_NAME" >>"$SETUP_LOG" 2>&1 || true + fi + for pid in "${GATEWAY_PID:-}" "${MIDDLEWARE_PID:-}" "${UPSTREAM_PID:-}"; do + if [[ -n "$pid" ]]; then + kill "$pid" 2>/dev/null || true + wait "$pid" 2>/dev/null || true + fi + done + if [[ "$status" -eq 0 ]]; then + rm -rf "$SMOKE_TMP_DIR" + else + echo "logs retained in $LOG_DIR" >&2 + fi + exit "$status" +} +trap cleanup EXIT + +choose_port_block() { + local start offset ok + for _ in {1..200}; do + start=$((20000 + RANDOM % 20000)) + ok=1 + for ((offset = 0; offset < 3; offset++)); do + if (exec 3<>"/dev/tcp/127.0.0.1/$((start + offset))") 2>/dev/null; then + ok=0 + break + fi + done + if [[ "$ok" == "1" ]]; then + printf '%s\n' "$start" + return + fi + done + echo "failed to find free local ports" >&2 + exit 1 +} + +PORT_BASE="$(choose_port_block)" +MIDDLEWARE_PORT="$PORT_BASE" +GATEWAY_PORT="$((PORT_BASE + 1))" +HEALTH_PORT="$((PORT_BASE + 2))" +GATEWAY_ENDPOINT="http://127.0.0.1:$GATEWAY_PORT" + +write_gateway_config() { + cat >"$GATEWAY_CONFIG" <&2 + tail -n 100 "$path" >&2 2>/dev/null || printf '(missing)\n' >&2 + done +} + +capture_logs() { + local container_id + container_id="$(docker ps -aq --format '{{.ID}} {{.Names}}' | awk -v name="--$SANDBOX_NAME-" 'index($2, name) && $2 ~ /-supervisor$/ { print $1; exit }')" + if [[ -n "$container_id" ]]; then + docker logs "$container_id" >"$SANDBOX_LOG" 2>&1 || true + fi +} + +fail() { + printf 'FAIL %s\n' "$1" >&2 + capture_logs + dump_logs + exit 1 +} + +run_setup_step() { + local label="$1" + shift + printf 'INFO %s\n' "$label" + printf '\n== %s ==\n' "$label" >>"$SETUP_LOG" + "$@" >>"$SETUP_LOG" 2>&1 || fail "$label" +} + +cargo_target_dir() { + cargo metadata --format-version=1 --no-deps --manifest-path "$1" | jq -er '.target_directory' +} + +generate_gateway_jwt_bundle() { + mkdir -p "$JWT_DIR" + openssl genpkey -algorithm ed25519 -out "$JWT_DIR/signing.pem" >/dev/null 2>&1 + openssl pkey -in "$JWT_DIR/signing.pem" -pubout -out "$JWT_DIR/public.pem" >/dev/null 2>&1 + printf '%s\n' "$RUN_ID" >"$JWT_DIR/kid" +} + +start_upstream() { + printf 'INFO starting package proxy stand-in on host.openshell.internal\n' + openssl req -x509 -newkey rsa:2048 -nodes -days 2 -subj "/CN=host.openshell.internal" \ + -addext "subjectAltName=DNS:host.openshell.internal" \ + -keyout "$SMOKE_TMP_DIR/key.pem" -out "$SMOKE_TMP_DIR/cert.pem" >/dev/null 2>&1 + UPSTREAM_TLS_CERT="$SMOKE_TMP_DIR/cert.pem" UPSTREAM_TLS_KEY="$SMOKE_TMP_DIR/key.pem" \ + python3 -u "$EXAMPLE_DIR/upstream.py" >"$UPSTREAM_LOG" 2>&1 & + UPSTREAM_PID=$! + for _ in {1..30}; do + kill -0 "$UPSTREAM_PID" 2>/dev/null || fail "upstream starts (ports 47080, 47081, 47443, 47900)" + if grep -q 'upstream ready' "$UPSTREAM_LOG"; then + return + fi + sleep 1 + done + fail "upstream is ready" +} + +start_middleware() { + printf 'INFO starting egress-usage service at %s:%s\n' "$SERVICE_HOST" "$MIDDLEWARE_PORT" + "$MIDDLEWARE_BIN" \ + --bind "0.0.0.0:$MIDDLEWARE_PORT" \ + >"$MIDDLEWARE_LOG" 2>&1 & + MIDDLEWARE_PID=$! + for _ in {1..60}; do + kill -0 "$MIDDLEWARE_PID" 2>/dev/null || fail "egress-usage service starts" + if (exec 3<>"/dev/tcp/$SERVICE_HOST/$MIDDLEWARE_PORT") 2>/dev/null; then + return + fi + sleep 1 + done + fail "egress-usage service is reachable" +} + +start_gateway() { + printf 'INFO starting gateway\n' + env -u OPENSHELL_DRIVERS -u OPENSHELL_COMPUTE_DRIVER "$GATEWAY_BIN" \ + --compute-driver docker \ + --config "$GATEWAY_CONFIG" \ + --bind-address 127.0.0.1 \ + --port "$GATEWAY_PORT" \ + --health-port "$HEALTH_PORT" \ + --metrics-port 0 \ + --log-level "${EGRESS_USAGE_SMOKE_LOG_LEVEL:-info}" \ + --disable-tls \ + --db-url "sqlite://$SMOKE_TMP_DIR/gateway.db" >"$GATEWAY_LOG" 2>&1 & + GATEWAY_PID=$! + for _ in {1..60}; do + kill -0 "$GATEWAY_PID" 2>/dev/null || fail "gateway starts with egress-usage" + if curl -fsS "http://127.0.0.1:$HEALTH_PORT/healthz" >/dev/null 2>&1; then + return + fi + sleep 1 + done + fail "gateway starts with egress-usage" +} + +create_sandbox() { + CLI=(env -u OPENSHELL_SANDBOX_POLICY "$CLI_BIN" --gateway-endpoint "$GATEWAY_ENDPOINT") + SANDBOX_CREATED=1 + run_setup_step "creating sandbox" \ + "${CLI[@]}" sandbox create --name "$SANDBOX_NAME" --from "$SANDBOX_IMAGE" --policy "$EXAMPLE_DIR/policy.yaml" --no-tty --detach -- sleep infinity +} + +sx() { + "${CLI[@]}" sandbox exec --name "$SANDBOX_NAME" --no-tty -- "$@" +} + +# Raw TCP client: the sandbox boundary mediates the connection transparently. +RAW_CLIENT=' +import socket, sys +host, port, count = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]) +received = 0 +try: + s = socket.create_connection((host, port), timeout=60) + s.sendall(("PULL %d\n" % count).encode()) + while received < count: + chunk = s.recv(65536) + if not chunk: + break + received += len(chunk) +except OSError as error: + print("error=%s" % error) +print("received=%d" % received) +sys.exit(0 if received == count else 4) +' + +raw_get() { + sx python3 -c "$RAW_CLIENT" "$1" "$2" "$3" +} + +http_get() { + sx curl -sS -o /dev/null -w 'status=%{http_code} size=%{size_download}\n' --max-time 60 "$@" +} + +run_suite() { + local out started elapsed + out="$(http_get http://host.openshell.internal:47080/bytes/1000000 2>&1)" || fail "baseline HTTP download: $out" + [[ "$out" == *"status=200 size=1000000"* ]] || fail "baseline HTTP download: $out" + printf 'PASS baseline HTTP 1MB under threshold (%s)\n' "$out" + + out="$(raw_get host.openshell.internal 47900 1000000 2>&1)" || fail "baseline raw TCP: $out" + printf 'PASS baseline raw TCP 1MB under threshold (%s)\n' "$out" + + out="$(http_get http://host.openshell.internal:47081/bytes/1000000 2>&1)" || fail "control host: $out" + printf 'PASS control key other_svc (%s)\n' "$out" + + printf 'INFO D2 burst: 200MB over tls: skip at 10MB/s, policy budget 20MB per 60s\n' + started=$(date +%s.%N) + out="$(http_get -k --limit-rate 10M https://host.openshell.internal:47443/bytes/200000000 2>&1)" && fail "D2 burst was not cut: $out" + elapsed=$(echo "$(date +%s.%N) - $started" | bc) + [[ "$out" != *"size=200000000"* ]] || fail "D2 burst was not cut: $out" + printf 'PASS D2 burst over tls: skip cut after %ss (%s)\n' "$elapsed" "$(echo "$out" | tr '\n' ' ')" + + out="$(raw_get host.openshell.internal 47900 1000 2>&1)" && fail "new raw TCP connection was not denied: $out" + [[ "$out" == *"received=0"* ]] || fail "new raw TCP connection denial: $out" + printf 'PASS new pkg_proxy raw TCP connection denied (%s)\n' "$out" + + out="$(http_get http://host.openshell.internal:47080/bytes/1000 2>&1)" || true + # Transparently mediated connections have no CONNECT to answer: a + # preflight denial closes the connection before any response. + [[ "$out" == *"status=403"* || "$out" == *"Empty reply"* ]] || fail "new HTTP request denied: $out" + printf 'PASS new pkg_proxy HTTP request denied (%s)\n' "$(echo "$out" | tr '\n' ' ')" + + out="$(http_get http://host.openshell.internal:47081/bytes/1000000 2>&1)" || fail "control host after trip: $out" + printf 'PASS control key other_svc still allowed (%s)\n' "$out" + + # Split transfer in a fresh sandbox: each connection ends before its first + # stats tick, so only the end totals count toward the window. + "${CLI[@]}" sandbox delete "$SANDBOX_NAME" >>"$SETUP_LOG" 2>&1 || true + SANDBOX_NAME="${SANDBOX_NAME}b" + create_sandbox + local count denied=0 + for count in $(seq 1 40); do + out="$(raw_get host.openshell.internal 47900 1000000 2>&1)" || true + if [[ "$out" == *"received=0"* ]]; then + denied="$count" + break + fi + done + [[ "$denied" -gt 0 ]] || fail "split transfer of 40 x 1MB raw TCP connections was not stopped" + printf 'PASS split transfer denied at connection %s (%s MB through)\n' "$denied" "$((denied - 1))" + + # The budget lives in the policy: changing the policy changes it. + sed 's/max_bytes: 20MB/max_bytes: lots/' "$EXAMPLE_DIR/policy.yaml" >"$SMOKE_TMP_DIR/invalid.yaml" + out="$("${CLI[@]}" policy set "$SANDBOX_NAME" --policy "$SMOKE_TMP_DIR/invalid.yaml" 2>&1)" \ + && fail "invalid budget was accepted: $out" + [[ "$out" == *"invalid byte size"* ]] || fail "invalid budget rejection: $out" + printf 'PASS invalid budget rejected by ValidateConfig (%s)\n' "$(echo "$out" | grep -m1 'invalid byte size')" + + out="$(http_get --limit-rate 10M http://host.openshell.internal:47081/bytes/15000000 2>&1)" \ + || fail "15MB on other_svc under its 50MB budget: $out" + printf 'PASS 15MB on other_svc under its 50MB policy budget (%s)\n' "$out" + + sed 's/max_bytes: 50MB/max_bytes: 10MB/' "$EXAMPLE_DIR/policy.yaml" >"$SMOKE_TMP_DIR/tight.yaml" + run_setup_step "lowering the other_svc budget to 10MB in the policy" \ + "${CLI[@]}" policy set "$SANDBOX_NAME" --policy "$SMOKE_TMP_DIR/tight.yaml" --wait + out="$(http_get --limit-rate 10M http://host.openshell.internal:47081/bytes/15000000 2>&1)" \ + && fail "15MB on other_svc was not cut after lowering its budget: $out" + printf 'PASS same 15MB on other_svc cut after the policy lowered its budget to 10MB (%s)\n' "$(echo "$out" | tr '\n' ' ')" + + grep -E 'TRIP|deny' "$MIDDLEWARE_LOG" || true + echo "ALL PASS egress-usage smoke" +} + +print_ready() { + cat < --wait + +Burst over tls: skip (cut at the policy budget, 20MB per 60s): + ${CLI[*]} sandbox exec --name $SANDBOX_NAME --no-tty -- curl -k --limit-rate 10M -o /dev/null https://host.openshell.internal:47443/bytes/200000000 + +Press Ctrl-C to clean up. +READY +} + +cd "$ROOT" +for command in cargo curl jq openssl mise docker bc python3; do + command -v "$command" >/dev/null 2>&1 || fail "missing required command: $command" +done +ROOT_TARGET_DIR="$(cargo_target_dir "$ROOT/Cargo.toml")" +EXAMPLE_TARGET_DIR="$(cargo_target_dir "$EXAMPLE_DIR/Cargo.toml")" +GATEWAY_BIN="$ROOT_TARGET_DIR/debug/openshell-gateway" +CLI_BIN="$ROOT_TARGET_DIR/debug/openshell" +MIDDLEWARE_BIN="$EXAMPLE_TARGET_DIR/debug/supervisor-middleware-egress-usage" +if [[ -z "${EGRESS_USAGE_SUPERVISOR_IMAGE:-}" ]]; then + run_setup_step "building gateway" cargo build --quiet -p openshell-gateway --bin openshell-gateway + run_setup_step "building Linux sandbox supervisor image" \ + env -u CI -u DOCKER_PLATFORM -u DOCKER_PUSH -u DOCKER_OUTPUT \ + CONTAINER_ENGINE=docker PREBUILT_AUTO_STAGE=1 \ + IMAGE_REGISTRY=localhost/openshell-egress-usage IMAGE_TAG="$RUN_ID" \ + mise run docker:build:supervisor + run_setup_step "building CLI" cargo build --quiet -p openshell-cli --bin openshell +fi +run_setup_step "building egress-usage service" cargo build --quiet --manifest-path "$EXAMPLE_DIR/Cargo.toml" +generate_gateway_jwt_bundle +start_upstream +start_middleware +start_gateway +create_sandbox + +if [[ "$RUN_TEST_SUITE" -eq 1 ]]; then + run_suite +else + print_ready + while kill -0 "$GATEWAY_PID" 2>/dev/null && kill -0 "$MIDDLEWARE_PID" 2>/dev/null; do sleep 1; done + fail "gateway or middleware exited" +fi diff --git a/examples/supervisor-middleware-egress-usage/src/config.rs b/examples/supervisor-middleware-egress-usage/src/config.rs new file mode 100644 index 0000000000..616df873c7 --- /dev/null +++ b/examples/supervisor-middleware-egress-usage/src/config.rs @@ -0,0 +1,302 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Limits from the policy's middleware config, with flag defaults. +//! +//! ```yaml +//! config: +//! max_bytes: 20MB +//! window: 60s +//! interval: 5s +//! block_for: 300s +//! direction: both +//! policies: +//! other_svc: +//! max_bytes: 5MB +//! ``` +//! +//! Every field is optional. Top-level fields override the service flags, and +//! `policies.` overrides them for one network policy key. + +use std::collections::HashMap; +use std::time::Duration; + +use prost_types::Struct; +use prost_types::value::Kind; + +use crate::usage::{Direction, Limits}; + +/// OpenShell clamps stats intervals to this range. +const MIN_INTERVAL: Duration = Duration::from_secs(1); +const MAX_INTERVAL: Duration = Duration::from_mins(1); + +#[derive(Debug, PartialEq, Eq)] +pub struct Config { + limits: Limits, + policies: HashMap, +} + +impl Config { + pub fn parse(config: Option<&Struct>, defaults: Limits) -> Result { + let Some(config) = config else { + return Ok(Self { + limits: defaults, + policies: HashMap::new(), + }); + }; + let limits = parse_limits(config, defaults, "config", true)?; + let mut policies = HashMap::new(); + if let Some(value) = config.fields.get("policies") { + let Some(Kind::StructValue(entries)) = &value.kind else { + return Err("config.policies must be a map of policy key to limits".into()); + }; + for (key, value) in &entries.fields { + let Some(Kind::StructValue(fields)) = &value.kind else { + return Err(format!("config.policies.{key} must be a map")); + }; + let path = format!("config.policies.{key}"); + policies.insert(key.clone(), parse_limits(fields, limits, &path, false)?); + } + } + Ok(Self { limits, policies }) + } + + pub fn limits_for(&self, policy_key: &str) -> Limits { + self.policies + .get(policy_key) + .copied() + .unwrap_or(self.limits) + } +} + +fn parse_limits( + fields: &Struct, + base: Limits, + path: &str, + top_level: bool, +) -> Result { + if let Some(field) = fields.fields.keys().find(|field| match field.as_str() { + "max_bytes" | "window" | "interval" | "block_for" | "direction" => false, + "policies" => !top_level, + _ => true, + }) { + return Err(format!("unsupported field '{path}.{field}'")); + } + let mut limits = base; + for (name, value) in &fields.fields { + let field = format!("{path}.{name}"); + match name.as_str() { + "max_bytes" => { + limits.max_bytes = value_with(value, &field, parse_bytes, |bytes| bytes)? + } + "window" => limits.window = duration(value, &field)?, + "interval" => limits.interval = duration(value, &field)?, + "block_for" => limits.block_for = duration(value, &field)?, + "direction" => { + let Some(Kind::StringValue(direction)) = &value.kind else { + return Err(format!("{field} must be 'both', 'in', or 'out'")); + }; + limits.direction = match direction.as_str() { + "both" => Direction::Both, + "in" => Direction::In, + "out" => Direction::Out, + _ => return Err(format!("{field} must be 'both', 'in', or 'out'")), + }; + } + _ => {} + } + } + if limits.max_bytes == 0 { + return Err(format!("{path}.max_bytes must be greater than zero")); + } + if !(MIN_INTERVAL..=MAX_INTERVAL).contains(&limits.interval) { + return Err(format!("{path}.interval must be between 1s and 60s")); + } + if limits.window < limits.interval { + return Err(format!("{path}.window must be at least the interval")); + } + Ok(limits) +} + +fn duration(value: &prost_types::Value, field: &str) -> Result { + value_with(value, field, parse_duration, Duration::from_secs) +} + +/// Accepts a string with units, or a number in base units (bytes, seconds). +fn value_with( + value: &prost_types::Value, + field: &str, + parse: fn(&str) -> Result, + from_number: impl Fn(u64) -> T, +) -> Result { + match &value.kind { + Some(Kind::StringValue(text)) => parse(text).map_err(|error| format!("{field}: {error}")), + Some(Kind::NumberValue(number)) if *number >= 0.0 && number.fract() == 0.0 => + { + #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)] + Ok(from_number(*number as u64)) + } + _ => Err(format!( + "{field} must be a string such as '20MB' or '60s', or a whole number" + )), + } +} + +pub fn parse_bytes(value: &str) -> Result { + let upper = value.trim().to_ascii_uppercase(); + let (digits, multiplier) = [ + ("GB", 1_000_000_000), + ("MB", 1_000_000), + ("KB", 1_000), + ("B", 1), + ] + .iter() + .find_map(|(suffix, multiplier)| { + upper + .strip_suffix(suffix) + .map(|digits| (digits.trim().to_string(), *multiplier)) + }) + .unwrap_or((upper.clone(), 1)); + digits + .parse::() + .map(|value| value * multiplier) + .map_err(|_| format!("invalid byte size '{value}'")) +} + +pub fn parse_duration(value: &str) -> Result { + let value = value.trim(); + let (digits, unit) = value.split_at( + value + .find(|c: char| !c.is_ascii_digit()) + .unwrap_or(value.len()), + ); + let number: u64 = digits + .parse() + .map_err(|_| format!("invalid duration '{value}'"))?; + match unit { + "" | "s" => Ok(Duration::from_secs(number)), + "ms" => Ok(Duration::from_millis(number)), + "m" => Ok(Duration::from_secs(number * 60)), + _ => Err(format!("invalid duration '{value}', use ms, s, or m")), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use prost_types::Value; + + fn defaults() -> Limits { + Limits { + max_bytes: 20_000_000, + window: Duration::from_secs(60), + interval: Duration::from_secs(5), + block_for: Duration::from_secs(300), + direction: Direction::Both, + } + } + + fn text(value: &str) -> Value { + Value { + kind: Some(Kind::StringValue(value.into())), + } + } + + fn map(fields: Vec<(&str, Value)>) -> Struct { + Struct { + fields: fields + .into_iter() + .map(|(name, value)| (name.to_string(), value)) + .collect(), + } + } + + #[test] + fn parses_sizes_and_durations() { + assert_eq!(parse_bytes("20MB").unwrap(), 20_000_000); + assert_eq!(parse_bytes("512").unwrap(), 512); + assert_eq!(parse_duration("5s").unwrap(), Duration::from_secs(5)); + assert_eq!(parse_duration("2m").unwrap(), Duration::from_secs(120)); + } + + #[test] + fn missing_config_uses_flag_defaults() { + let config = Config::parse(None, defaults()).unwrap(); + assert_eq!(config.limits_for("pkg_proxy"), defaults()); + } + + #[test] + fn policy_overrides_inherit_top_level_fields() { + let config = map(vec![ + ("max_bytes", text("10MB")), + ( + "interval", + Value { + kind: Some(Kind::NumberValue(1.0)), + }, + ), + ( + "policies", + Value { + kind: Some(Kind::StructValue(map(vec![( + "other_svc", + Value { + kind: Some(Kind::StructValue(map(vec![ + ("max_bytes", text("5MB")), + ("direction", text("in")), + ]))), + }, + )]))), + }, + ), + ]); + let config = Config::parse(Some(&config), defaults()).unwrap(); + let pkg = config.limits_for("pkg_proxy"); + assert_eq!(pkg.max_bytes, 10_000_000); + assert_eq!(pkg.interval, Duration::from_secs(1)); + let other = config.limits_for("other_svc"); + assert_eq!(other.max_bytes, 5_000_000); + assert_eq!(other.interval, Duration::from_secs(1)); + assert_eq!(other.direction, Direction::In); + } + + #[test] + fn rejects_bad_config() { + for (config, error) in [ + (map(vec![("max_bytes", text("lots"))]), "invalid byte size"), + (map(vec![("max_bytes", text("0"))]), "greater than zero"), + (map(vec![("interval", text("100ms"))]), "between 1s and 60s"), + (map(vec![("window", text("2s"))]), "at least the interval"), + ( + map(vec![("direction", text("up"))]), + "'both', 'in', or 'out'", + ), + ( + map(vec![("budget", text("1MB"))]), + "unsupported field 'config.budget'", + ), + ( + map(vec![( + "policies", + Value { + kind: Some(Kind::StructValue(map(vec![( + "other_svc", + Value { + kind: Some(Kind::StructValue(map(vec![("policies", text("x"))]))), + }, + )]))), + }, + )]), + "unsupported field 'config.policies.other_svc.policies'", + ), + ] { + let result = Config::parse(Some(&config), defaults()); + assert!( + result + .as_ref() + .is_err_and(|message| message.contains(error)), + "{result:?} should contain {error}" + ); + } + } +} diff --git a/examples/supervisor-middleware-egress-usage/src/main.rs b/examples/supervisor-middleware-egress-usage/src/main.rs new file mode 100644 index 0000000000..68519ee39f --- /dev/null +++ b/examples/supervisor-middleware-egress-usage/src/main.rs @@ -0,0 +1,209 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Example OpenShell TCP middleware: egress usage. +//! +//! Binds `TCP_CONNECTION` at `PRE_CONNECT` (`TcpHooks.EvaluateConnection`) +//! and selects STATS mode for every connection. It sums upstream bytes into a +//! sliding window per sandbox and network policy key. When a window passes +//! `max_bytes`, it closes that key's open connections and denies new ones for +//! `block_for`. +//! +//! Limits come from the policy's `config` (see `config.rs`), so changing the +//! policy changes the budget. The flags are the defaults for unset fields. + +mod config; +mod usage; + +use std::net::SocketAddr; +use std::time::Duration; + +use clap::Parser; +use openshell_core::extension_protocol::{ + ExtensionFamily, tcp_hooks_middleware_metadata, validate_gateway_metadata, +}; +use openshell_core::middleware::{TcpHookResultStream, WebSocketResponseStream}; +use openshell_core::proto::middleware::v1::supervisor_middleware_server::{ + SupervisorMiddleware, SupervisorMiddlewareServer, +}; +use openshell_core::proto::middleware::v1::tcp_hooks_server::{TcpHooks, TcpHooksServer}; +use openshell_core::proto::{ + HttpRequestEvaluation, HttpRequestResult, MiddlewareBinding, MiddlewareDescribeRequest, + MiddlewareManifest, SupervisorMiddlewareOperation, SupervisorMiddlewarePhase, TcpHookEvent, + ValidateConfigRequest, ValidateConfigResponse, WebSocketSessionEvent, +}; +use tonic::transport::Server; +use tonic::{Request, Response, Status}; + +use crate::config::{parse_bytes, parse_duration}; +use crate::usage::{Direction, Limits, Usage}; + +const MANIFEST_NAME: &str = "example/egress-usage"; + +#[derive(Debug, Parser)] +#[command(about = "Run the example OpenShell egress-usage TCP middleware")] +/// Limit flags are defaults; the policy's middleware config overrides them. +struct Cli { + /// Address on which to serve plaintext gRPC. + #[arg(long, default_value = "127.0.0.1:50061")] + bind: SocketAddr, + /// Bytes allowed per window per sandbox and policy key, such as 20MB. + #[arg(long, default_value = "20MB", value_parser = parse_bytes)] + max_bytes: u64, + /// Sliding window length, such as 60s. + #[arg(long, default_value = "60s", value_parser = parse_duration)] + window: Duration, + /// Stats interval requested from OpenShell, such as 5s. + #[arg(long, default_value = "5s", value_parser = parse_duration)] + interval: Duration, + /// How long a tripped key stays blocked, such as 300s. + #[arg(long, default_value = "300s", value_parser = parse_duration)] + block_for: Duration, + /// Which bytes count toward the window. + #[arg(long, value_enum, default_value_t = Direction::Both)] + direction: Direction, +} + +#[derive(Clone)] +struct EgressUsage { + usage: Usage, +} + +#[tonic::async_trait] +impl SupervisorMiddleware for EgressUsage { + type EvaluateWebSocketSessionStream = WebSocketResponseStream; + + async fn describe( + &self, + request: Request, + ) -> Result, Status> { + let manifest = MiddlewareManifest { + name: MANIFEST_NAME.into(), + service_version: env!("CARGO_PKG_VERSION").into(), + bindings: vec![MiddlewareBinding { + operation: SupervisorMiddlewareOperation::TcpConnection as i32, + phase: SupervisorMiddlewarePhase::PreConnect as i32, + ..Default::default() + }], + expected_audience: String::new(), + extension: Some(tcp_hooks_middleware_metadata( + MANIFEST_NAME, + openshell_core::VERSION, + )), + }; + validate_gateway_metadata( + ExtensionFamily::SupervisorMiddleware, + MANIFEST_NAME, + manifest.extension.as_ref(), + request.into_inner().gateway, + ) + .map_err(|error| Status::failed_precondition(error.to_string()))?; + Ok(Response::new(manifest)) + } + + async fn validate_config( + &self, + request: Request, + ) -> Result, Status> { + // The gateway calls this when a policy is set, so a bad budget is + // rejected before any sandbox loads it. + let response = match self.usage.config(request.into_inner().config.as_ref()) { + Ok(_) => ValidateConfigResponse { + valid: true, + reason: String::new(), + }, + Err(reason) => ValidateConfigResponse { + valid: false, + reason, + }, + }; + Ok(Response::new(response)) + } + + async fn evaluate_http_request( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unimplemented( + "egress usage implements TCP hooks only", + )) + } + + async fn evaluate_web_socket_session( + &self, + _request: Request>, + ) -> Result, Status> { + Err(Status::unimplemented( + "egress usage implements TCP hooks only", + )) + } +} + +#[tonic::async_trait] +impl TcpHooks for EgressUsage { + type EvaluateConnectionStream = TcpHookResultStream; + + async fn evaluate_connection( + &self, + request: Request>, + ) -> Result, Status> { + Ok(Response::new( + self.usage.connection_stream(request.into_inner()), + )) + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let cli = Cli::parse(); + let limits = Limits { + max_bytes: cli.max_bytes, + window: cli.window, + interval: cli.interval, + block_for: cli.block_for, + direction: cli.direction, + }; + println!( + "serving {MANIFEST_NAME} on http://{} max_bytes={} window={:?} interval={:?} block_for={:?} direction={:?}", + cli.bind, + limits.max_bytes, + limits.window, + limits.interval, + limits.block_for, + limits.direction + ); + let service = EgressUsage { + usage: Usage::new(limits), + }; + Server::builder() + .add_service(SupervisorMiddlewareServer::new(service.clone())) + .add_service(TcpHooksServer::new(service)) + .serve(cli.bind) + .await?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn example_policy_is_valid() { + let policy = openshell_policy::parse_sandbox_policy(include_str!("../policy.yaml")) + .expect("example policy must parse"); + openshell_policy::validate_sandbox_policy(&policy).expect("example policy must be valid"); + let defaults = Cli::parse_from(["egress-usage"]); + let usage = Usage::new(Limits { + max_bytes: defaults.max_bytes, + window: defaults.window, + interval: defaults.interval, + block_for: defaults.block_for, + direction: defaults.direction, + }); + let config = usage + .config(policy.network_middlewares["egress-usage"].config.as_ref()) + .expect("example config must be valid"); + assert_eq!(config.limits_for("pkg_proxy").max_bytes, 20_000_000); + assert_eq!(config.limits_for("other_svc").max_bytes, 50_000_000); + } +} diff --git a/examples/supervisor-middleware-egress-usage/src/usage.rs b/examples/supervisor-middleware-egress-usage/src/usage.rs new file mode 100644 index 0000000000..7ecc31a2ab --- /dev/null +++ b/examples/supervisor-middleware-egress-usage/src/usage.rs @@ -0,0 +1,287 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Sliding egress windows per sandbox and policy key. + +use std::collections::{HashMap, VecDeque}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use openshell_core::middleware::TcpHookResultStream; +use openshell_core::proto::{ + MiddlewareDiagnostics, TcpClose, TcpHookEvent, TcpHookResult, TcpInspect, TcpPreflight, + TcpPreflightResult, TcpStatsMode, tcp_hook_event, tcp_hook_result, tcp_inspect, + tcp_preflight_result, +}; +use tokio::sync::mpsc; +use tokio_stream::StreamExt; +use tokio_stream::wrappers::ReceiverStream; +use tonic::Status; + +use crate::config::Config; + +/// Reason code sent with every close and deny. +pub const REASON_CODE: &str = "egress_budget_exceeded"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] +pub enum Direction { + /// bytes_in plus bytes_out. + Both, + /// bytes_in only (downloads). + In, + /// bytes_out only (uploads). + Out, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Limits { + pub max_bytes: u64, + pub window: Duration, + pub interval: Duration, + pub block_for: Duration, + pub direction: Direction, +} + +/// Sandbox id and network policy key. +type Key = (String, String); + +#[derive(Default)] +struct Window { + samples: VecDeque<(Instant, u64)>, + total: u64, +} + +impl Window { + fn add(&mut self, now: Instant, bytes: u64, window: Duration) -> u64 { + while let Some(&(at, bytes)) = self.samples.front() { + if now.duration_since(at) <= window { + break; + } + self.samples.pop_front(); + self.total -= bytes; + } + if bytes > 0 { + self.samples.push_back((now, bytes)); + self.total += bytes; + } + self.total + } +} + +type ResultSender = mpsc::Sender>; + +#[derive(Default)] +struct State { + windows: HashMap, + blocked_until: HashMap, + open: HashMap>, +} + +#[derive(Clone)] +pub struct Usage { + /// Flag values, used where the policy config leaves a field unset. + defaults: Limits, + state: Arc>, + next_id: Arc, +} + +fn close_result() -> TcpHookResult { + TcpHookResult { + result: Some(tcp_hook_result::Result::Close(TcpClose { + diagnostics: Some(MiddlewareDiagnostics { + reason: "sandbox egress over budget".into(), + reason_code: REASON_CODE.into(), + ..Default::default() + }), + })), + } +} + +fn key_of(preflight: &TcpPreflight) -> Key { + let sandbox = preflight + .context + .as_ref() + .map(|context| context.sandbox_id.clone()) + .unwrap_or_default(); + let policy = preflight + .policy_keys + .first() + .cloned() + .unwrap_or_else(|| preflight.endpoint_id.clone()); + (sandbox, policy) +} + +impl Usage { + pub fn new(defaults: Limits) -> Self { + Self { + defaults, + state: Arc::default(), + next_id: Arc::default(), + } + } + + /// Parses a policy config, filling unset fields from the flags. + pub fn config(&self, config: Option<&prost_types::Struct>) -> Result { + Config::parse(config, self.defaults) + } +} + +fn counted(limits: &Limits, bytes_out: u64, bytes_in: u64) -> u64 { + match limits.direction { + Direction::Both => bytes_out + bytes_in, + Direction::In => bytes_in, + Direction::Out => bytes_out, + } +} + +impl Usage { + pub fn connection_stream( + &self, + mut events: tonic::Streaming, + ) -> TcpHookResultStream { + let (tx, rx) = mpsc::channel(4); + let usage = self.clone(); + tokio::spawn(async move { + let Some(Ok(TcpHookEvent { + event: Some(tcp_hook_event::Event::Preflight(preflight)), + })) = events.next().await + else { + let _ = tx + .send(Err(Status::invalid_argument( + "first event must be preflight", + ))) + .await; + return; + }; + let key = key_of(&preflight); + // The policy config is read per connection, so a policy update + // applies to connections opened after it loads. + let limits = match usage.config(preflight.config.as_ref()) { + Ok(config) => config.limits_for(&key.1), + Err(error) => { + let _ = tx.send(Err(Status::invalid_argument(error))).await; + return; + } + }; + let id = usage.next_id.fetch_add(1, Ordering::Relaxed); + let blocked = { + let mut state = usage.state.lock().unwrap(); + let now = Instant::now(); + let blocked = state + .blocked_until + .get(&key) + .is_some_and(|until| *until > now); + if !blocked { + state + .open + .entry(key.clone()) + .or_default() + .insert(id, tx.clone()); + } + blocked + }; + if blocked { + println!( + "deny sandbox={} policy={} host={}:{} reason={REASON_CODE}", + key.0, key.1, preflight.host, preflight.port + ); + let _ = tx.send(Ok(close_result())).await; + return; + } + let interval = prost_types::Duration::try_from(limits.interval).ok(); + let _ = tx + .send(Ok(TcpHookResult { + result: Some(tcp_hook_result::Result::PreflightResult( + TcpPreflightResult { + decision: Some(tcp_preflight_result::Decision::Inspect(TcpInspect { + mode: Some(tcp_inspect::Mode::Stats(TcpStatsMode { interval })), + })), + diagnostics: None, + }, + )), + })) + .await; + let mut last = 0_u64; + while let Some(Ok(event)) = events.next().await { + let (bytes_out, bytes_in, end) = match event.event { + Some(tcp_hook_event::Event::Stats(stats)) => { + (stats.bytes_out, stats.bytes_in, false) + } + Some(tcp_hook_event::Event::End(end)) => (end.bytes_out, end.bytes_in, true), + _ => continue, + }; + let counted = counted(&limits, bytes_out, bytes_in); + let delta = counted.saturating_sub(last); + last = counted; + usage.record(&key, &preflight, &limits, delta); + if end { + break; + } + } + let mut state = usage.state.lock().unwrap(); + if let Some(open) = state.open.get_mut(&key) { + open.remove(&id); + } + }); + Box::pin(ReceiverStream::new(rx)) + } + + /// Add `delta` to the key's window and trip it when over the limit. + fn record(&self, key: &Key, preflight: &TcpPreflight, limits: &Limits, delta: u64) { + let now = Instant::now(); + let mut state = self.state.lock().unwrap(); + let total = state + .windows + .entry(key.clone()) + .or_default() + .add(now, delta, limits.window); + let already_blocked = state + .blocked_until + .get(key) + .is_some_and(|until| *until > now); + if total <= limits.max_bytes || already_blocked { + return; + } + state + .blocked_until + .insert(key.clone(), now + limits.block_for); + let open = state.open.remove(key).unwrap_or_default(); + println!( + "TRIP sandbox={} policy={} host={}:{} window_bytes={total} max_bytes={} over_by={} closing={} block_for={:?}", + key.0, + key.1, + preflight.host, + preflight.port, + limits.max_bytes, + total - limits.max_bytes, + open.len(), + limits.block_for, + ); + for sender in open.into_values() { + // Never block the lock holder: a full result channel already has + // a pending result, and the connection is denied on its next one. + let _ = sender.try_send(Ok(close_result())); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn window_expires_old_samples() { + let mut window = Window::default(); + let start = Instant::now(); + assert_eq!(window.add(start, 10, Duration::from_secs(60)), 10); + assert_eq!( + window.add(start + Duration::from_secs(30), 5, Duration::from_secs(60)), + 15 + ); + assert_eq!( + window.add(start + Duration::from_secs(61), 1, Duration::from_secs(60)), + 6 + ); + } +} diff --git a/examples/supervisor-middleware-egress-usage/upstream.py b/examples/supervisor-middleware-egress-usage/upstream.py new file mode 100644 index 0000000000..7417171f53 --- /dev/null +++ b/examples/supervisor-middleware-egress-usage/upstream.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Package proxy stand-in for the egress-usage example. + +Runs on the host. Docker sandboxes reach it as host.openshell.internal, so +policy keys are told apart by port. + +- 47080 (pkg_proxy) and 47081 (other_svc): plain HTTP. GET /bytes/ streams + n bytes. PUT or POST consumes the body and returns its length. +- 47443 (pkg_proxy): the same HTTP handler over TLS (cert and key from + UPSTREAM_TLS_CERT and UPSTREAM_TLS_KEY), for tls: skip passthrough. +- 47900 (pkg_proxy): raw TCP. Send "PULL \\n" to receive n bytes, or + "PUSH\\n" followed by bytes and a half-close to receive the count. A stream + that starts with an HTTP method is sniffed as HTTP, hence the verbs. +""" + +import os +import socketserver +import ssl +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +CHUNK = b"x" * 65536 + + +def stream_bytes(write, count): + while count > 0: + piece = CHUNK[: min(count, len(CHUNK))] + write(piece) + count -= len(piece) + + +class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_GET(self): + parts = self.path.strip("/").split("/") + if len(parts) == 2 and parts[0] == "bytes" and parts[1].isdigit(): + count = int(parts[1]) + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(count)) + self.end_headers() + try: + stream_bytes(self.wfile.write, count) + except (BrokenPipeError, ConnectionResetError): + pass + return + body = b"ok\n" + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _consume(self): + length = int(self.headers.get("Content-Length", "0")) + received = 0 + while received < length: + data = self.rfile.read(min(65536, length - received)) + if not data: + break + received += len(data) + body = f"{received}\n".encode() + self.send_response(200) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + do_PUT = _consume + do_POST = _consume + + def log_message(self, format, *args): + print(f"{self.client_address[0]} {format % args}", flush=True) + + +class RawHandler(socketserver.StreamRequestHandler): + def handle(self): + line = self.rfile.readline().decode(errors="replace").strip() + try: + if line.startswith("PULL "): + stream_bytes(self.wfile.write, int(line[5:])) + elif line == "PUSH": + received = 0 + while True: + data = self.rfile.read1(65536) + if not data: + break + received += len(data) + self.wfile.write(f"{received}\n".encode()) + except (BrokenPipeError, ConnectionResetError): + pass + print(f"raw {self.client_address[0]} {line}", flush=True) + + +class RawServer(socketserver.ThreadingTCPServer): + allow_reuse_address = True + daemon_threads = True + + +def main(): + bind = os.environ.get("UPSTREAM_BIND", "127.0.0.1") + servers = [ + ThreadingHTTPServer((bind, 47080), Handler), + ThreadingHTTPServer((bind, 47081), Handler), + RawServer((bind, 47900), RawHandler), + ] + cert, key = os.environ.get("UPSTREAM_TLS_CERT"), os.environ.get("UPSTREAM_TLS_KEY") + if cert and key: + tls = ThreadingHTTPServer((bind, 47443), Handler) + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(cert, key) + tls.socket = context.wrap_socket(tls.socket, server_side=True) + servers.append(tls) + for server in servers[1:]: + threading.Thread(target=server.serve_forever, daemon=True).start() + print("upstream ready", flush=True) + servers[0].serve_forever() + + +if __name__ == "__main__": + main() diff --git a/proto/sandbox.proto b/proto/sandbox.proto index 80344bd4a1..2369bba049 100644 --- a/proto/sandbox.proto +++ b/proto/sandbox.proto @@ -534,4 +534,5 @@ enum SupervisorMiddlewareBinding { SUPERVISOR_MIDDLEWARE_BINDING_HTTP_REQUEST_PRE_CREDENTIALS = 1; SUPERVISOR_MIDDLEWARE_BINDING_HTTP_RESPONSE_PRE_RETURN = 2; SUPERVISOR_MIDDLEWARE_BINDING_WEBSOCKET_MESSAGE_PRE_CREDENTIALS = 3; + SUPERVISOR_MIDDLEWARE_BINDING_TCP_CONNECTION_PRE_CONNECT = 4; } diff --git a/proto/supervisor_middleware.proto b/proto/supervisor_middleware.proto index 14ffe09f21..342def3b04 100644 --- a/proto/supervisor_middleware.proto +++ b/proto/supervisor_middleware.proto @@ -60,6 +60,19 @@ service HttpHooks { returns (stream HttpResponseHookResult); } +// TcpHooks observes and controls admitted upstream TCP connections for one +// operator-run middleware. It shares SupervisorMiddleware's registration, +// discovery, configuration, and endpoint. +service TcpHooks { + // One stream per admitted upstream connection and middleware stage. + // OpenShell calls it for the TCP_CONNECTION bindings of a service whose + // manifest lists openshell.supervisor-middleware.tcp-hooks in + // extension.required_capabilities. The first event is a preflight, sent + // after network-policy admission and before OpenShell dials the upstream. + // The first result answers it. A close result may follow at any time. + rpc EvaluateConnection(stream TcpHookEvent) returns (stream TcpHookResult); +} + message MiddlewareDescribeRequest { // Gateway or supervisor protocol metadata. Middleware must reject unmet // requirements before accepting traffic. @@ -916,6 +929,10 @@ enum SupervisorMiddlewareOperation { // hooks capability, and the deprecated HttpResponsePreReturn.Evaluate // otherwise. SUPERVISOR_MIDDLEWARE_OPERATION_HTTP_RESPONSE = 3; + // An admitted upstream TCP connection at PRE_CONNECT. OpenShell calls + // TcpHooks.EvaluateConnection; the service must require the + // openshell.supervisor-middleware.tcp-hooks capability. + SUPERVISOR_MIDDLEWARE_OPERATION_TCP_CONNECTION = 4; } // Ordered phase within a supervisor operation. @@ -923,6 +940,8 @@ enum SupervisorMiddlewarePhase { SUPERVISOR_MIDDLEWARE_PHASE_UNSPECIFIED = 0; SUPERVISOR_MIDDLEWARE_PHASE_PRE_CREDENTIALS = 1; SUPERVISOR_MIDDLEWARE_PHASE_PRE_RETURN = 2; + // After network-policy admission and before OpenShell dials the upstream. + SUPERVISOR_MIDDLEWARE_PHASE_PRE_CONNECT = 3; } // WebSocketSessionEvent is one ordered event in a stage-local stream. @@ -1183,3 +1202,124 @@ message HttpRequestResult { // OpenShell may return this code to the requester, unlike free-form reason. string reason_code = 8; } + +// TCP hooks. One ordered event in a TcpHooks.EvaluateConnection stream. +// Preflight is always first. stats events follow only in STATS mode. end is +// best effort and has no result. +message TcpHookEvent { + oneof event { + TcpPreflight preflight = 1; + TcpStats stats = 2; + TcpEnd end = 3; + } +} + +// One result in a TcpHooks.EvaluateConnection stream. The first result must +// be preflight_result or close. close may also follow a preflight result at +// any time to abort the open connection. +message TcpHookResult { + oneof result { + TcpPreflightResult preflight_result = 1; + TcpClose close = 2; + } +} + +// How OpenShell carries the connection's bytes after admission. +enum TcpInspection { + TCP_INSPECTION_UNSPECIFIED = 0; + // CONNECT tunnel with tls: skip: raw bytes, usually TLS ciphertext. + TCP_INSPECTION_TLS_PASSTHROUGH = 1; + // CONNECT tunnel whose payload OpenShell detects after the tunnel opens: + // terminated TLS, plaintext HTTP, or raw TCP. + TCP_INSPECTION_CONNECT_TUNNEL = 2; + // Transparent TCP redirected without a proxy request. + TCP_INSPECTION_TRANSPARENT_TCP = 3; + // Plain HTTP request sent to the forward proxy. + TCP_INSPECTION_HTTP_FORWARD = 4; +} + +// Observation mode a TCP stage can select at preflight. +enum TcpMode { + TCP_MODE_UNSPECIFIED = 0; + // Cumulative byte totals at the stage's interval. The stage never sees + // connection content. + TCP_MODE_STATS = 1; +} + +message TcpPreflight { + // Destination hostname selected by network policy. + string host = 1; + // Destination TCP port. + uint32 port = 2; + // Reserved: TLS server name, when known before dial. Currently empty. + string sni = 3; + TcpInspection inspection = 4; + // Network policy keys that admitted the connection. + repeated string policy_keys = 5; + // Stable endpoint identity, currently "host:port". + string endpoint_id = 6; + // Sandbox, connection, and originating process identity. request_id + // identifies this connection. + RequestContext context = 7; + // Built-in middleware name or operator-owned registration name. + string middleware_name = 8; + // Validated service configuration. + google.protobuf.Struct config = 9; + // Modes this connection is eligible for. + repeated TcpMode permitted_modes = 10; + TcpLimits limits = 11; + // Phase of the TCP_CONNECTION binding this stream serves: PRE_CONNECT. + SupervisorMiddlewarePhase phase = 12; +} + +message TcpLimits { + // Shortest and longest stats interval OpenShell accepts. + google.protobuf.Duration min_stats_interval = 1; + google.protobuf.Duration max_stats_interval = 2; +} + +message TcpPreflightResult { + oneof decision { + // Relay without this stage observing the connection. Only end follows. + TcpContinue continue_without_inspection = 1; + // Observe the connection in one permitted mode. + TcpInspect inspect = 2; + } + MiddlewareDiagnostics diagnostics = 3; +} + +message TcpContinue {} + +message TcpInspect { + oneof mode { + TcpStatsMode stats = 1; + } +} + +message TcpStatsMode { + // Stats interval, clamped to the offered limits. Unset uses the minimum. + google.protobuf.Duration interval = 1; +} + +// Cumulative totals since the upstream connection opened, counted at the +// upstream socket: what left for and arrived from the upstream after every +// transform. Sent only when a total changed since the previous stats event. +message TcpStats { + uint64 bytes_out = 1; + uint64 bytes_in = 2; + google.protobuf.Duration elapsed = 3; +} + +// Final totals. Best effort. +message TcpEnd { + uint64 bytes_out = 1; + uint64 bytes_in = 2; + google.protobuf.Duration duration = 3; + MiddlewareSessionEndReason reason = 4; +} + +// As the first result, denies the connection before OpenShell dials the +// upstream. Later, aborts the open connection in both directions. +message TcpClose { + MiddlewareDiagnostics diagnostics = 1; +}