admin/CC-Switch
1
0
Fork 0
mirror of https://github.com/farion1231/cc-switch.git synced 2026-07-24 12:44:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(proxy/gemini): gate generic REST path suffixes behind Google host whitelist

Browse Source 
`should_normalize_gemini_full_url` previously treated any full URL whose
path ends with `/v1`, `/v1/models`, `/models`, `/v1/openai`, or `/openai`
as a structured Gemini endpoint and rewrote it to
`/v1beta/models/{model}:generateContent`. These are ubiquitous REST
conventions — opaque relays such as `https://relay.example/custom/v1`
legitimately use them for fixed endpoints — so the rewrite silently
routed traffic to the wrong upstream path.

Split the predicate into two layers:

- **Unconditional**: `matches_structured_gemini_models_path` (i.e. a
  `/models/...:generateContent` method call anywhere in the path), the
  Google-specific `/v1beta*` family, and the deep OpenAI-compat paths
  (`/v1beta/openai/chat/completions`, `/openai/chat/completions`, and
  their `responses` siblings). These remain host-agnostic because the
  path grammar itself is Gemini-specific.
- **Google-host gated**: `/v1`, `/v1/models`, `/models`, `/v1/openai`,
  `/openai`. Only normalized when the host is one of
  `generativelanguage.googleapis.com`, `aiplatform.googleapis.com`, or a
  real `*-aiplatform.googleapis.com` Vertex regional endpoint. The match
  is exact/suffix (not `contains`), so lookalike hosts like
  `aiplatform.example.com` are correctly treated as opaque relays.

Tests (8 new in `gemini_url::tests`):
- Four opaque-relay cases: `/custom/v1`, `/custom/models`,
  `/custom/v1/models`, `/custom/openai` — all preserved as-is.
- Three Google-host counter-cases: `/v1`, `/models`, and
  `us-central1-aiplatform.googleapis.com/v1` still normalize.
- One lookalike safety case: `aiplatform.example.com/v1` is NOT
  treated as Google.

All 849 lib tests pass; cargo fmt + clippy -D warnings clean.

Addresses Codex review P2 on #1918.
This commit is contained in:
Jason
2026-04-16 17:35:11 +08:00
parent baaefcc2fc
commit 1f3ebe9c62
1 changed files with 192 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/src/proxy/gemini_url.rs
+192 -17
View File
@@ -56,35 +56,81 @@ fn should_normalize_gemini_full_url(base_url: &str) -> bool {
.map_or(base_url, |(base, _)| base)
.trim_end_matches('/');
let (base_without_query, _) = split_query(base_url);
let (_, path) = split_origin_and_path(base_without_query);
let (origin, path) = split_origin_and_path(base_without_query);
if path.is_empty() || path == "/" {
return true;
}
let path = path.trim_end_matches('/');
// Only classify a path as a "structured Gemini endpoint" — safe to
// rewrite to /v1beta/models/{model}:method — when the `/models/`
// segment is actually followed by a Gemini `*:generateContent` /
// `*:streamGenerateContent` method call. Bare `path.contains("/v1/models/")`
// would also match legitimate opaque relay routes like
// `/v1/models/invoke`, which are *not* Gemini-structured and must be
// preserved as-is.
matches_structured_gemini_models_path(path)
let on_google_host = is_google_gemini_host(extract_host(origin));
// Unconditional layer: the `/models/...:generateContent` structure is
// Gemini-specific, and `/v1beta*` / deep OpenAI-compat paths
// (`/v1beta/openai/...`, `/openai/chat/completions`, `/openai/responses`)
// are specific enough to Gemini or AI-Studio that treating them as
// structured on any host is safe.
if matches_structured_gemini_models_path(path)
|| path.ends_with("/v1beta")
|| path.ends_with("/v1")
|| path.ends_with("/v1beta/models")
|| path.ends_with("/v1/models")
|| path.ends_with("/models")
|| path.ends_with("/v1beta/openai")
|| path.ends_with("/v1/openai")
|| path.ends_with("/openai")
|| path.ends_with("/v1beta/openai/chat/completions")
|| path.ends_with("/v1/openai/chat/completions")
|| path.ends_with("/openai/chat/completions")
|| path.ends_with("/v1beta/openai/responses")
|| path.ends_with("/v1/openai/responses")
|| path.ends_with("/openai/chat/completions")
|| path.ends_with("/openai/responses")
|| path.ends_with("/v1/openai/chat/completions")
|| path.ends_with("/v1/openai/responses")
{
return true;
}
// Generic REST suffixes (`/v1`, `/models`, etc.) legitimately appear on
// opaque relays — e.g. `https://relay.example/custom/v1` is a relay
// root, not a Gemini base. Only treat these as "structured Gemini" when
// the host itself is Google's Gemini or Vertex AI endpoint.
if on_google_host
&& (path.ends_with("/v1")
|| path.ends_with("/v1/models")
|| path.ends_with("/models")
|| path.ends_with("/v1/openai")
|| path.ends_with("/openai"))
{
return true;
}
false
}
/// Extract the host portion of an origin like `https://host:port` or
/// `https://host`. Returns an empty string if no host can be found (e.g.
/// bare `http://`).
fn extract_host(origin: &str) -> &str {
let after_scheme = origin.split_once("://").map_or(origin, |(_, rest)| rest);
// authority may carry credentials (`user:pass@host`) and a port
// (`host:port`). Strip userinfo first, then port.
let without_userinfo = after_scheme
.rsplit_once('@')
.map_or(after_scheme, |(_, h)| h);
let without_port = without_userinfo
.split_once(':')
.map_or(without_userinfo, |(h, _)| h);
// Strip trailing `/` defensively (split_origin_and_path already handled
// it, but this helper may be reused elsewhere).
without_port.trim_end_matches('/')
}
/// Returns true when `host` is one of Google's Gemini / Vertex AI endpoints.
/// Case-insensitive. Requires exact match or a real `-aiplatform.googleapis.com`
/// subdomain suffix — not a substring match, so lookalikes like
/// `aiplatform.example.com` are rejected.
fn is_google_gemini_host(host: &str) -> bool {
if host.is_empty() {
return false;
}
let host_lower = host.to_ascii_lowercase();
host_lower == "generativelanguage.googleapis.com"
|| host_lower == "aiplatform.googleapis.com"
|| host_lower.ends_with("-aiplatform.googleapis.com")
}
fn split_query(input: &str) -> (&str, Option<&str>) {
@@ -326,4 +372,133 @@ mod tests {
"https://relay.example/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
);
}
// ------------------------------------------------------------------
// Google-host whitelist tests (generic REST suffix handling)
//
// Generic REST conventions like `/v1`, `/models`, `/openai` legitimately
// appear on opaque relays. `should_normalize_gemini_full_url` only
// treats these as structured Gemini endpoints when the host itself is
// Google's Gemini or Vertex AI endpoint.
// ------------------------------------------------------------------
/// Regression: a relay whose fixed path ends with `/v1` (a ubiquitous
/// REST convention) used to be rewritten to
/// `/v1beta/models/{model}:generateContent`, dropping the relay's own
/// `/v1` endpoint.
#[test]
fn preserves_opaque_full_url_with_v1_suffix() {
let url = resolve_gemini_native_url(
"https://relay.example/custom/v1",
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
true,
);
assert_eq!(url, "https://relay.example/custom/v1?alt=sse");
}
/// Companion case: bare `/models` suffix on a non-Google host is a
/// generic REST path, not a Gemini-structured endpoint.
#[test]
fn preserves_opaque_full_url_with_models_suffix() {
let url = resolve_gemini_native_url(
"https://relay.example/custom/models",
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
true,
);
assert_eq!(url, "https://relay.example/custom/models?alt=sse");
}
/// Companion case: `/v1/models` — same ambiguity as `/models`, with the
/// version prefix. Must stay as-is on non-Google hosts.
#[test]
fn preserves_opaque_full_url_with_v1_models_suffix() {
let url = resolve_gemini_native_url(
"https://relay.example/custom/v1/models",
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
true,
);
assert_eq!(url, "https://relay.example/custom/v1/models?alt=sse");
}
/// Companion case: a relay that exposes an `/openai` compatibility
/// surface without the deep `/openai/chat/completions` path. Must stay
/// as-is on non-Google hosts.
#[test]
fn preserves_opaque_full_url_with_openai_suffix_on_non_google_host() {
let url = resolve_gemini_native_url(
"https://relay.example/custom/openai",
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
true,
);
assert_eq!(url, "https://relay.example/custom/openai?alt=sse");
}
/// Counter-case: `/v1` on the official Gemini host must still be
/// normalized to the full `/v1beta/models/...` endpoint — users who
/// paste `https://generativelanguage.googleapis.com/v1` as their base
/// URL expect the proxy to resolve the method path.
#[test]
fn normalizes_google_host_with_v1_suffix() {
let url = resolve_gemini_native_url(
"https://generativelanguage.googleapis.com/v1",
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
true,
);
assert_eq!(
url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
);
}
/// Counter-case: `/models` on the official Gemini host is recognized
/// and normalized.
#[test]
fn normalizes_google_host_with_models_suffix() {
let url = resolve_gemini_native_url(
"https://generativelanguage.googleapis.com/models",
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
true,
);
assert_eq!(
url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
);
}
/// Counter-case: Vertex AI regional endpoints live under
/// `*-aiplatform.googleapis.com`. Those should also be treated as
/// Google-host for the whitelist.
#[test]
fn normalizes_vertex_aiplatform_host_with_v1_suffix() {
let url = resolve_gemini_native_url(
"https://us-central1-aiplatform.googleapis.com/v1",
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
true,
);
assert_eq!(
url,
"https://us-central1-aiplatform.googleapis.com/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse"
);
}
/// Safety: the Google-host whitelist must do an exact/suffix match, not
/// a `contains`. A lookalike host like `aiplatform.example.com` must
/// NOT be treated as Google.
#[test]
fn preserves_non_google_aiplatform_lookalike_host() {
let url = resolve_gemini_native_url(
"https://aiplatform.example.com/v1",
"/v1beta/models/gemini-2.5-flash:streamGenerateContent?alt=sse",
true,
);
assert_eq!(url, "https://aiplatform.example.com/v1?alt=sse");
}
}