starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-07 00:05:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

codex: scope cookies

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-06-15 11:03:56 +08:00
co-authored by Codex
parent 44c3306023
commit e8557932bd
10 changed files with 278 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/codex/accounts/cookies/repository.rs
+70 -4
View File
@@ -73,9 +73,19 @@ impl CookieRepository {
&self,
account_id: &str,
request_domain: &str,
) -> CookieResult<Option<String>> {
self.cookie_header_for_request(account_id, request_domain, "/")
.await
}
pub async fn cookie_header_for_request(
&self,
account_id: &str,
request_domain: &str,
request_path: &str,
) -> CookieResult<Option<String>> {
let rows = sqlx::query(
"select domain, name, value_cipher, expires_at from account_cookies where account_id = ? order by name asc",
"select domain, name, value_cipher, path, expires_at from account_cookies where account_id = ?",
)
.bind(account_id)
.fetch_all(&self.pool)
@@ -87,18 +97,38 @@ impl CookieRepository {
if !domain_matches(request_domain, &domain) {
continue;
}
let path = row.get::<String, _>("path");
if !path_matches(request_path, &path) {
continue;
}
if cookie_is_expired(row.get::<Option<String>, _>("expires_at").as_deref(), now) {
continue;
}
let name = row.get::<String, _>("name");
let value_cipher = row.get::<String, _>("value_cipher");
let value = self.secret_box.decrypt(&value_cipher)?;
pairs.push(format!("{name}={}", value.expose_secret()));
pairs.push(CookieHeaderPair {
path_len: path.len(),
name: name.clone(),
value: format!("{name}={}", value.expose_secret()),
});
}
if pairs.is_empty() {
Ok(None)
} else {
Ok(Some(pairs.join("; ")))
pairs.sort_by(|left, right| {
right
.path_len
.cmp(&left.path_len)
.then_with(|| left.name.cmp(&right.name))
});
Ok(Some(
pairs
.into_iter()
.map(|pair| pair.value)
.collect::<Vec<_>>()
.join("; "),
))
}
}
@@ -139,6 +169,12 @@ struct ParsedCookie {
expires_at: Option<String>,
}
struct CookieHeaderPair {
path_len: usize,
name: String,
value: String,
}
fn parse_set_cookie(raw: &str) -> Option<ParsedCookie> {
let mut parts = raw.split(';').map(str::trim);
let (name, value) = parts.next()?.split_once('=')?;
@@ -151,7 +187,7 @@ fn parse_set_cookie(raw: &str) -> Option<ParsedCookie> {
};
match attribute.trim().to_ascii_lowercase().as_str() {
"domain" => domain = value.trim_start_matches('.').to_string(),
"path" => path = value.to_string(),
"path" => path = normalize_cookie_path(value),
"max-age" => {
if let Ok(seconds) = value.parse::<i64>() {
expires_at = Some(max_age_expires_at(seconds));
@@ -207,6 +243,36 @@ fn domain_matches(request_domain: &str, cookie_domain: &str) -> bool {
.is_some_and(|prefix| prefix.ends_with('.'))
}
fn normalize_cookie_path(path: &str) -> String {
let path = path.trim();
if path.starts_with('/') {
path.to_string()
} else {
"/".to_string()
}
}
fn path_matches(request_path: &str, cookie_path: &str) -> bool {
let request_path = normalize_request_path(request_path);
let cookie_path = normalize_cookie_path(cookie_path);
request_path == cookie_path
|| (request_path.starts_with(&cookie_path)
&& (cookie_path.ends_with('/')
|| request_path
.as_bytes()
.get(cookie_path.len())
.is_some_and(|byte| *byte == b'/')))
}
fn normalize_request_path(path: &str) -> String {
let path = path.trim();
if path.starts_with('/') {
path.to_string()
} else {
format!("/{path}")
}
}
fn cookie_is_expired(expires_at: Option<&str>, now: DateTime<Utc>) -> bool {
expires_at
.and_then(parse_cookie_expires_at)
src/codex/accounts/service/health.rs
+3 -1
View File
@@ -12,6 +12,7 @@ use crate::{
model::AccountStatus,
repository::{AccountRepository, StoredAccount},
},
codex::gateway::transport::endpoints::primary_usage_request_path,
codex::gateway::transport::http_client::{
build_reqwest_client, CodexBackendClient, CodexClientError, CodexRequestContext,
},
@@ -111,10 +112,11 @@ pub(super) async fn fetch_account_usage(
async fn account_cookie_header(service: &AccountService, account_id: &str) -> Option<String> {
let domain = request_domain(&service.config.api.base_url)?;
let request_path = primary_usage_request_path(&service.config.api.base_url);
service
.cookie_repository
.as_ref()?
.cookie_header(account_id, &domain)
.cookie_header_for_request(account_id, &domain, &request_path)
.await
.ok()
.flatten()
src/codex/gateway/transport/endpoints.rs
+61
View File
@@ -0,0 +1,61 @@
use reqwest::Url;
pub(crate) const CODEX_RESPONSES_PATH: &str = "/codex/responses";
pub(crate) const CODEX_RESPONSES_COMPACT_PATH: &str = "/codex/responses/compact";
pub(crate) const CODEX_USAGE_PATH: &str = "/codex/usage";
pub(crate) const CODEX_USAGE_API_PATH: &str = "/api/codex/usage";
pub(crate) const WHAM_USAGE_PATH: &str = "/wham/usage";
pub(crate) fn endpoint_url(base_url: &str, endpoint_path: &str) -> String {
format!(
"{}/{}",
base_url.trim_end_matches('/'),
endpoint_path.trim_start_matches('/')
)
}
pub(crate) fn endpoint_request_path(base_url: &str, endpoint_path: &str) -> String {
let endpoint_path = endpoint_path.trim_start_matches('/');
let base_path = Url::parse(base_url)
.ok()
.map(|url| url.path().trim_end_matches('/').to_string())
.filter(|path| !path.is_empty())
.unwrap_or_default();
if base_path.is_empty() {
format!("/{endpoint_path}")
} else {
format!("{base_path}/{endpoint_path}")
}
}
pub(crate) fn usage_endpoint_urls(base_url: &str) -> Vec<String> {
usage_endpoint_paths(base_url)
.into_iter()
.map(|path| endpoint_url(base_url, path))
.collect()
}
pub(crate) fn primary_usage_request_path(base_url: &str) -> String {
let endpoint_path = usage_endpoint_paths(base_url)
.into_iter()
.next()
.unwrap_or(CODEX_USAGE_API_PATH);
endpoint_request_path(base_url, endpoint_path)
}
fn usage_endpoint_paths(base_url: &str) -> Vec<&'static str> {
if has_backend_api_base_path(base_url) {
vec![WHAM_USAGE_PATH, CODEX_USAGE_PATH]
} else {
vec![CODEX_USAGE_API_PATH, CODEX_USAGE_PATH]
}
}
fn has_backend_api_base_path(base_url: &str) -> bool {
Url::parse(base_url).ok().is_some_and(|url| {
url.path()
.split('/')
.any(|segment| segment == "backend-api")
})
}
src/codex/gateway/transport/http_client.rs
+6 -13
View File
@@ -20,6 +20,9 @@ use uuid::Uuid;
use crate::{
codex::gateway::fingerprint::model::Fingerprint,
codex::gateway::transport::{
endpoints::{
endpoint_url, usage_endpoint_urls, CODEX_RESPONSES_COMPACT_PATH, CODEX_RESPONSES_PATH,
},
headers::build_ordered_codex_base_headers,
sse::SseError,
types::{CodexCompactRequest, CodexResponsesRequest},
@@ -491,7 +494,7 @@ impl CodexBackendClient {
request: &CodexResponsesRequest,
headers: HeaderMap,
) -> CodexClientResult<ReqwestResponse> {
let url = format!("{}/codex/responses", self.base_url);
let url = endpoint_url(&self.base_url, CODEX_RESPONSES_PATH);
Ok(self
.client
.post(url)
@@ -506,7 +509,7 @@ impl CodexBackendClient {
request: &CodexCompactRequest,
headers: HeaderMap,
) -> CodexClientResult<ReqwestResponse> {
let url = format!("{}/codex/responses/compact", self.base_url);
let url = endpoint_url(&self.base_url, CODEX_RESPONSES_COMPACT_PATH);
Ok(self
.client
.post(url)
@@ -517,17 +520,7 @@ impl CodexBackendClient {
}
fn usage_endpoints(&self) -> Vec<String> {
if self.base_url.contains("/backend-api") {
vec![
format!("{}/wham/usage", self.base_url),
format!("{}/codex/usage", self.base_url),
]
} else {
vec![
format!("{}/api/codex/usage", self.base_url),
format!("{}/codex/usage", self.base_url),
]
}
usage_endpoint_urls(&self.base_url)
}
fn request_headers(&self, context: CodexRequestContext<'_>) -> CodexClientResult<HeaderMap> {
src/codex/gateway/transport/mod.rs
+1
View File
@@ -1,3 +1,4 @@
pub mod endpoints;
pub mod headers;
pub mod http_client;
pub mod rate_limits;
src/codex/gateway/transport/websocket/mod.rs
+2 -1
View File
@@ -14,6 +14,7 @@ use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
use reqwest::{header::HeaderMap, StatusCode};
use crate::codex::gateway::transport::{
endpoints::{endpoint_url, CODEX_RESPONSES_PATH},
rate_limits::{parse_rate_limits_event_raw, rate_limits_to_header_pairs, ParsedRateLimits},
types::CodexResponsesRequest,
};
@@ -471,7 +472,7 @@ fn build_ws_request(
}
fn websocket_url(base_url: &str) -> String {
let url = format!("{}/codex/responses", base_url.trim_end_matches('/'));
let url = endpoint_url(base_url, CODEX_RESPONSES_PATH);
if let Some(rest) = url.strip_prefix("https://") {
format!("wss://{rest}")
} else if let Some(rest) = url.strip_prefix("http://") {
src/codex/serving/dispatch/mod.rs
+37 -26
View File
@@ -48,6 +48,10 @@ use crate::{
codex::gateway::installation_id::get_installation_id,
codex::gateway::protocol::codex_to_openai::openai_error,
codex::gateway::transport::{
endpoints::{
endpoint_request_path, primary_usage_request_path, CODEX_RESPONSES_COMPACT_PATH,
CODEX_RESPONSES_PATH,
},
http_client::{
build_reqwest_client, CodexBackendClient, CodexBackendStream,
CodexBackendWebSocketStream, CodexClientError, CodexCompactResponse,
@@ -90,7 +94,6 @@ use self::reasoning_replay::ReasoningReplayCache;
use self::stream_audit::{StreamAudit, WebSocketStreamAudit};
use self::{
limits::apply_rate_limit_headers_with_deps,
routing::request_domain,
stream::{
completed_response_metadata, ensure_stream_metadata, has_terminal_sse_event,
premature_close_failed_event, TupleStreamReconverter,
@@ -471,6 +474,28 @@ async fn record_failed_request_attempt_with_deps(
}
}
async fn upstream_cookie_header_for_endpoint(
deps: &CodexUpstreamDependencies,
account_id: &str,
endpoint_path: &str,
) -> Option<String> {
let request_path = endpoint_request_path(&deps.config.api.base_url, endpoint_path);
upstream_cookie_header_for_request_path(deps, account_id, &request_path).await
}
async fn upstream_cookie_header_for_request_path(
deps: &CodexUpstreamDependencies,
account_id: &str,
request_path: &str,
) -> Option<String> {
let repo = deps.cookie_repository.as_ref()?;
let domain = routing::request_domain(&deps.config.api.base_url)?;
repo.cookie_header_for_request(account_id, &domain, request_path)
.await
.ok()
.flatten()
}
struct DirtyQuotaVerification {
limit_reached: bool,
}
@@ -495,11 +520,9 @@ async fn fetch_dirty_account_usage_with_deps(
account: &Account,
request_id: &str,
) -> Result<Value, CodexClientError> {
let request_domain = request_domain(&deps.config.api.base_url);
let cookie_header = match (deps.cookie_repository.as_ref(), request_domain.as_deref()) {
(Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(),
_ => None,
};
let usage_path = primary_usage_request_path(&deps.config.api.base_url);
let cookie_header =
upstream_cookie_header_for_request_path(deps, &account.id, &usage_path).await;
let client = CodexBackendClient::new(
build_reqwest_client(deps.config.tls.force_http11)?,
deps.config.api.base_url.clone(),
@@ -1157,11 +1180,8 @@ async fn send_codex_request(
account: &Account,
request_id: &str,
) -> Result<crate::codex::gateway::transport::http_client::CodexBackendResponse, CodexClientError> {
let request_domain = request_domain(&deps.config.api.base_url);
let cookie_header = match (deps.cookie_repository.as_ref(), request_domain.as_deref()) {
(Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(),
_ => None,
};
let cookie_header =
upstream_cookie_header_for_endpoint(deps, &account.id, CODEX_RESPONSES_PATH).await;
let account_scope = &account.id;
let identity = build_conversation_identity(
@@ -1206,11 +1226,8 @@ async fn send_compact_request(
account: &Account,
request_id: &str,
) -> Result<CodexCompactResponse, CodexClientError> {
let request_domain = request_domain(&deps.config.api.base_url);
let cookie_header = match (deps.cookie_repository.as_ref(), request_domain.as_deref()) {
(Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(),
_ => None,
};
let cookie_header =
upstream_cookie_header_for_endpoint(deps, &account.id, CODEX_RESPONSES_COMPACT_PATH).await;
let installation_id = get_installation_id(Some(&deps.config.database.url));
let client = CodexBackendClient::new(
@@ -1261,11 +1278,8 @@ async fn send_codex_stream_request(
account: &Account,
request_id: &str,
) -> Result<CodexBackendStream, CodexClientError> {
let request_domain = request_domain(&deps.config.api.base_url);
let cookie_header = match (deps.cookie_repository.as_ref(), request_domain.as_deref()) {
(Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(),
_ => None,
};
let cookie_header =
upstream_cookie_header_for_endpoint(deps, &account.id, CODEX_RESPONSES_PATH).await;
// Build conversation identity for session affinity
// Use account.id as the scope since entry_id doesn't exist
@@ -1371,11 +1385,8 @@ async fn send_codex_websocket_stream_request(
account: &Account,
request_id: &str,
) -> Result<CodexBackendWebSocketStream, CodexClientError> {
let request_domain = request_domain(&deps.config.api.base_url);
let cookie_header = match (deps.cookie_repository.as_ref(), request_domain.as_deref()) {
(Some(repo), Some(domain)) => repo.cookie_header(&account.id, domain).await.ok().flatten(),
_ => None,
};
let cookie_header =
upstream_cookie_header_for_endpoint(deps, &account.id, CODEX_RESPONSES_PATH).await;
let account_scope = &account.id;
let identity = build_conversation_identity(