diff --git a/crates/aisix-proxy/AGENTS.md b/crates/aisix-proxy/AGENTS.md index fcf104b1..d9cf25a3 100644 --- a/crates/aisix-proxy/AGENTS.md +++ b/crates/aisix-proxy/AGENTS.md @@ -28,18 +28,31 @@ onto whatever the executor runs next on that thread. written against it silently never runs for group traffic, and nothing errors — requests keep succeeding on a target that should have been excluded. -Decide, and encode the decision at the call site: - -- **Binds each target** (anything protecting the upstream behind it — rate limits, - cooldown, health, timeouts): resolve it from the attempt model *inside* the - dispatch loop, in all four group-capable endpoints (chat, messages, - count_tokens, responses) and in both the streaming and non-streaming branches. - A limit-shaped gate should skip the target and let dispatch continue rather than - failing the whole request — see `quota::reserve_routing_target`. -- **Binds the requested entry** (anything scoped to the alias the caller named — - `allowed_cidrs`, guardrail attachment): keep it pre-dispatch, and say so in the - user-facing docs, because the group/member split is otherwise invisible. - -A reservation-shaped gate additionally must not double-charge: `reserve_routing_target` -returns `None` for non-routing dispatch, whose model layers the pre-dispatch -`quota::enforce*` already reserved. +**The default is that a per-model gate binds each target.** Anything an operator +configures ON a model — rate limits, `allowed_cidrs`, cooldown, health, timeouts — +is a statement about that model, and reaching it through a group must not strip it. +Two gates are deliberately entry-scoped instead: guardrail attachment (resolved +from `model_id` before dispatch, by design) and the group's own copy of any of the +above. Anything else that only checks `model_entry` / `virtual_entry` is a bug. + +Two shapes, both already implemented — copy the nearest one: + +- **Filter the candidate set** (static per-caller predicates like `allowed_cidrs`): + drop ineligible targets in `routing::resolve_attempt_models` *before* the strategy + picks, so `max_fallbacks` budgets attempts across reachable targets and a + metric-based strategy ranks only those. Empty result → the gate's own error. + Do NOT fold these into `filter_attempt_models`: its + `when_all_unavailable: try_anyway` policy hands back the unfiltered list, which + would defeat an allowlist. See `routing::targets_allowed_for_ip`. +- **Check per attempt** (dynamic/stateful gates like a rate-limit reservation): + resolve from the attempt model *inside* the dispatch loop, in all four + group-capable endpoints (chat, messages, count_tokens, responses) and in both the + streaming and non-streaming branches; skip the target and continue rather than + failing the whole request. See `quota::reserve_routing_target`, which also shows + the non-double-charge rule: it returns `None` for non-routing dispatch, whose + model layers the pre-dispatch `quota::enforce*` already reserved. + +Whichever shape, the group's own gate stays enforced pre-dispatch — the two tiers +are additive, not either/or — and a caller-visible rejection must keep the +direct-model envelope (`ModelIpRestricted` names no model and no CIDR), so a group +never becomes a probe for which members exist. diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index abe52bc2..a84c3f53 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1029,6 +1029,7 @@ async fn dispatch( .as_deref() .unwrap_or(auth.entry.id.as_str()), ), + source_ip: &client.source_ip, }, ) .map_err(&with_model)?; diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index a6f19d1f..1c6151c7 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -185,6 +185,7 @@ async fn dispatch( .as_deref() .unwrap_or(auth.entry.id.as_str()), ), + source_ip: &client.source_ip, }, )?; let retry_on_429 = model_entry diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 9e229e93..2dc17252 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -638,6 +638,7 @@ async fn dispatch( .as_deref() .unwrap_or(auth.entry.id.as_str()), ), + source_ip: &client.source_ip, }, )?; diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 686cd340..efc0a049 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -517,6 +517,7 @@ async fn dispatch( .as_deref() .unwrap_or(auth.entry.id.as_str()), ), + source_ip: &client.source_ip, }, )?; let retry_on_429 = model_entry diff --git a/crates/aisix-proxy/src/routing.rs b/crates/aisix-proxy/src/routing.rs index 9e84cd28..37abc150 100644 --- a/crates/aisix-proxy/src/routing.rs +++ b/crates/aisix-proxy/src/routing.rs @@ -424,14 +424,46 @@ pub(crate) fn filter_attempt_models( } /// Per-request routing inputs threaded into [`resolve_attempt_models`]: the -/// tags that gate tag/metadata routing and the stability key for sticky -/// (A/B / canary) weighted selection. Tags come from request headers; the -/// stability key is the routing-key header when present, otherwise the +/// tags that gate tag/metadata routing, the stability key for sticky +/// (A/B / canary) weighted selection, and the caller's resolved source IP +/// for the per-target client-IP allowlist. Tags come from request headers; +/// the stability key is the routing-key header when present, otherwise the /// caller's API key id. +/// +/// `source_ip` defaults to the empty string, which +/// [`aisix_core::Model::ip_allowed`] treats as "not in range" — so a caller +/// that forgets to thread it fails closed on restricted targets rather than +/// silently disabling the allowlist. #[derive(Clone, Copy, Default)] pub(crate) struct RoutingRequest<'a> { pub tags: &'a [String], pub stability_key: Option<&'a str>, + pub source_ip: &'a str, +} + +/// Drop the targets whose own `allowed_cidrs` excludes `source_ip`. +/// +/// Deliberately NOT folded into [`filter_attempt_models`]: that filter's +/// `when_all_unavailable: try_anyway` policy hands back the *unfiltered* +/// candidate list, which would send a request to a target the operator just +/// declared off-limits for this caller. An allowlist has no "try anyway". +fn targets_allowed_for_ip( + snapshot: &AisixSnapshot, + targets: Vec, + source_ip: &str, +) -> Vec { + targets + .into_iter() + .filter(|t| { + // An unresolvable name is left in place so the resolution loop + // below still reports it as a config error, rather than being + // silently swallowed here as an IP rejection. + snapshot + .models + .get_by_name(&t.model) + .is_none_or(|entry| entry.value.ip_allowed(source_ip)) + }) + .collect() } /// Resolve the ordered list of concrete Models a request will attempt. @@ -469,6 +501,19 @@ pub(crate) fn resolve_attempt_models( req.tags ))); } + // Client-IP pre-filter (AISIX-Cloud#1087 follow-up): a target whose own + // `allowed_cidrs` excludes this caller is not a candidate. Applied BEFORE + // the strategy picks, so `max_fallbacks` budgets attempts across the + // targets this caller may actually reach, and a metric-based strategy + // ranks only those. The group's own `allowed_cidrs` is separately enforced + // pre-dispatch by `dispatch::check_ip_access`; this adds the member tier + // that a group previously bypassed entirely. + let eligible = targets_allowed_for_ip(snapshot, eligible, req.source_ip); + if eligible.is_empty() { + // Report the name the caller asked for, not the excluded members — + // matching `ModelForbidden`, and without disclosing group internals. + return Err(ProxyError::ModelIpRestricted(virtual_name.to_string())); + } let filtered_routing = Routing { targets: eligible, ..routing.clone() @@ -656,6 +701,88 @@ mod tests { ); } + // ───────────────── per-target client-IP allowlist ───────────────── + + fn ip_snapshot(models: &[(&str, Option>)]) -> AisixSnapshot { + let table = aisix_core::snapshot::ResourceTable::default(); + for (i, (name, cidrs)) in models.iter().enumerate() { + let model: Model = serde_json::from_value(serde_json::json!({ + "display_name": name, + "provider": "openai", + "model_name": "up", + "provider_key_id": "pk-1", + "allowed_cidrs": cidrs, + })) + .unwrap(); + table.insert(aisix_core::ResourceEntry::new(format!("m-{i}"), model, 1)); + } + AisixSnapshot { + models: table, + ..Default::default() + } + } + + #[test] + fn ip_filter_drops_only_the_out_of_range_target() { + let snap = ip_snapshot(&[("restricted", Some(vec!["10.0.0.0/8"])), ("open", None)]); + let targets = vec![tagged("restricted", &[]), tagged("open", &[])]; + + // In range → both stay candidates. + assert_eq!( + model_names(&targets_allowed_for_ip(&snap, targets.clone(), "10.1.2.3")), + vec!["restricted", "open"] + ); + // Out of range → the restricted member drops out, the group still serves. + assert_eq!( + model_names(&targets_allowed_for_ip(&snap, targets, "8.8.8.8")), + vec!["open"] + ); + } + + #[test] + fn ip_filter_empties_when_every_target_excludes_the_caller() { + // The caller turns an empty result into a 403 rather than dispatching. + let snap = ip_snapshot(&[ + ("a", Some(vec!["10.0.0.0/8"])), + ("b", Some(vec!["192.168.0.0/16"])), + ]); + let targets = vec![tagged("a", &[]), tagged("b", &[])]; + assert!(targets_allowed_for_ip(&snap, targets, "8.8.8.8").is_empty()); + } + + #[test] + fn ip_filter_fails_closed_on_an_unattributable_source_ip() { + // Mirrors `Model::ip_allowed`: an empty/unparseable IP can never + // satisfy a configured allowlist, so a request whose peer address + // was lost must not reach a restricted target. + let snap = ip_snapshot(&[("restricted", Some(vec!["10.0.0.0/8"]))]); + let targets = vec![tagged("restricted", &[])]; + assert!(targets_allowed_for_ip(&snap, targets, "").is_empty()); + } + + #[test] + fn ip_filter_keeps_unresolvable_names_for_the_config_error_path() { + // A target naming a Model that isn't in the snapshot must surface as + // the existing "does not resolve to a Model" config error, not be + // silently swallowed here as an IP rejection. + let snap = ip_snapshot(&[("known", None)]); + let targets = vec![tagged("ghost", &[])]; + assert_eq!( + model_names(&targets_allowed_for_ip(&snap, targets, "8.8.8.8")), + vec!["ghost"] + ); + } + + #[test] + fn ip_filter_is_a_noop_when_no_target_restricts() { + let snap = ip_snapshot(&[("a", None), ("b", None)]); + let targets = vec![tagged("a", &[]), tagged("b", &[])]; + assert_eq!( + model_names(&targets_allowed_for_ip(&snap, targets, "8.8.8.8")), + vec!["a", "b"] + ); + } + #[test] fn eligible_tagged_no_match_no_default_is_empty() { // The caller turns an empty result into a "no target matches tags" error. diff --git a/tests/e2e/src/cases/model-ip-restriction-e2e.test.ts b/tests/e2e/src/cases/model-ip-restriction-e2e.test.ts index 9e14ad04..bc1cb8a1 100644 --- a/tests/e2e/src/cases/model-ip-restriction-e2e.test.ts +++ b/tests/e2e/src/cases/model-ip-restriction-e2e.test.ts @@ -24,6 +24,14 @@ import { // AC-1 block — out-of-range XFF → 403 `code: "ip_restricted"`, upstream // untouched (rejected pre-dispatch). // AC-2 isolation — same external IP: restricted model 403, unrestricted 200. +// +// AISIX-Cloud#1087 follow-up: the same per-model allowlist must hold when +// the model is reached as a Model Group target. Pre-fix a group bypassed +// its members' `allowed_cidrs` entirely — only the named alias was checked +// — so adding a restricted model to a group silently published it to every +// caller. Post-fix an out-of-range target drops out of the candidate set +// (the group still serves from the remaining targets), and a group whose +// every target excludes the caller returns 403. const CALLER_PLAINTEXT = "sk-model-ip-caller"; const CALLER_KEY_HASH = createHash("sha256") @@ -32,6 +40,17 @@ const CALLER_KEY_HASH = createHash("sha256") const RESTRICTED_MODEL = "ip-restricted-model"; const OPEN_MODEL = "ip-open-model"; +// Groups for the #1087 follow-up: one that can fall back to an +// unrestricted member, one whose every member excludes the caller. +const MIXED_GROUP = "ip-mixed-group"; +const ALL_RESTRICTED_GROUP = "ip-all-restricted-group"; +const RESTRICTED_MODEL_2 = "ip-restricted-model-2"; +// The mixed group's open member gets its OWN upstream returning a marker +// string. Sharing the restricted member's upstream would make the test pass +// pre-fix too: both members answer 200, so only the response content can +// prove WHICH member the group actually dispatched to. +const GROUP_OPEN_MODEL = "ip-group-open-model"; +const GROUP_OPEN_MARKER = "served-by-open-member"; async function chat( app: SpawnedApp, @@ -55,6 +74,7 @@ async function chat( describe("model IP restriction e2e (#557): allowed_cidrs gate before upstream", () => { let app: SpawnedApp | undefined; let upstream: OpenAiUpstream | undefined; + let groupOpenUpstream: OpenAiUpstream | undefined; let etcdReachable = false; beforeAll(async () => { @@ -63,6 +83,22 @@ describe("model IP restriction e2e (#557): allowed_cidrs gate before upstream", if (!etcdReachable) return; upstream = await startOpenAiUpstream(); + groupOpenUpstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-group-open", + object: "chat.completion", + created: 0, + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: GROUP_OPEN_MARKER }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }); // 127.0.0.1 (the loopback e2e client) is the trusted proxy, so the // gateway honours `x-forwarded-for` and treats the forwarded value as // the real client IP. @@ -91,15 +127,59 @@ describe("model IP restriction e2e (#557): allowed_cidrs gate before upstream", model_name: "gpt-4o-mini", provider_key_id: pk.id, }); + // Second restricted model on a range the test client is never in, so a + // group of two restricted members has no reachable target. + await seed.createModel({ + display_name: RESTRICTED_MODEL_2, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + allowed_cidrs: ["192.168.0.0/16"], + }); + const groupOpenPk = await seed.createProviderKey({ + display_name: "group-open-pk", + secret: "sk-mock", + api_base: `${groupOpenUpstream!.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: GROUP_OPEN_MODEL, + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: groupOpenPk.id, + }); + // Restricted member FIRST so an out-of-range caller only succeeds if the + // member's allowlist actually removed it from the candidate set. + await seed.createModel({ + display_name: MIXED_GROUP, + routing: { + strategy: "failover", + targets: [{ model: RESTRICTED_MODEL }, { model: GROUP_OPEN_MODEL }], + }, + }); + await seed.createModel({ + display_name: ALL_RESTRICTED_GROUP, + routing: { + strategy: "failover", + targets: [{ model: RESTRICTED_MODEL }, { model: RESTRICTED_MODEL_2 }], + }, + }); await seed.createApiKey({ key_hash: CALLER_KEY_HASH, - allowed_models: [RESTRICTED_MODEL, OPEN_MODEL], + allowed_models: [ + RESTRICTED_MODEL, + RESTRICTED_MODEL_2, + OPEN_MODEL, + GROUP_OPEN_MODEL, + MIXED_GROUP, + ALL_RESTRICTED_GROUP, + ], }); }); afterAll(async () => { await app?.exit(); await upstream?.close(); + await groupOpenUpstream?.close(); }); test( @@ -150,4 +230,73 @@ describe("model IP restriction e2e (#557): allowed_cidrs gate before upstream", }, 60_000, ); + + test( + "model group: an out-of-range member drops out and the group serves from the rest", + async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + // Probe from IN range so readiness doesn't depend on the very + // exclusion this test is about. + await waitConfigPropagation(async () => { + const r = await chat(app!, MIXED_GROUP, "10.1.2.3"); + await r.text(); + return r.status === 200; + }); + + // Out-of-range caller: the restricted FIRST target is excluded, so the + // group is served by the open member instead. Pre-fix the group ignored + // the member allowlist and dispatched straight to the restricted one — + // which also answers 200, so the marker in the body is what makes this + // assertion discriminate between fixed and broken. + const restrictedHitsBefore = upstream!.receivedRequests.length; + const served = await chat(app, MIXED_GROUP, "114.114.114.114"); + expect(served.status).toBe(200); + const body = (await served.json()) as { + choices?: Array<{ message?: { content?: string } }>; + }; + expect(body.choices?.[0]?.message?.content).toBe(GROUP_OPEN_MARKER); + // The excluded member is never attempted — not attempted-then-failed-over. + expect(upstream!.receivedRequests.length).toBe(restrictedHitsBefore); + }, + 60_000, + ); + + test( + "model group: 403 when every member excludes the caller, upstream untouched", + async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + await waitConfigPropagation(async () => { + const r = await chat(app!, ALL_RESTRICTED_GROUP, "10.1.2.3"); + await r.text(); + return r.status === 200; + }); + + const hitsBefore = upstream.receivedRequests.length; + const blocked = await chat(app, ALL_RESTRICTED_GROUP, "114.114.114.114"); + expect(blocked.status).toBe(403); + const body = (await blocked.json()) as { + error?: { code?: string; message?: string }; + }; + expect(body.error?.code).toBe("ip_restricted"); + // Same generic envelope as the direct-model rejection: no model name + // and no CIDR reaches the caller, so a probe can't enumerate which + // members exist or what ranges they allow (the #557 rule, which the + // group path must not weaken by naming the target it excluded). + const message = body.error?.message ?? ""; + expect(message).toBe( + "Access denied: your client IP is not allowed to access this model", + ); + for (const internal of [RESTRICTED_MODEL, RESTRICTED_MODEL_2, "192.168"]) { + expect(message).not.toContain(internal); + } + expect(upstream.receivedRequests.length).toBe(hitsBefore); + }, + 60_000, + ); });