Forward executor network policy decisions for auditing (#38670)

## What changed

- Add a best-effort `network/policyDecision` notification for final domain and non-domain policy decisions made by executor-local proxies.
- Validate notifications against the active process on the controller and emit audit events with controller-trusted session and execution metadata.
- Reserve outbound RPC capacity so audit notifications cannot block control messages, and expose valid `chatgpt-account-id` header values for audit attribution.

## Testing

- Cover notification serialization, proxy decision capture, executor-to-controller delivery, trusted metadata handling, and reserved RPC capacity.

GitOrigin-RevId: a39f96a6b3d9401c03d54eaef5b9a6d3fe0da78b
This commit is contained in:
viyatb-oai
2026-08-14 23:36:23 +00:00
committed by copyberry
parent a186f5484d
commit 15fde8c1f2
15 changed files with 655 additions and 19 deletions

View File

@@ -12,6 +12,7 @@ use arc_swap::ArcSwap;
use arc_swap::ArcSwapOption;
use codex_exec_server_protocol::JSONRPCNotification;
use codex_network_proxy::NetworkPolicyDecider;
use codex_network_proxy::NetworkProxyAuditMetadata;
use futures::FutureExt;
use futures::future::BoxFuture;
use serde_json::Value;
@@ -106,6 +107,8 @@ use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeParams;
use crate::protocol::InitializeResponse;
use crate::protocol::NETWORK_POLICY_DECISION_METHOD;
use crate::protocol::NetworkPolicyDecisionNotification;
use crate::protocol::ProcessOutputChunk;
use crate::protocol::ProcessSandboxType;
use crate::protocol::ProcessSignal;
@@ -123,6 +126,7 @@ use crate::rpc_server_requests::MAX_IN_FLIGHT_SERVER_CALLS;
use codex_http_client::HttpClientFactory;
pub(crate) mod http_client;
mod network_policy_audit;
#[path = "client_recovery.rs"]
mod recovery;
pub(crate) use recovery::is_retryable_recovery_error;
@@ -189,8 +193,13 @@ pub(crate) struct SessionState {
ordered_events: StdMutex<OrderedSessionEvents>,
recoverable: AtomicBool,
next_write_id: AtomicU64,
network_policy_controller: ArcSwapOption<NetworkPolicyDecisionController>,
network_policy_cancelled: CancellationToken,
network_policy: NetworkPolicyState,
}
struct NetworkPolicyState {
controller: ArcSwapOption<NetworkPolicyDecisionController>,
cancelled: CancellationToken,
audit: Option<NetworkPolicyAuditContext>,
}
#[derive(Clone)]
@@ -199,6 +208,11 @@ struct NetworkPolicyDecisionController {
timeout: Duration,
}
struct NetworkPolicyAuditContext {
metadata: NetworkProxyAuditMetadata,
execution_id: Option<String>,
}
#[derive(Default)]
struct OrderedSessionEvents {
last_published_seq: u64,
@@ -932,10 +946,20 @@ impl ExecServerClient {
}
let process_id = params.process_id.clone();
let state = Arc::new(SessionState::new(/*recoverable*/ false));
let mut state = SessionState::new(/*recoverable*/ false);
state.network_policy.audit =
params
.network_proxy
.as_ref()
.map(|launch| NetworkPolicyAuditContext {
metadata: launch.audit_metadata.clone(),
execution_id: launch.execution_id.clone(),
});
let state = Arc::new(state);
if let Some(controller) = network_policy_controller.as_ref() {
state
.network_policy_controller
.network_policy
.controller
.store(Some(Arc::new(controller.clone())));
}
if let Err(error) = self.inner.insert_session(&process_id, Arc::clone(&state)) {
@@ -1185,8 +1209,11 @@ impl SessionState {
ordered_events: StdMutex::new(OrderedSessionEvents::default()),
recoverable: AtomicBool::new(recoverable),
next_write_id: AtomicU64::new(1),
network_policy_controller: ArcSwapOption::empty(),
network_policy_cancelled: CancellationToken::new(),
network_policy: NetworkPolicyState {
controller: ArcSwapOption::empty(),
cancelled: CancellationToken::new(),
audit: None,
},
}
}
@@ -1448,7 +1475,7 @@ impl Session {
}
pub(crate) fn cancel_network_policy_decisions(&self) {
self.state.network_policy_cancelled.cancel();
self.state.network_policy.cancelled.cancel();
}
pub(crate) async fn unregister(&self) {
@@ -1505,8 +1532,8 @@ impl Inner {
let mut next_sessions = sessions.as_ref().clone();
next_sessions.remove(process_id);
self.sessions.store(Arc::new(next_sessions));
expected.network_policy_cancelled.cancel();
expected.network_policy_controller.store(None);
expected.network_policy.cancelled.cancel();
expected.network_policy.controller.store(None);
}
fn take_all_sessions(&self) -> HashMap<ProcessId, Arc<SessionState>> {
@@ -1545,8 +1572,8 @@ fn fail_all_sessions(inner: &Arc<Inner>, message: String) {
let sessions = inner.take_all_sessions();
for (_, session) in sessions {
session.network_policy_cancelled.cancel();
session.network_policy_controller.store(None);
session.network_policy.cancelled.cancel();
session.network_policy.controller.store(None);
// Sessions synthesize a closed read response and emit a pushed Failed
// event. That covers both polling consumers and streaming consumers
// such as environment-backed MCP stdio.
@@ -1616,6 +1643,24 @@ async fn handle_server_notification(
.handle_http_body_delta_notification(notification.params)
.await?;
}
NETWORK_POLICY_DECISION_METHOD => {
let Ok(params) = serde_json::from_value::<NetworkPolicyDecisionNotification>(
notification.params.unwrap_or(Value::Null),
) else {
debug!("ignoring malformed exec-server network policy decision notification");
return Ok(());
};
let Some(session) = inner.get_session(&params.process_id) else {
debug!("ignoring network policy decision for an unknown exec-server process");
return Ok(());
};
let Some(context) = session.network_policy.audit.as_ref() else {
return Ok(());
};
if !network_policy_audit::emit_network_policy_decision(context, &params) {
debug!("ignoring invalid exec-server network policy decision notification");
}
}
other => {
debug!("ignoring unknown exec-server notification: {other}");
}

View File

@@ -0,0 +1,81 @@
use super::NetworkPolicyAuditContext;
use crate::protocol::ExecServerNetworkProtocol;
use crate::protocol::MAX_NETWORK_POLICY_HOST_BYTES;
use crate::protocol::MAX_NETWORK_POLICY_PROCESS_ID_BYTES;
use crate::protocol::MAX_NETWORK_POLICY_REASON_BYTES;
use crate::protocol::NetworkPolicyDecisionNotification;
const MAX_NETWORK_POLICY_METHOD_BYTES: usize = 32;
const MAX_NETWORK_POLICY_CLIENT_BYTES: usize = 256;
const MAX_NETWORK_POLICY_TIMESTAMP_BYTES: usize = 64;
pub(super) fn emit_network_policy_decision(
context: &NetworkPolicyAuditContext,
decision: &NetworkPolicyDecisionNotification,
) -> bool {
if decision.process_id.is_empty()
|| decision.process_id.len() > MAX_NETWORK_POLICY_PROCESS_ID_BYTES
|| decision.host.is_empty()
|| decision.host.len() > MAX_NETWORK_POLICY_HOST_BYTES
|| decision.host.chars().any(char::is_control)
|| decision.host.chars().any(char::is_whitespace)
|| decision.reason.len() > MAX_NETWORK_POLICY_REASON_BYTES
|| decision.reason.chars().any(char::is_control)
|| !matches!(decision.scope.as_str(), "domain" | "non_domain")
|| !matches!(decision.decision.as_str(), "allow" | "deny" | "ask")
|| !matches!(
decision.source.as_str(),
"baseline_policy" | "mode_guard" | "proxy_state" | "decider"
)
|| decision.timestamp.is_empty()
|| decision.timestamp.len() > MAX_NETWORK_POLICY_TIMESTAMP_BYTES
|| decision.timestamp.chars().any(char::is_control)
|| decision.method.as_ref().is_some_and(|method| {
method.len() > MAX_NETWORK_POLICY_METHOD_BYTES
|| method.chars().any(char::is_control)
|| method.chars().any(char::is_whitespace)
})
|| decision.client.as_ref().is_some_and(|client| {
client.len() > MAX_NETWORK_POLICY_CLIENT_BYTES
|| client.chars().any(char::is_control)
|| client.chars().any(char::is_whitespace)
})
{
return false;
}
let protocol = match decision.protocol {
ExecServerNetworkProtocol::Http => "http",
ExecServerNetworkProtocol::HttpsConnect => "https_connect",
ExecServerNetworkProtocol::Socks5Tcp => "socks5_tcp",
ExecServerNetworkProtocol::Socks5Udp => "socks5_udp",
};
let metadata = &context.metadata;
tracing::event!(
target: "codex_otel.network_proxy",
tracing::Level::INFO,
event.name = "codex.network_proxy.policy_decision",
event.timestamp = decision.timestamp,
conversation.id = metadata.conversation_id.as_deref(),
app.version = metadata.app_version.as_deref(),
auth_mode = metadata.auth_mode.as_deref(),
originator = metadata.originator.as_deref(),
user.account_id = metadata.user_account_id.as_deref(),
user.email = metadata.user_email.as_deref(),
terminal.type = metadata.terminal_type.as_deref(),
model = metadata.model.as_deref(),
slug = metadata.slug.as_deref(),
network.policy.scope = decision.scope,
network.policy.decision = decision.decision,
network.policy.source = decision.source,
network.policy.reason = decision.reason,
network.transport.protocol = protocol,
server.address = decision.host,
server.port = decision.port,
http.request.method = decision.method.as_deref().unwrap_or("none"),
client.address = decision.client.as_deref().unwrap_or("unknown"),
execution.id = context.execution_id.as_deref(),
network.policy.override = decision.policy_override,
);
true
}

View File

@@ -1,7 +1,9 @@
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCNotification;
use codex_exec_server_protocol::JSONRPCRequest;
use codex_exec_server_protocol::JSONRPCResponse;
use codex_exec_server_protocol::RequestId;
@@ -10,15 +12,21 @@ use codex_http_client::OutboundProxyPolicy;
use codex_network_proxy::NetworkDecision;
use codex_network_proxy::NetworkPolicyDecider;
use codex_network_proxy::NetworkPolicyRequest;
use codex_network_proxy::NetworkProxyAuditMetadata;
use codex_utils_path_uri::PathUri;
use pretty_assertions::assert_eq;
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::time::timeout;
use tracing::instrument::WithSubscriber;
use tracing_subscriber::prelude::*;
use super::super::LazyRemoteExecServerClient;
use super::super::NetworkPolicyAuditContext;
use super::super::NetworkPolicyDecisionController;
use super::super::SessionState;
use super::super::handle_server_notification;
use super::accept_websocket;
use super::complete_websocket_initialize;
use super::read_jsonrpc_websocket;
@@ -31,7 +39,9 @@ use crate::protocol::ExecParams;
use crate::protocol::ExecServerNetworkPolicyDecision;
use crate::protocol::ExecServerNetworkPolicyRequest;
use crate::protocol::ExecServerNetworkProtocol;
use crate::protocol::NETWORK_POLICY_DECISION_METHOD;
use crate::protocol::NETWORK_POLICY_REQUEST_METHOD;
use crate::protocol::NetworkPolicyDecisionNotification;
use crate::protocol::NetworkPolicyRequestParams;
use crate::protocol::NetworkPolicyRequestResponse;
use crate::rpc_server_requests::MAX_IN_FLIGHT_SERVER_CALLS;
@@ -76,6 +86,129 @@ async fn read_decision(
.decision
}
#[tokio::test(flavor = "current_thread")]
async fn policy_decisions_reject_forged_process_and_use_trusted_controller_metadata() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let websocket_url = format!("ws://{}", listener.local_addr().expect("listener address"));
let (release_tx, release_rx) = oneshot::channel();
let (initialized_tx, initialized_rx) = oneshot::channel();
let server = tokio::spawn(async move {
let mut websocket = accept_websocket(&listener).await;
complete_websocket_initialize(
&mut websocket,
"audit-session",
/*expected_resume_session_id*/ None,
)
.await;
initialized_tx
.send(())
.expect("client should await completed WebSocket initialization");
release_rx.await.expect("server should be released");
});
let logs = Arc::new(Mutex::new(Vec::new()));
let writer_logs = Arc::clone(&logs);
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.with_ansi(false)
.with_writer(move || AuditLogWriter(Arc::clone(&writer_logs))),
);
async move {
let client = LazyRemoteExecServerClient::new(
ExecServerTransportParams::websocket_url(websocket_url, Duration::from_secs(1)),
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
)
.get()
.await
.expect("client should connect");
initialized_rx
.await
.expect("server should complete WebSocket initialization");
let mut state = SessionState::new(/*recoverable*/ true);
state.network_policy.audit = Some(NetworkPolicyAuditContext {
metadata: NetworkProxyAuditMetadata {
conversation_id: Some("trusted-conversation".to_string()),
user_account_id: Some("trusted-account".to_string()),
..NetworkProxyAuditMetadata::default()
},
execution_id: Some("trusted-execution".to_string()),
});
client
.inner
.insert_session(&ProcessId::from("trusted-process"), Arc::new(state))
.expect("trusted process should register");
for (process_id, host) in [
("forged-process", "forged.example"),
("trusted-process", "trusted.example"),
] {
handle_server_notification(
&client.inner,
JSONRPCNotification {
method: NETWORK_POLICY_DECISION_METHOD.to_string(),
params: Some(
serde_json::to_value(NetworkPolicyDecisionNotification {
process_id: ProcessId::from(process_id),
timestamp: "2026-08-11T12:00:00.000Z".to_string(),
scope: "domain".to_string(),
decision: "deny".to_string(),
source: "baseline_policy".to_string(),
reason: "not_allowed".to_string(),
protocol: ExecServerNetworkProtocol::HttpsConnect,
host: host.to_string(),
port: 443,
method: None,
client: None,
policy_override: false,
})
.expect("network policy decision should serialize"),
),
},
)
.await
.expect("controller should handle network policy notification");
}
let output = String::from_utf8(
logs.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone(),
)
.expect("audit log should be UTF-8");
assert!(!output.contains("forged.example"));
for expected in [
"trusted-conversation",
"trusted-account",
"trusted-execution",
] {
assert!(
output.contains(expected),
"missing `{expected}` in {output}"
);
}
release_tx.send(()).expect("server should be released");
}
.with_subscriber(subscriber)
.await;
server.await.expect("server should finish");
}
struct AuditLogWriter(Arc<Mutex<Vec<u8>>>);
impl std::io::Write for AuditLogWriter {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
self.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.extend_from_slice(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
#[tokio::test]
async fn abandoned_process_start_unregisters_and_cleans_up() {
let listener = TcpListener::bind("127.0.0.1:0")
@@ -155,7 +288,8 @@ async fn abandoned_process_start_unregisters_and_cleans_up() {
Arc::new(|_request: NetworkPolicyRequest| async { NetworkDecision::Allow });
let decider_weak = Arc::downgrade(&decider);
state
.network_policy_controller
.network_policy
.controller
.store(Some(Arc::new(NetworkPolicyDecisionController {
decider,
timeout: Duration::from_secs(30),
@@ -163,7 +297,7 @@ async fn abandoned_process_start_unregisters_and_cleans_up() {
start.abort();
assert!(start.await.is_err_and(|error| error.is_cancelled()));
assert!(state.network_policy_cancelled.is_cancelled());
assert!(state.network_policy.cancelled.is_cancelled());
assert!(client.inner.get_session(&process_id).is_none());
assert!(decider_weak.upgrade().is_none());
@@ -305,7 +439,7 @@ async fn policy_requests_use_process_decider_and_cancel_on_unregister() {
}
}
});
session.state.network_policy_controller.store(Some(Arc::new(
session.state.network_policy.controller.store(Some(Arc::new(
NetworkPolicyDecisionController {
decider,
timeout: Duration::from_secs(30),

View File

@@ -646,10 +646,10 @@ impl ExecServerClient {
.flatten();
let controller = session
.as_ref()
.and_then(|session| session.network_policy_controller.load_full());
.and_then(|session| session.network_policy.controller.load_full());
let process_cancelled = session
.as_ref()
.map(|session| session.network_policy_cancelled.clone());
.map(|session| session.network_policy.cancelled.clone());
let expected_session = session.as_ref().map(Arc::downgrade);
let policy_request =
(process_id_valid && host_valid).then_some(NetworkPolicyRequest {

View File

@@ -8,6 +8,9 @@ use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_exec_server_protocol::JSONRPCErrorError;
use codex_network_proxy::NetworkPolicyAuditEvent;
use codex_network_proxy::NetworkPolicyAuditObserver;
use codex_network_proxy::NetworkProtocol;
use codex_network_proxy::NetworkProxyHandle;
use codex_protocol::config_types::EnvironmentVariablePattern;
use codex_protocol::config_types::ShellEnvironmentPolicy;
@@ -46,7 +49,10 @@ use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecOutputStream;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::ExecServerNetworkProtocol;
use crate::protocol::MAX_NETWORK_POLICY_PROCESS_ID_BYTES;
use crate::protocol::NETWORK_POLICY_DECISION_METHOD;
use crate::protocol::NetworkPolicyDecisionNotification;
use crate::protocol::ProcessOutputChunk;
use crate::protocol::ProcessSandboxType;
use crate::protocol::ProcessSignal;
@@ -294,11 +300,44 @@ impl LocalProcess {
process_shutdown.clone(),
)
});
let network_policy_audit_observer = params.network_proxy.as_ref().map(|_| {
let process_id = process_id.clone();
let inner = Arc::downgrade(&self.inner);
Arc::new(move |event: NetworkPolicyAuditEvent| {
let Some(inner) = inner.upgrade() else {
return;
};
let Some(notifications) = notification_sender(&inner) else {
return;
};
let notification = NetworkPolicyDecisionNotification {
process_id: process_id.clone(),
timestamp: event.timestamp,
scope: event.scope,
decision: event.decision,
source: event.source,
reason: event.reason,
protocol: match event.protocol {
NetworkProtocol::Http => ExecServerNetworkProtocol::Http,
NetworkProtocol::HttpsConnect => ExecServerNetworkProtocol::HttpsConnect,
NetworkProtocol::Socks5Tcp => ExecServerNetworkProtocol::Socks5Tcp,
NetworkProtocol::Socks5Udp => ExecServerNetworkProtocol::Socks5Udp,
},
host: event.host,
port: event.port,
method: event.method,
client: event.client,
policy_override: event.policy_override,
};
let _ = notifications.try_notify(NETWORK_POLICY_DECISION_METHOD, &notification);
}) as NetworkPolicyAuditObserver
});
let prepared = prepare_exec_request(
&params,
child_env(&params),
self.runtime_paths.as_ref(),
network_policy_decider,
network_policy_audit_observer,
)
.await?;
if prepared.command.is_empty() {
@@ -1112,6 +1151,102 @@ mod tests {
}
}
#[cfg(unix)]
#[tokio::test]
async fn executor_proxy_sends_final_network_policy_notification() {
let (outgoing_tx, mut outgoing_rx) = mpsc::channel(NOTIFICATION_CHANNEL_CAPACITY);
let backend = LocalProcess::with_runtime_paths(
RpcNotificationSender::new(outgoing_tx),
ExecServerTelemetry::default(),
/*runtime_paths*/ None,
);
let proxy_config = RemoteNetworkProxyConfig::from_effective_config(&NetworkProxyConfig {
enabled: true,
..NetworkProxyConfig::default()
})
.expect("build remote network proxy config");
let mut params = test_exec_params(HashMap::new());
params.process_id = ProcessId::from("audit-process");
params.argv = vec![
"/bin/sh".to_string(),
"-c".to_string(),
"printf '%s\\n' \"$HTTP_PROXY\"; exec sleep 60".to_string(),
];
params.network_proxy = Some(
RemoteNetworkProxyLaunchConfig::new(proxy_config)
.for_execution("environment-1".to_string(), "execution-1".to_string()),
);
backend
.exec(params)
.await
.expect("start process with proxy");
let output = backend
.exec_read(ReadParams {
process_id: ProcessId::from("audit-process"),
after_seq: None,
max_bytes: None,
wait_ms: Some(1_000),
})
.await
.expect("read executor proxy address");
let proxy_addr = String::from_utf8(
output
.chunks
.into_iter()
.find(|chunk| matches!(chunk.stream, ExecOutputStream::Stdout))
.expect("executor proxy address output")
.chunk
.into_inner(),
)
.expect("UTF-8 proxy address");
let proxy_addr = proxy_addr
.trim()
.strip_prefix("http://")
.expect("HTTP executor proxy address");
let mut stream = tokio::net::TcpStream::connect(proxy_addr)
.await
.expect("connect to executor proxy");
stream
.write_all(b"CONNECT 8.8.8.8:443 HTTP/1.1\r\nHost: 8.8.8.8:443\r\n\r\n")
.await
.expect("write CONNECT request");
let mut response = [0_u8; 256];
let response_len = timeout(Duration::from_secs(2), stream.read(&mut response))
.await
.expect("proxy response timeout")
.expect("read proxy response");
assert!(String::from_utf8_lossy(&response[..response_len]).starts_with("HTTP/1.1 403"));
let notification = timeout(Duration::from_secs(2), async {
loop {
match outgoing_rx.recv().await {
Some(RpcServerOutboundMessage::Notification(notification))
if notification.method == NETWORK_POLICY_DECISION_METHOD =>
{
break serde_json::from_value::<NetworkPolicyDecisionNotification>(
notification
.params
.expect("network policy notification params"),
)
.expect("deserialize network policy notification");
}
Some(_) => {}
None => panic!("outbound notifications closed"),
}
}
})
.await
.expect("network policy notification timeout");
assert_eq!(notification.process_id, ProcessId::from("audit-process"));
assert_eq!(notification.decision, "deny");
assert_eq!(notification.host, "8.8.8.8");
assert_eq!(
notification.protocol,
ExecServerNetworkProtocol::HttpsConnect
);
backend.shutdown().await;
}
fn telemetry_backend() -> (
LocalProcess,
codex_otel::MetricsClient,

View File

@@ -4,6 +4,7 @@ use std::sync::Arc;
use codex_exec_server_protocol::JSONRPCErrorError;
use codex_network_proxy::CUSTOM_CA_ENV_KEYS;
use codex_network_proxy::ManagedNetworkSandboxContext;
use codex_network_proxy::NetworkPolicyAuditObserver;
use codex_network_proxy::NetworkPolicyDecider;
use codex_network_proxy::NetworkProxy;
use codex_network_proxy::NetworkProxyHandle;
@@ -79,6 +80,7 @@ pub(crate) async fn prepare_exec_request(
env: HashMap<String, String>,
runtime_paths: Option<&ExecServerRuntimePaths>,
network_policy_decider: Option<Arc<dyn NetworkPolicyDecider>>,
network_policy_audit_observer: Option<NetworkPolicyAuditObserver>,
) -> Result<PreparedExecRequest, JSONRPCErrorError> {
#[cfg(target_os = "windows")]
let mut env = env;
@@ -102,6 +104,7 @@ pub(crate) async fn prepare_exec_request(
network_proxy,
env,
network_policy_decider,
network_policy_audit_observer,
)
.await?;
let Some(sandbox_context) = params.sandbox.as_ref() else {
@@ -295,6 +298,7 @@ async fn prepare_managed_network(
network_proxy: Option<&RemoteNetworkProxyLaunchConfig>,
env: HashMap<String, String>,
network_policy_decider: Option<Arc<dyn NetworkPolicyDecider>>,
network_policy_audit_observer: Option<NetworkPolicyAuditObserver>,
) -> Result<
(
HashMap<String, String>,
@@ -307,8 +311,11 @@ async fn prepare_managed_network(
let Some(network_proxy) = network_proxy.cloned() else {
return Ok((env, managed_network.cloned(), None, None));
};
let state = NetworkProxyState::from_remote_launch_config(network_proxy)
let mut state = NetworkProxyState::from_remote_launch_config(network_proxy)
.map_err(|err| invalid_params(format!("invalid network proxy config: {err}")))?;
if let Some(observer) = network_policy_audit_observer {
state.set_policy_audit_observer(observer);
}
let mut builder = NetworkProxy::builder().state(Arc::new(state));
if let Some(network_policy_decider) = network_policy_decider {
builder = builder.policy_decider_arc(network_policy_decider);

View File

@@ -72,6 +72,7 @@ async fn sandbox_request_wraps_native_argv_on_executor() {
HashMap::new(),
Some(&runtime_paths),
/*network_policy_decider*/ None,
/*network_policy_audit_observer*/ None,
)
.await
.expect("prepare sandboxed request");
@@ -140,6 +141,7 @@ async fn sandbox_request_routes_custom_arg0_to_inner_helper() {
HashMap::new(),
Some(&runtime_paths),
/*network_policy_decider*/ None,
/*network_policy_audit_observer*/ None,
)
.await
.expect("prepare sandboxed request");
@@ -203,6 +205,7 @@ async fn sandbox_request_allows_prepared_managed_proxy_port() {
HashMap::new(),
Some(&runtime_paths),
/*network_policy_decider*/ None,
/*network_policy_audit_observer*/ None,
)
.await
.expect("prepare managed-network sandbox request");
@@ -243,6 +246,7 @@ async fn native_request_preserves_native_launch_fields() {
env.clone(),
/*runtime_paths*/ None,
/*network_policy_decider*/ None,
/*network_policy_audit_observer*/ None,
)
.await
.expect("prepare native request");
@@ -295,6 +299,7 @@ async fn native_request_handles_remote_proxy_config_for_platform() {
let prepared = prepare_exec_request(
&params, env, /*runtime_paths*/ None, /*network_policy_decider*/ None,
/*network_policy_audit_observer*/ None,
)
.await
.expect("prepare request with executor-local proxy");
@@ -367,6 +372,7 @@ async fn disabled_remote_proxy_config_is_rejected_before_exporting_ports() {
HashMap::new(),
/*runtime_paths*/ None,
/*network_policy_decider*/ None,
/*network_policy_audit_observer*/ None,
)
.await
.err()
@@ -426,6 +432,7 @@ async fn managed_network_honors_windows_sandbox_level(windows_sandbox_level: Win
HashMap::new(),
Some(&runtime_paths),
/*network_policy_decider*/ None,
/*network_policy_audit_observer*/ None,
)
.await;

View File

@@ -37,6 +37,7 @@ use crate::rpc_server_requests::RpcServerRequestSender;
pub(crate) const SESSION_ALREADY_ATTACHED_ERROR_CODE: i64 = -32010;
const MAX_IN_FLIGHT_REGULAR_CALLS: usize = 1024;
const RESERVED_CLEANUP_CALLS: usize = 1;
const RESERVED_OUTBOUND_CONTROL_MESSAGES: usize = 16;
#[derive(Debug)]
pub(crate) enum RpcCallError {
@@ -153,6 +154,25 @@ impl RpcNotificationSender {
.await
.map_err(|_| internal_error("RPC connection closed while sending notification".into()))
}
pub(crate) fn try_notify<P: Serialize>(&self, method: &str, params: &P) -> bool {
let Ok(permit) = self.outgoing_tx.try_reserve() else {
return false;
};
if self.outgoing_tx.capacity() < RESERVED_OUTBOUND_CONTROL_MESSAGES {
return false;
}
let Ok(params) = serde_json::to_value(params) else {
return false;
};
permit.send(RpcServerOutboundMessage::Notification(
JSONRPCNotification {
method: method.to_string(),
params: Some(params),
},
));
true
}
}
pub(crate) struct RpcRouter<S> {
@@ -793,6 +813,7 @@ mod tests {
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::sync::mpsc;
use tokio::task::JoinSet;
use tokio::time::timeout;
use tracing::Instrument;
@@ -800,12 +821,28 @@ mod tests {
use tracing_subscriber::prelude::*;
use super::MAX_IN_FLIGHT_REGULAR_CALLS;
use super::RESERVED_OUTBOUND_CONTROL_MESSAGES;
use super::RpcCallError;
use super::RpcClient;
use super::RpcNotificationSender;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
use crate::connection::JsonRpcTransport;
#[tokio::test]
async fn best_effort_notifications_preserve_outbound_control_capacity() {
let (outgoing_tx, _outgoing_rx) = mpsc::channel(RESERVED_OUTBOUND_CONTROL_MESSAGES + 2);
let notifications = RpcNotificationSender::new(outgoing_tx);
assert!(notifications.try_notify("network/policyDecision", &serde_json::json!({"n": 1})));
assert!(notifications.try_notify("network/policyDecision", &serde_json::json!({"n": 2})));
assert!(!notifications.try_notify("network/policyDecision", &serde_json::json!({"n": 3})));
notifications
.response(RequestId::Integer(7), serde_json::json!({"ok": true}))
.await
.expect("reserved capacity must remain available for controller responses");
}
async fn read_jsonrpc_line<R>(lines: &mut tokio::io::Lines<BufReader<R>>) -> JSONRPCMessage
where
R: tokio::io::AsyncRead + Unpin,