Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -56,3 +56,6 @@ apps/gateway/target/

# DD skills (installed per-user via pup)
.claude/skills/dd-*

.claude/worktrees

1 change: 1 addition & 0 deletions apps/gateway/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions apps/gateway/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,10 @@ dashmap = "6"
# catastrophic backtracking (no ReDoS). Compiled patterns are cached in a DashMap.
regex = "1"

# Linear-time byte-substring search for body condition matching (attacker-
# controlled request bodies). Already in the tree via `regex`.
memchr = "2"

# Base64 (for gateway auth)
base64 = "0.22"

Expand Down
533 changes: 514 additions & 19 deletions apps/gateway/src/budget.rs

Large diffs are not rendered by default.

696 changes: 690 additions & 6 deletions apps/gateway/src/condition_match.rs

Large diffs are not rendered by default.

7 changes: 5 additions & 2 deletions apps/gateway/src/connect.rs
Original file line number Diff line number Diff line change
Expand Up @@ -745,8 +745,11 @@ impl PolicyEngine {
}
resolved_session_policy = session_policy;
resolved_connection_id = connection_id;
}
if resolved_provider.is_none() {
// Attribute the provider dispatched to `apply_resource_scope`
// to the SAME serving connection as its session_policy and
// connection_id — mirroring the single-connection paths.
// Setting it on the first-with-rules would decouple the
// provider from the scope actually enforced.
resolved_provider = Some(provider);
}
match (earliest_expires_at, token_expires_at) {
Expand Down
75 changes: 75 additions & 0 deletions apps/gateway/src/db.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,14 @@ pub(crate) struct SecretRow {
pub metadata: Option<serde_json::Value>,
}

/// A budget row from the `budgets` table (cost cap on a secret for an org).
#[derive(Debug, FromRow)]
pub(crate) struct BudgetRow {
pub secret_id: String,
pub limit_cents: i32,
pub period: String,
}

/// A user row from the `users` table.
#[derive(Debug, FromRow)]
pub(crate) struct UserRow {
Expand Down Expand Up @@ -344,6 +352,73 @@ pub(crate) async fn find_secrets_by_org(
.context("querying secrets by organization_id")
}

/// Load budgets for the given org and secret ids (host-matched metered LLM
/// secrets). Returns one row per bound secret (0/1 in practice per host).
pub(crate) async fn find_budgets_for_secrets(
pool: &PgPool,
organization_id: &str,
secret_ids: &[String],
) -> Result<Vec<BudgetRow>> {
sqlx::query_as::<_, BudgetRow>(
r#"SELECT secret_id, limit_cents, period
FROM budgets
WHERE organization_id = $1 AND secret_id = ANY($2)"#,
)
.bind(organization_id)
.bind(secret_ids)
.fetch_all(pool)
.await
.context("querying budgets for secrets")
}

/// Read the durable accumulated spend (nano-dollars) for a `(secret, org,
/// period)` window. `None` when the window has no recorded spend yet.
pub(crate) async fn read_budget_spend(
pool: &PgPool,
secret_id: &str,
organization_id: &str,
period: &str,
) -> Result<Option<i64>> {
let row: Option<(i64,)> = sqlx::query_as(
r#"SELECT spent_nanos FROM budget_spends
WHERE secret_id = $1 AND organization_id = $2 AND period = $3"#,
)
.bind(secret_id)
.bind(organization_id)
.bind(period)
.fetch_optional(pool)
.await
.context("reading budget spend")?;
Ok(row.map(|r| r.0))
}

/// Accumulate `delta_nanos` into the durable spend floor for a `(secret, org,
/// period)` window and return the new total. The durable floor is rehydrated
/// into the hot counter on cache miss so a flush can't silently refill a budget.
pub(crate) async fn upsert_budget_spend(
pool: &PgPool,
secret_id: &str,
organization_id: &str,
period: &str,
delta_nanos: i64,
) -> Result<i64> {
let row: (i64,) = sqlx::query_as(
r#"INSERT INTO budget_spends (secret_id, organization_id, period, spent_nanos, updated_at)
VALUES ($1, $2, $3, $4, NOW())
ON CONFLICT (secret_id, organization_id, period)
DO UPDATE SET spent_nanos = budget_spends.spent_nanos + $4, updated_at = NOW()
RETURNING spent_nanos"#,
)
.bind(secret_id)
.bind(organization_id)
.bind(period)
.bind(delta_nanos)
.fetch_one(pool)
.await
.context("upserting budget spend")?;
Ok(row.0)
}

/// Update a secret's encrypted value (used for token refresh).
pub(crate) async fn update_secret_value(
pool: &PgPool,
Expand Down
5 changes: 5 additions & 0 deletions apps/gateway/src/gateway.rs
Original file line number Diff line number Diff line change
Expand Up @@ -985,6 +985,8 @@ async fn handle_http_proxy(
let mut resolved_body_transform: Option<crate::apps::BodyTransform> = None;
// Granular-access policy of the connection that wins injection (if any).
let mut resolved_session_policy: Option<serde_json::Value> = None;
// Provider of that connection — dispatches the resource-scope gate.
let mut resolved_provider: Option<String> = None;
// Id of the connection that wins injection — same attribution law as
// `resolved_session_policy`; unlike `connection_label` below, it MUST be
// threaded (policy decisions bind to it).
Expand Down Expand Up @@ -1014,13 +1016,15 @@ async fn handle_http_proxy(
finalizer,
body_transform,
session_policy,
provider,
connection_id: winning_connection_id,
..
}) => {
app_rules = rules;
resolved_finalizer = finalizer;
resolved_body_transform = body_transform;
resolved_session_policy = session_policy;
resolved_provider = Some(provider);
resolved_connection_id = winning_connection_id;
}
Ok(AppConnectionResult::Ambiguous { connections }) => {
Expand Down Expand Up @@ -1104,6 +1108,7 @@ async fn handle_http_proxy(
finalizer: resolved_finalizer,
body_transform: resolved_body_transform,
claim_token: resolved.claim_token,
provider: resolved_provider,
session_policy: resolved_session_policy,
winning_connection_id: resolved_connection_id,
budget_bindings: resolved.budget_bindings,
Expand Down
99 changes: 68 additions & 31 deletions apps/gateway/src/gateway/forward.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ use crate::apps;
use crate::cache::CacheStore;
use crate::default_interceptions;
use crate::inject;
use crate::policy::{self, PolicyDecision};
use crate::policy::{self, BodyCapture, MatchInput, PolicyDecision};
use crate::policy_engine;

use super::hooks;
Expand Down Expand Up @@ -162,36 +162,43 @@ pub(crate) async fn forward_request(
default_interceptions::match_target(super::strip_port(host), &path, &method)
.filter(|_| content_length_at_most(req.headers(), MAX_DEFAULT_INTERCEPT_BODY));

// Buffer the request body for condition matching, when the request guard needs
// to inspect it (e.g. Dropbox folder scoping reads the JSON body), or for a
// matched default interception. In OSS, both predicates return false → zero
// overhead unless a default interception matched.
let (condition_buffer, req) = if crate::policy_engine::needs_body_buffer(&rules.policy_rules_v2)
|| hooks::needs_request_body(rules, host, method.as_str(), &path)
{
let (parts, incoming) = req.into_parts();
let (buf, fwd_body) =
crate::condition_match::prepare_body(incoming, method.as_str(), &url).await?;
(buf, hyper::Request::from_parts(parts, fwd_body))
} else if default_target.is_some() {
// OSS-safe: fully buffer the known-small body, keeping the bytes for both
// the interception check and (if it declines) forwarding.
let (parts, incoming) = req.into_parts();
let bytes = incoming
.collect()
.await
.context("buffering request body for default interception")?
.to_bytes();
let req = hyper::Request::from_parts(parts, reqwest::Body::from(bytes.clone()));
(Some(bytes.to_vec()), req)
} else {
(None, req.map(reqwest::Body::wrap))
};
// Buffer the request body for condition matching, when a body-conditioned
// rule could govern this host (`needs_body_buffer` is host-scoped), when the
// request guard needs to inspect it (e.g. Dropbox folder scoping reads the
// JSON body), or for a matched default interception. Unconditioned traffic
// keeps streaming → zero overhead.
let (capture, req) =
if crate::policy_engine::needs_body_buffer(&rules.policy_rules_v2, policy_host)
|| crate::policy_engine::needs_scope_body(
rules.provider.as_deref().unwrap_or(""),
host,
rules.session_policy.as_ref(),
)
|| hooks::needs_request_body(rules, host, method.as_str(), &path)
{
let (parts, incoming) = req.into_parts();
let (capture, fwd_body) =
crate::condition_match::prepare_body(incoming, method.as_str(), &url).await?;
(capture, hyper::Request::from_parts(parts, fwd_body))
} else if default_target.is_some() {
// OSS-safe: fully buffer the known-small body, keeping the bytes for both
// the interception check and (if it declines) forwarding.
let (parts, incoming) = req.into_parts();
let bytes = incoming
.collect()
.await
.context("buffering request body for default interception")?
.to_bytes();
let req = hyper::Request::from_parts(parts, reqwest::Body::from(bytes.clone()));
(BodyCapture::Full(bytes.to_vec()), req)
} else {
(BodyCapture::None, req.map(reqwest::Body::wrap))
};

// Answer a matched default interception before any forwarding. A handler that
// declines (e.g. a real refresh token) falls through to normal forwarding.
if let Some(target) = default_target {
if let Some(synth) = target.handle(condition_buffer.as_deref().unwrap_or(&[])) {
if let Some(synth) = target.handle(capture.bytes().unwrap_or(&[])) {
info!(method = %method, url = %url, "default interception — serving synthetic response");
return Ok(response::json(synth.status, synth.body));
}
Expand All @@ -218,6 +225,12 @@ pub(crate) async fn forward_request(
));
}

// The per-request condition input: the captured body (only a FULL capture is
// matchable — a truncated one fails closed) plus the pre-injection request
// headers. Borrows `req`; its last use is the `evaluate` call below, which NLL
// releases before `req.into_parts()` later.
let match_input = MatchInput::from_capture(&capture, req.headers());

// The first-match engine over the published `policy_rules_v2` is authoritative.
// `policy_host` is the pre-rewrite rule-match host; `is_llm_host(host)` is the
// effective host for the deny-default carve.
Expand All @@ -226,7 +239,7 @@ pub(crate) async fn forward_request(
policy_host,
method.as_str(),
&path,
condition_buffer.as_deref(),
&match_input,
has_injections,
policy::is_llm_host(host),
rules.winning_connection_id.as_deref(),
Expand All @@ -235,6 +248,30 @@ pub(crate) async fn forward_request(
)
.await;

// ── Granular resource-scope gate (Tier 3b) ────────────────────────────────
// Tighten the engine's decision by the winning connection's granular scope
// (GitHub repositories / Dropbox folders). Run unconditionally on the final
// decision from EITHER engine — session policy is a property of the
// connection, not the rule generation, so it must enforce on legacy /
// pre-cutover projects too. A monotone tightening: an allow-family verdict
// for an out-of-scope or indeterminate resource becomes `Blocked`; an
// existing block is untouched. A scope block is authored by no rule, so it
// drops the matched-rule attribution.
let (decision, matched_rule) = {
let (decision, scope_blocked) = policy_engine::apply_resource_scope(
decision,
rules.provider.as_deref().unwrap_or(""),
host,
rules.session_policy.as_ref(),
&path,
&match_input,
);
if scope_blocked {
warn!(method = %method, url = %url, "BLOCKED by resource scope");
}
(decision, if scope_blocked { None } else { matched_rule })
};

// ── Early return for block / rate-limit / default-deny (no body needed) ───
match &decision {
PolicyDecision::BlockedByDefaultPolicy => {
Expand Down Expand Up @@ -346,7 +383,7 @@ pub(crate) async fn forward_request(
method.as_str(),
&path,
&headers,
condition_buffer.as_deref(),
capture.bytes(),
)
.await
{
Expand Down Expand Up @@ -378,8 +415,8 @@ pub(crate) async fn forward_request(
// Peek a bounded prefix of the body for the summary + preview, then
// build the forwarding body. If condition buffering already captured
// the body, reuse that buffer instead of peeking the stream again.
let (summary_bytes, fwd_body): (Cow<'_, [u8]>, reqwest::Body) = if let Some(ref buf) =
condition_buffer
let (summary_bytes, fwd_body): (Cow<'_, [u8]>, reqwest::Body) = if let Some(buf) =
capture.bytes()
{
// Body already buffered for condition matching — borrow its prefix
// for the summary instead of copying it again.
Expand Down
Loading