Skip to content

Commit f027d07

Browse files
lavkushryclaude
andcommitted
feat(tool-broker): extract connector execution into standalone aegis-tool-broker binary
Phase 1 of the tool-broker extraction: the gateway process no longer links aegis-tool-broker-connectors (removed from src/Cargo.toml production deps, kept as a dev-dependency for the one deliberate mock-server test exception). It never resolves a real provider credential or runs a shell command itself anymore. POST /v1/broker/execute now calls the new standalone aegis-tool-broker binary over HTTP with a required service-to-service bearer token (AEGIS_TOOL_BROKER_API_TOKEN); with no broker configured it fails closed with 501, matching the OIDC/signed-policy-bundle precedent. The gateway is unchanged for authorization, hash-bound approval consumption, and receipt emission -- only the connector-execution step moved process. Ships with a Helm chart (helm/aegis-tool-broker/) and a docker-compose.full.yml service, both wired for local/dev use. Known Phase-1 gap (documented in Tool_Broker.md's Honest Scope, not glossed over): the gateway consumes the approval before the broker HTTP call, so a broker-unreachable failure burns the approval without executing the action. Covered by execute_fails_closed_when_the_broker_is_unreachable. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 41bf4e0 commit f027d07

29 files changed

Lines changed: 1713 additions & 121 deletions

Cargo.lock

Lines changed: 22 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ members = [
1515
"bins/aegis-cage-runner",
1616
"bins/aegis-egress-proxy",
1717
"bins/aegis-llm-gateway",
18+
"bins/aegis-tool-broker",
1819
]
1920
# src/fuzz has its own cargo-fuzz-managed build (ASan/sancov instrumentation,
2021
# a separate target dir) and must never be built by `cargo build/test

bins/aegis-tool-broker/Cargo.toml

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
[package]
2+
name = "aegis-tool-broker"
3+
version = "0.1.0"
4+
edition = "2021"
5+
description = "AegisAgent standalone tool-broker execution engine (Phase 1 extraction): resolves credentials and runs connectors in its own process, so the gateway never links this crate's connectors or holds a real provider credential."
6+
7+
[[bin]]
8+
name = "aegis-tool-broker"
9+
path = "src/main.rs"
10+
11+
[dependencies]
12+
aegis-tool-broker-core = { path = "../../lib/tool-broker-core" }
13+
aegis-tool-broker-connectors = { path = "../../lib/tool-broker-connectors" }
14+
serde = { workspace = true }
15+
serde_json = { workspace = true }
16+
tokio = { workspace = true }
17+
tracing = { workspace = true }
18+
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
19+
thiserror = { workspace = true }
20+
axum = { workspace = true }
21+
clap = { version = "4.6", features = ["derive", "env"] }
22+
reqwest = { version = "0.12.28", default-features = false, features = ["rustls-tls", "json"] }
23+
sha2 = { workspace = true }
24+
hex = "=0.4.3"
25+
26+
[dev-dependencies]
27+
tempfile = "3.10"
28+
tower = { version = "0.4", features = ["util"] }
29+
async-trait = "0.1.89"

bins/aegis-tool-broker/Dockerfile

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# aegis-tool-broker: standalone tool-broker execution engine (Phase 1
2+
# extraction). A separate binary/image from the gateway — the gateway no
3+
# longer links aegis-tool-broker-connectors at all, so this is the only
4+
# process that resolves a real provider credential or executes a connector.
5+
FROM rust:1.96-bookworm AS builder
6+
WORKDIR /build
7+
8+
# Network resilience for the crate fetch, mirroring src/Dockerfile.
9+
ENV CARGO_NET_RETRY=10
10+
ENV CARGO_HTTP_MULTIPLEXING=false
11+
12+
# #1160: repo-root .cargo/config.toml sets `--cfg tokio_unstable` — copy it
13+
# so this build (run from /build, mirroring CI's repo-root invocation)
14+
# resolves the same workspace config the gateway's Dockerfile does.
15+
COPY .cargo ./.cargo
16+
COPY Cargo.toml Cargo.lock ./
17+
COPY src ./src
18+
COPY lib ./lib
19+
COPY bins ./bins
20+
RUN cargo build --release -p aegis-tool-broker --bin aegis-tool-broker
21+
RUN strip /build/target/release/aegis-tool-broker
22+
23+
# distroless/cc-debian12 (no shell, no package manager), matching the
24+
# gateway image's hardening posture. ca-certificates copied explicitly — the
25+
# GithubConnector's real mode and HttpConnector both make outbound HTTPS
26+
# calls and need a real CA store to verify them.
27+
FROM gcr.io/distroless/cc-debian12:nonroot
28+
WORKDIR /app
29+
COPY --from=builder /etc/ssl/certs /etc/ssl/certs
30+
COPY --from=builder /build/target/release/aegis-tool-broker /usr/local/bin/aegis-tool-broker
31+
32+
ENV RUST_LOG=info
33+
34+
# Unlike aegis-egress-proxy (a CONNECT proxy with no HTTP health route),
35+
# this binary is a real HTTP JSON API — `--healthcheck` self-GETs /livez and
36+
# exits 0/1, used here instead of pulling curl/wget into the distroless
37+
# final image.
38+
HEALTHCHECK --interval=10s --timeout=3s --start-period=10s --retries=6 \
39+
CMD ["/usr/local/bin/aegis-tool-broker", "--healthcheck"]
40+
41+
ENTRYPOINT ["/usr/local/bin/aegis-tool-broker"]
42+
CMD ["--listen", "0.0.0.0:8899"]

bins/aegis-tool-broker/src/auth.rs

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
//! Bearer-token auth for the one privileged route (`POST /v1/execute`).
2+
//!
3+
//! Unlike the gateway's optional admin key (which treats a loopback bind as
4+
//! evidence of safe operator intent), this binary's entire surface *is* the
5+
//! privileged action it guards — there is no "safe to leave open" mode, so
6+
//! the token is a required CLI/env field (`Cli::api_token` has no
7+
//! `default_value`) and every `/v1/execute` call is checked unconditionally.
8+
9+
use axum::extract::State;
10+
use axum::http::{header, Request, StatusCode};
11+
use axum::middleware::Next;
12+
use axum::response::{IntoResponse, Response};
13+
use std::sync::Arc;
14+
15+
/// SHA-256 digest equality rather than a raw `==` — mirrors the gateway's
16+
/// `admin_decision` rationale: a timing side-channel on a raw string
17+
/// comparison can leak the configured token one byte at a time.
18+
pub fn tokens_match(configured: &str, provided: &str) -> bool {
19+
fn hash(token: &str) -> String {
20+
use sha2::{Digest, Sha256};
21+
hex::encode(Sha256::digest(token.as_bytes()))
22+
}
23+
hash(configured) == hash(provided)
24+
}
25+
26+
pub async fn require_bearer_token(
27+
State(configured_token): State<Arc<String>>,
28+
request: Request<axum::body::Body>,
29+
next: Next,
30+
) -> Response {
31+
let provided = request
32+
.headers()
33+
.get(header::AUTHORIZATION)
34+
.and_then(|h| h.to_str().ok())
35+
.and_then(|v| v.strip_prefix("Bearer "));
36+
37+
match provided {
38+
Some(token) if tokens_match(&configured_token, token) => next.run(request).await,
39+
_ => (StatusCode::UNAUTHORIZED, "invalid or missing bearer token").into_response(),
40+
}
41+
}
42+
43+
#[cfg(test)]
44+
mod tests {
45+
use super::*;
46+
47+
#[test]
48+
fn matching_tokens_are_equal() {
49+
assert!(tokens_match("secret-token", "secret-token"));
50+
}
51+
52+
#[test]
53+
fn mismatched_tokens_are_not_equal() {
54+
assert!(!tokens_match("secret-token", "wrong-token"));
55+
}
56+
57+
#[test]
58+
fn empty_provided_token_never_matches_a_configured_one() {
59+
assert!(!tokens_match("secret-token", ""));
60+
}
61+
62+
#[test]
63+
fn two_empty_tokens_are_equal_but_the_server_never_configures_an_empty_one() {
64+
// Documents the invariant this relies on: Cli::api_token has no
65+
// default_value, so an empty configured token is unreachable in
66+
// production -- clap itself refuses to start the process.
67+
assert!(tokens_match("", ""));
68+
}
69+
}

bins/aegis-tool-broker/src/dto.rs

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
//! Wire types for `POST /v1/execute`. Deliberately a separate DTO from
2+
//! `aegis_tool_broker_connectors`' own types (matching this codebase's
3+
//! existing convention for gateway<->satellite contracts, e.g.
4+
//! `routes/egress.rs`'s `EgressCheckRequest` vs. `aegis-egress-proxy`'s own
5+
//! `GatewayCheckRequest` — two independently-defined, structurally-matching
6+
//! types, not a shared crate) so this binary's public HTTP contract can
7+
//! evolve without dragging the gateway's `Cargo.toml` along.
8+
9+
use aegis_tool_broker_core::BrokerAction;
10+
use serde::{Deserialize, Serialize};
11+
12+
#[derive(Debug, Deserialize)]
13+
pub struct ExecuteRequest {
14+
pub tool_name: String,
15+
pub connector_type: String,
16+
/// Opaque reference (e.g. `env:GITHUB_TOKEN`), or `None` for a tool
17+
/// registered without a credential. Never a secret value.
18+
pub credential_ref: Option<String>,
19+
/// `active` | `disabled` — from the gateway's `broker_tools` row.
20+
pub tool_status: String,
21+
pub action: BrokerAction,
22+
pub consumed_approval: Option<ConsumedApprovalWire>,
23+
}
24+
25+
#[derive(Debug, Deserialize)]
26+
pub struct ConsumedApprovalWire {
27+
pub approval_id: String,
28+
pub action_hash: String,
29+
}
30+
31+
#[derive(Debug, Serialize)]
32+
pub struct ExecuteResponse {
33+
pub output: serde_json::Value,
34+
}
35+
36+
#[derive(Debug, Serialize)]
37+
pub struct ErrorBody {
38+
pub error: String,
39+
pub kind: &'static str,
40+
}
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
//! Maps `aegis_tool_broker_connectors::ExecuteError` to an HTTP response.
2+
//! Status-code mapping is unchanged from the gateway's own (pre-extraction)
3+
//! `broker_execute_error_response` in `routes/broker.rs` — only relocated.
4+
5+
use crate::dto::ErrorBody;
6+
use aegis_tool_broker_connectors::ExecuteError;
7+
use axum::http::StatusCode;
8+
use axum::response::{IntoResponse, Json, Response};
9+
10+
pub fn execute_error_response(e: ExecuteError) -> Response {
11+
let (status, kind) = match &e {
12+
ExecuteError::ToolNotActive { .. } => (StatusCode::FORBIDDEN, "tool_not_active"),
13+
ExecuteError::UnknownConnectorType { .. } => {
14+
(StatusCode::NOT_IMPLEMENTED, "unknown_connector_type")
15+
}
16+
ExecuteError::ApprovalRequired => (StatusCode::FORBIDDEN, "approval_required"),
17+
ExecuteError::ApprovalActionMismatch { .. } => {
18+
(StatusCode::FORBIDDEN, "approval_action_mismatch")
19+
}
20+
ExecuteError::Credential(_) => (StatusCode::SERVICE_UNAVAILABLE, "credential"),
21+
ExecuteError::Connector(_) => (StatusCode::SERVICE_UNAVAILABLE, "connector"),
22+
};
23+
(
24+
status,
25+
Json(ErrorBody {
26+
error: e.to_string(),
27+
kind,
28+
}),
29+
)
30+
.into_response()
31+
}
Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
//! Builds the [`BrokerExecutor`] this binary serves. Moved verbatim (Phase 1
2+
//! extraction) from the gateway's `routes/broker.rs::default_broker_executor`.
3+
4+
use std::sync::Arc;
5+
6+
use aegis_tool_broker_connectors::{
7+
BrokerExecutor, ConnectorRegistry, FilesystemConnector, GithubConnector, GithubMode,
8+
HttpConnector, ShellConnector,
9+
};
10+
use aegis_tool_broker_core::EnvCredentialResolver;
11+
use tracing::error;
12+
13+
/// Builds the [`BrokerExecutor`] this process serves. GitHub runs in mock
14+
/// mode unless `AEGIS_GITHUB_API_BASE` is set (real mode); `HttpConnector`
15+
/// is always registered (HTTPS-only). The filesystem and shell connectors
16+
/// are opt-in via `AEGIS_BROKER_WORKSPACE` — with no configured workspace
17+
/// there is nothing safe to scope them to, so they're simply absent from
18+
/// the registry (an execute against `filesystem`/`shell` then fails closed
19+
/// with `UnknownConnectorType`, not a wide-open default).
20+
pub fn build_broker_executor() -> Arc<BrokerExecutor> {
21+
let github_mode = match std::env::var("AEGIS_GITHUB_API_BASE") {
22+
Ok(base_url) if !base_url.trim().is_empty() => GithubMode::Real {
23+
base_url: base_url.trim_end_matches('/').to_string(),
24+
},
25+
_ => GithubMode::Mock,
26+
};
27+
let mut registry = ConnectorRegistry::default()
28+
.register(Arc::new(GithubConnector::new(github_mode)))
29+
.register(Arc::new(HttpConnector::new()));
30+
31+
if let Ok(workspace) = std::env::var("AEGIS_BROKER_WORKSPACE") {
32+
if !workspace.trim().is_empty() {
33+
match FilesystemConnector::new(&workspace) {
34+
Ok(fs) => registry = registry.register(Arc::new(fs)),
35+
Err(e) => error!(
36+
"AEGIS_BROKER_WORKSPACE {:?} unusable for filesystem connector: {}",
37+
workspace, e
38+
),
39+
}
40+
match ShellConnector::new(&workspace) {
41+
Ok(shell) => registry = registry.register(Arc::new(shell)),
42+
Err(e) => error!(
43+
"AEGIS_BROKER_WORKSPACE {:?} unusable for shell connector: {}",
44+
workspace, e
45+
),
46+
}
47+
}
48+
}
49+
50+
Arc::new(BrokerExecutor::new(
51+
registry,
52+
Arc::new(EnvCredentialResolver),
53+
))
54+
}
55+
56+
#[cfg(test)]
57+
mod tests {
58+
use super::*;
59+
60+
#[tokio::test]
61+
async fn github_defaults_to_mock_mode_when_api_base_unset() {
62+
std::env::remove_var("AEGIS_GITHUB_API_BASE");
63+
std::env::remove_var("AEGIS_BROKER_WORKSPACE");
64+
let executor = build_broker_executor();
65+
// Mock mode never makes a real network call; a benign read must
66+
// succeed synchronously without any credential configured.
67+
let output = executor
68+
.execute(
69+
&aegis_tool_broker_connectors::BrokerToolBinding {
70+
tool_name: "gh".to_string(),
71+
connector_type: "github".to_string(),
72+
credential_ref: None,
73+
status: "active".to_string(),
74+
},
75+
&aegis_tool_broker_core::BrokerAction {
76+
tool: "gh".to_string(),
77+
action: "read".to_string(),
78+
resource: None,
79+
mutates_state: false,
80+
parameters: serde_json::json!({"path": "/repos/acme/api/issues"}),
81+
},
82+
None,
83+
)
84+
.await;
85+
assert!(output.is_ok());
86+
}
87+
88+
#[tokio::test]
89+
async fn filesystem_and_shell_are_absent_without_a_configured_workspace() {
90+
std::env::remove_var("AEGIS_BROKER_WORKSPACE");
91+
let executor = build_broker_executor();
92+
let err = executor
93+
.execute(
94+
&aegis_tool_broker_connectors::BrokerToolBinding {
95+
tool_name: "fs".to_string(),
96+
connector_type: "filesystem".to_string(),
97+
credential_ref: None,
98+
status: "active".to_string(),
99+
},
100+
&aegis_tool_broker_core::BrokerAction {
101+
tool: "fs".to_string(),
102+
action: "read".to_string(),
103+
resource: None,
104+
mutates_state: false,
105+
parameters: serde_json::json!({}),
106+
},
107+
None,
108+
)
109+
.await
110+
.expect_err("filesystem must be unregistered without AEGIS_BROKER_WORKSPACE");
111+
assert!(matches!(
112+
err,
113+
aegis_tool_broker_connectors::ExecuteError::UnknownConnectorType { .. }
114+
));
115+
}
116+
117+
#[tokio::test]
118+
async fn filesystem_and_shell_are_registered_with_a_configured_workspace() {
119+
let dir = tempfile::tempdir().unwrap();
120+
std::env::set_var("AEGIS_BROKER_WORKSPACE", dir.path());
121+
let executor = build_broker_executor();
122+
let output = executor
123+
.execute(
124+
&aegis_tool_broker_connectors::BrokerToolBinding {
125+
tool_name: "shell".to_string(),
126+
connector_type: "shell".to_string(),
127+
credential_ref: None,
128+
status: "active".to_string(),
129+
},
130+
&aegis_tool_broker_core::BrokerAction {
131+
tool: "shell".to_string(),
132+
action: "run".to_string(),
133+
resource: None,
134+
mutates_state: false,
135+
parameters: serde_json::json!({"command": ["/usr/bin/env"]}),
136+
},
137+
None,
138+
)
139+
.await;
140+
assert!(output.is_ok());
141+
std::env::remove_var("AEGIS_BROKER_WORKSPACE");
142+
}
143+
}

0 commit comments

Comments
 (0)