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

fix: harden account scheduling and client context

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-04 19:56:32 +08:00
co-authored by Codex
parent 80f5ebaf13
commit 583a61da7d
13 changed files with 855 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/admin/auth/session.rs
+11 -11
View File
@@ -4,7 +4,7 @@ use axum::{
extract::{FromRequestParts, State},
http::{header::SET_COOKIE, request::Parts, HeaderMap, HeaderValue, StatusCode},
response::{IntoResponse, Response},
Json,
Extension, Json,
};
use serde::{Deserialize, Serialize};
@@ -13,6 +13,7 @@ use crate::{
auth::service::AdminSessionError,
response::{AdminEnvelope, AdminError, AdminResponse},
},
http::middleware::request_id::ClientIp,
infra::time::china_rfc3339,
runtime::state::AppState,
};
@@ -57,10 +58,10 @@ struct AdminSessionStatusData {
/// `POST /api/admin/login`
pub(crate) async fn login(
State(state): State<AppState>,
headers: HeaderMap,
client_ip: Option<Extension<ClientIp>>,
Json(payload): Json<AdminLoginRequest>,
) -> Result<Response, AdminError> {
let source = admin_login_source(&headers);
let source = admin_login_source(optional_client_ip(client_ip));
let login_result = state
.services
.admin_sessions
@@ -191,12 +192,11 @@ fn admin_api_key_header(headers: &HeaderMap) -> Option<String> {
(!value.is_empty()).then(|| value.to_string())
}
fn admin_login_source(headers: &HeaderMap) -> String {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("unknown")
.to_string()
fn optional_client_ip(client_ip: Option<Extension<ClientIp>>) -> Option<String> {
let Extension(client_ip) = client_ip?;
Some(client_ip.as_str().to_string())
}
fn admin_login_source(fallback_client_ip: Option<String>) -> String {
fallback_client_ip.unwrap_or_else(|| "unknown".to_string())
}
backend/src/http/middleware/request_id.rs
+34 -31
View File
@@ -14,14 +14,8 @@ use uuid::Uuid;
/// 请求 ID 头。
pub const REQUEST_ID_HEADER: &str = "x-request-id";
const MAX_REQUEST_ID_LEN: usize = 128;
const CF_CONNECTING_IP_HEADER: &str = "cf-connecting-ip";
const REAL_IP_HEADER: &str = "x-real-ip";
const X_FORWARDED_FOR_HEADER: &str = "x-forwarded-for";
const UNTRUSTED_CLIENT_IP_HEADERS: &[&str] = &[
CF_CONNECTING_IP_HEADER,
REAL_IP_HEADER,
X_FORWARDED_FOR_HEADER,
];
/// 可信反向代理配置。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -77,6 +71,17 @@ impl RequestId {
}
}
/// 由连接来源和可信代理链解析出的客户端 IP。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientIp(String);
impl ClientIp {
/// 返回 IP 字符串。
pub fn as_str(&self) -> &str {
&self.0
}
}
fn valid_request_id(value: &str) -> bool {
!value.is_empty()
&& value.len() <= MAX_REQUEST_ID_LEN
@@ -117,10 +122,10 @@ async fn run_with_request_id(mut request: Request, next: Next) -> Response {
}
fn attach_request_context(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
attach_real_ip_from_connection(trusted_proxies, request);
attach_client_ip_from_connection(trusted_proxies, request);
}
fn attach_real_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
fn attach_client_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
let peer_ip = request
.extensions()
.get::<ConnectInfo<SocketAddr>>()
@@ -128,47 +133,45 @@ fn attach_real_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request:
let client_ip =
peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers(), trusted_proxies));
for header in UNTRUSTED_CLIENT_IP_HEADERS {
request.headers_mut().remove(*header);
}
let Some(client_ip) = client_ip else {
return;
};
if let Ok(value) = HeaderValue::from_str(&client_ip.to_string()) {
request.headers_mut().insert(REAL_IP_HEADER, value);
}
request.extensions_mut().insert(ClientIp(client_ip));
}
fn client_ip_from_peer(
peer_ip: IpAddr,
headers: &HeaderMap,
trusted_proxies: &TrustedProxyConfig,
) -> IpAddr {
) -> String {
if trusted_proxies.is_trusted(peer_ip) {
trusted_forwarded_client_ip(headers).unwrap_or(peer_ip)
trusted_forwarded_client_ip(headers).unwrap_or_else(|| peer_ip.to_string())
} else {
peer_ip
peer_ip.to_string()
}
}
fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option<IpAddr> {
header_ip(headers, CF_CONNECTING_IP_HEADER)
.or_else(|| first_forwarded_for_ip(headers))
.or_else(|| header_ip(headers, REAL_IP_HEADER))
fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
first_forwarded_for_value(headers).or_else(|| header_string(headers, REAL_IP_HEADER))
}
fn header_ip(headers: &HeaderMap, name: &str) -> Option<IpAddr> {
parse_ip(headers.get(name)?.to_str().ok()?)
fn header_string(headers: &HeaderMap, name: &str) -> Option<String> {
headers
.get(name)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn first_forwarded_for_ip(headers: &HeaderMap) -> Option<IpAddr> {
let value = headers.get(X_FORWARDED_FOR_HEADER)?.to_str().ok()?;
parse_ip(value.split(',').next()?)
}
fn parse_ip(value: &str) -> Option<IpAddr> {
value.trim().parse().ok()
fn first_forwarded_for_value(headers: &HeaderMap) -> Option<String> {
let value = header_string(headers, X_FORWARDED_FOR_HEADER)?;
value
.split(',')
.next()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn parse_trusted_proxy_entry(entry: &str) -> Result<IpNet, TrustedProxyConfigError> {
backend/src/proxy/dispatch/responses/service.rs
+134
View File
@@ -1231,6 +1231,140 @@ impl ResponseDispatchService {
history_recovery_used = true;
continue;
}
Err(ResponseDispatchError::Upstream(error))
if is_rate_limit_upstream_error(&error) =>
{
self.account_pool.release(&release_account_id).await;
rate_limited_count += 1;
last_rate_limit_error = Some(upstream_error_body(&error));
last_exhausted_account_class = Some(ExhaustedAccountClass::RateLimited);
let cooldown_until = rate_limit_cooldown_until(&error, Utc::now());
self.account_pool
.mark_quota_limited_until(&release_account_id, cooldown_until)
.await;
excluded_account_ids.push(release_account_id);
continue;
}
Err(ResponseDispatchError::Upstream(error))
if is_quota_exhausted_upstream_error(&error) =>
{
self.account_pool.release(&release_account_id).await;
quota_exhausted_count += 1;
last_quota_error = Some(upstream_error_body(&error));
last_exhausted_account_class =
Some(ExhaustedAccountClass::QuotaExhausted);
self.account_pool
.set_status(&release_account_id, AccountStatus::QuotaExhausted)
.await;
excluded_account_ids.push(release_account_id);
continue;
}
Err(ResponseDispatchError::Upstream(error))
if is_auth_upstream_error(&error) =>
{
self.account_pool.release(&release_account_id).await;
let upstream_error = upstream_error_body(&error);
let account_status = auth_failure_account_status(&error);
match account_status {
AccountStatus::Disabled => {
disabled_count += 1;
last_disabled_auth_error = Some(upstream_error);
last_exhausted_account_class =
Some(ExhaustedAccountClass::Disabled);
}
AccountStatus::Banned => {
banned_count += 1;
last_banned_status_code =
Some(upstream_error_http_status(&error));
last_banned_auth_error = Some(upstream_error);
last_exhausted_account_class =
Some(ExhaustedAccountClass::Banned);
}
_ => {
expired_count += 1;
last_auth_error = Some(upstream_error);
last_exhausted_account_class =
Some(ExhaustedAccountClass::Expired);
}
}
self.account_pool
.set_status(&release_account_id, account_status)
.await;
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
continue;
}
Err(ResponseDispatchError::Upstream(error))
if is_cloudflare_challenge_upstream_error(&error) =>
{
self.account_pool.release(&release_account_id).await;
cloudflare_challenge_count += 1;
last_cloudflare_challenge_error =
Some(cloudflare_challenge_error_message().to_string());
last_exhausted_account_class =
Some(ExhaustedAccountClass::CloudflareChallenge);
self.cloudflare
.apply_challenge(self.account_pool.as_ref(), &release_account_id)
.await;
excluded_account_ids.push(release_account_id);
continue;
}
Err(ResponseDispatchError::Upstream(error))
if is_cloudflare_path_block_upstream_error(&error) =>
{
self.account_pool.release(&release_account_id).await;
cloudflare_path_block_count += 1;
last_cloudflare_path_block_error =
Some(cloudflare_path_block_error_message().to_string());
last_exhausted_account_class =
Some(ExhaustedAccountClass::CloudflarePathBlocked);
self.cloudflare
.apply_path_block(self.account_pool.as_ref(), &release_account_id)
.await;
excluded_account_ids.push(release_account_id);
continue;
}
Err(ResponseDispatchError::Upstream(error))
if is_model_unsupported_upstream_error(&error) =>
{
self.account_pool.release(&release_account_id).await;
let upstream_error = upstream_error_body(&error);
if model_unsupported_retry_used {
return_stream_dispatch_error!(
ResponseDispatchError::ModelUnsupported {
count: model_unsupported_count + 1,
upstream_error,
},
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(transport))
);
}
model_unsupported_count += 1;
last_model_unsupported_error = Some(upstream_error);
last_exhausted_account_class =
Some(ExhaustedAccountClass::ModelUnsupported);
model_unsupported_retry_used = true;
excluded_account_ids.push(release_account_id);
continue;
}
Err(ResponseDispatchError::Upstream(error))
if is_banned_upstream_error(&error) =>
{
self.account_pool.release(&release_account_id).await;
banned_count += 1;
last_banned_status_code = Some(upstream_error_http_status(&error));
last_banned_auth_error = Some(upstream_error_body(&error));
last_exhausted_account_class = Some(ExhaustedAccountClass::Banned);
self.account_pool
.set_status(&release_account_id, AccountStatus::Banned)
.await;
excluded_account_ids.push(release_account_id);
continue;
}
Err(error) => {
self.account_pool.release(&release_account_id).await;
if let ResponseDispatchError::Upstream(upstream_error) = &error {
backend/src/proxy/openai/responses.rs
-53
View File
@@ -9,7 +9,6 @@ use axum::{
};
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use std::net::IpAddr;
use crate::{
http::middleware::request_id::RequestId,
@@ -612,7 +611,6 @@ async fn handle_compact_responses(
}
};
let mut compact_request = translate_response_to_compact(openai_request);
compact_request.client_ip = client_ip_from_headers(&headers);
compact_request.client_user_agent = user_agent_from_headers(&headers);
match state
@@ -694,7 +692,6 @@ fn fill_string_from_header(field: &mut Option<String>, headers: &HeaderMap, name
}
pub(crate) fn attach_client_context(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
request.client_ip = client_ip_from_headers(headers);
request.client_user_agent = user_agent_from_headers(headers);
}
@@ -702,56 +699,6 @@ pub(crate) fn user_agent_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, "user-agent")
}
pub(crate) fn client_ip_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, "cf-connecting-ip")
.or_else(|| header_string(headers, "x-real-ip"))
.or_else(|| forwarded_client_ip(headers))
.and_then(|value| normalize_client_ip(&value))
}
fn forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
let value = header_string(headers, "x-forwarded-for")?;
let parts = value
.split(',')
.map(str::trim)
.filter(|part| !part.is_empty())
.collect::<Vec<_>>();
parts
.iter()
.find(|part| !is_private_ip(part))
.or_else(|| parts.first())
.map(|part| (*part).to_string())
}
fn normalize_client_ip(value: &str) -> Option<String> {
let trimmed = value.trim();
if trimmed.is_empty() {
return None;
}
trimmed
.parse::<std::net::SocketAddr>()
.map(|addr| addr.ip().to_string())
.ok()
.or_else(|| Some(trimmed.to_string()))
}
fn is_private_ip(value: &str) -> bool {
let Some(ip) = normalize_client_ip(value).and_then(|value| value.parse::<IpAddr>().ok()) else {
return false;
};
match ip {
IpAddr::V4(ip) => {
ip.is_private()
|| ip.is_loopback()
|| ip.is_link_local()
|| ip.octets()[0] == 100 && (64..=127).contains(&ip.octets()[1])
}
IpAddr::V6(ip) => ip.is_loopback() || ip.is_unique_local() || ip.is_unicast_link_local(),
}
}
fn openai_subagent_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, OPENAI_SUBAGENT_HEADER).and_then(|value| {
if matches!(
backend/src/upstream/accounts/pool.rs
+120 -23
View File
@@ -154,13 +154,19 @@ pub struct DistinctPlanAccount {
#[derive(Debug)]
struct AccountAcquireWithStatusRefresh {
acquired: Option<AcquiredAccount>,
expired_account_ids: Vec<String>,
refreshed_accounts: Vec<RefreshedAccountState>,
}
#[derive(Debug)]
struct DistinctPlanAccountsWithStatusRefresh {
accounts: Vec<DistinctPlanAccount>,
expired_account_ids: Vec<String>,
refreshed_accounts: Vec<RefreshedAccountState>,
}
#[derive(Debug)]
struct AccountCapacitySummaryWithStatusRefresh {
summary: AccountCapacitySummary,
refreshed_accounts: Vec<RefreshedAccountState>,
}
/// 账号池容量摘要。
@@ -176,6 +182,71 @@ pub struct AccountCapacitySummary {
pub available_slots: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct RuntimeAccountStateSnapshot {
status: AccountStatus,
quota_limit_reached: bool,
quota_verify_required: bool,
quota_cooldown_until: Option<DateTime<Utc>>,
cloudflare_cooldown_until: Option<DateTime<Utc>>,
window_request_count: u64,
window_input_tokens: u64,
window_output_tokens: u64,
window_cached_tokens: u64,
window_image_input_tokens: u64,
window_image_output_tokens: u64,
window_image_request_count: u64,
window_image_request_failed_count: u64,
window_started_at: Option<DateTime<Utc>>,
window_reset_at: Option<DateTime<Utc>>,
limit_window_seconds: Option<u64>,
}
impl From<&Account> for RuntimeAccountStateSnapshot {
fn from(account: &Account) -> Self {
Self {
status: account.status,
quota_limit_reached: account.quota_limit_reached,
quota_verify_required: account.quota_verify_required,
quota_cooldown_until: account.quota_cooldown_until,
cloudflare_cooldown_until: account.cloudflare_cooldown_until,
window_request_count: account.window_request_count,
window_input_tokens: account.window_input_tokens,
window_output_tokens: account.window_output_tokens,
window_cached_tokens: account.window_cached_tokens,
window_image_input_tokens: account.window_image_input_tokens,
window_image_output_tokens: account.window_image_output_tokens,
window_image_request_count: account.window_image_request_count,
window_image_request_failed_count: account.window_image_request_failed_count,
window_started_at: account.window_started_at,
window_reset_at: account.window_reset_at,
limit_window_seconds: account.limit_window_seconds,
}
}
}
impl RuntimeAccountStateSnapshot {
fn usage_window_changed(&self, other: &Self) -> bool {
self.window_request_count != other.window_request_count
|| self.window_input_tokens != other.window_input_tokens
|| self.window_output_tokens != other.window_output_tokens
|| self.window_cached_tokens != other.window_cached_tokens
|| self.window_image_input_tokens != other.window_image_input_tokens
|| self.window_image_output_tokens != other.window_image_output_tokens
|| self.window_image_request_count != other.window_image_request_count
|| self.window_image_request_failed_count != other.window_image_request_failed_count
|| self.window_started_at != other.window_started_at
|| self.window_reset_at != other.window_reset_at
|| self.limit_window_seconds != other.limit_window_seconds
}
}
#[derive(Debug, Clone)]
struct RefreshedAccountState {
account: Account,
sync_usage_window: bool,
}
/// 运行时窗口用量增量。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountWindowUsageDelta {
@@ -419,7 +490,7 @@ impl AccountPool {
request: &AccountAcquireRequest,
) -> AccountAcquireWithStatusRefresh {
self.cleanup_stale_slots(request.now);
let expired_account_ids = self.refresh_account_statuses(request.now);
let refreshed_accounts = self.refresh_account_statuses(request.now);
let candidates = self.candidates(request);
let selected = if let Some(preferred_account_id) = &request.preferred_account_id {
let preferred = candidates
@@ -458,7 +529,7 @@ impl AccountPool {
let Some((selected, previous_slot_at, selection_source)) = selected else {
return AccountAcquireWithStatusRefresh {
acquired: None,
expired_account_ids,
refreshed_accounts,
};
};
tracing::info!(
@@ -485,13 +556,15 @@ impl AccountPool {
account: selected,
previous_slot_at,
}),
expired_account_ids,
refreshed_accounts,
}
}
fn refresh_account_statuses(&mut self, now: DateTime<Utc>) -> Vec<String> {
fn refresh_account_statuses(&mut self, now: DateTime<Utc>) -> Vec<RefreshedAccountState> {
let mut refreshed_accounts = Vec::new();
let mut expired_account_ids = Vec::new();
for (account_id, account) in &mut self.accounts {
let before = RuntimeAccountStateSnapshot::from(&*account);
if account.status == AccountStatus::Active && access_token_expired(account, now) {
account.status = AccountStatus::Expired;
expired_account_ids.push(account_id.clone());
@@ -499,11 +572,18 @@ impl AccountPool {
refresh_quota_window(account, now);
refresh_cloudflare_cooldown(account, now);
}
let after = RuntimeAccountStateSnapshot::from(&*account);
if before != after {
refreshed_accounts.push(RefreshedAccountState {
account: account.clone(),
sync_usage_window: before.usage_window_changed(&after),
});
}
}
for account_id in &expired_account_ids {
self.slots.remove(account_id);
}
expired_account_ids
refreshed_accounts
}
/// 释放指定账号的一个在途槽位。
@@ -527,8 +607,15 @@ impl AccountPool {
/// 计算账号池容量摘要。
pub fn capacity_summary(&mut self, now: DateTime<Utc>) -> AccountCapacitySummary {
self.capacity_summary_with_status_refresh(now).summary
}
fn capacity_summary_with_status_refresh(
&mut self,
now: DateTime<Utc>,
) -> AccountCapacitySummaryWithStatusRefresh {
self.cleanup_stale_slots(now);
self.refresh_account_statuses(now);
let refreshed_accounts = self.refresh_account_statuses(now);
let is_capacity_account = |account: &Account| {
account.status == AccountStatus::Active
&& AccountService::quota_available_at(account, now, self.options.skip_quota_limited)
@@ -550,11 +637,14 @@ impl AccountPool {
.map(|(_, slots)| slots.len().min(self.options.max_concurrent_per_account))
.sum();
AccountCapacitySummary {
max_concurrent_per_account: self.options.max_concurrent_per_account,
total_slots,
used_slots,
available_slots: total_slots.saturating_sub(used_slots),
AccountCapacitySummaryWithStatusRefresh {
summary: AccountCapacitySummary {
max_concurrent_per_account: self.options.max_concurrent_per_account,
total_slots,
used_slots,
available_slots: total_slots.saturating_sub(used_slots),
},
refreshed_accounts,
}
}
@@ -569,7 +659,7 @@ impl AccountPool {
now: DateTime<Utc>,
) -> DistinctPlanAccountsWithStatusRefresh {
self.cleanup_stale_slots(now);
let expired_account_ids = self.refresh_account_statuses(now);
let refreshed_accounts = self.refresh_account_statuses(now);
let mut by_plan = IndexMap::<String, Vec<Account>>::new();
for account in self.accounts.values() {
@@ -615,7 +705,7 @@ impl AccountPool {
DistinctPlanAccountsWithStatusRefresh {
accounts,
expired_account_ids,
refreshed_accounts,
}
}
@@ -874,7 +964,14 @@ impl RuntimeAccountPoolService {
/// 读取账号池容量摘要。
pub async fn capacity_summary(&self, now: DateTime<Utc>) -> AccountCapacitySummary {
self.pool.lock().await.capacity_summary(now)
let refresh = self
.pool
.lock()
.await
.capacity_summary_with_status_refresh(now);
self.persist_runtime_account_states(refresh.refreshed_accounts)
.await;
refresh.summary
}
/// 使用当前时间读取账号池容量摘要。
@@ -964,7 +1061,7 @@ impl RuntimeAccountPoolService {
/// 按请求上下文获取可用账号。
pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option<AcquiredAccount> {
let refresh = self.pool.lock().await.acquire_with_status_refresh(request);
self.persist_expired_statuses(refresh.expired_account_ids)
self.persist_runtime_account_states(refresh.refreshed_accounts)
.await;
refresh.acquired
}
@@ -976,22 +1073,22 @@ impl RuntimeAccountPoolService {
.lock()
.await
.distinct_plan_accounts_with_status_refresh(now);
self.persist_expired_statuses(refresh.expired_account_ids)
self.persist_runtime_account_states(refresh.refreshed_accounts)
.await;
refresh.accounts
}
async fn persist_expired_statuses(&self, account_ids: Vec<String>) {
for account_id in account_ids {
async fn persist_runtime_account_states(&self, refreshed_states: Vec<RefreshedAccountState>) {
for refreshed in refreshed_states {
if let Err(error) = self
.store
.set_status(&account_id, AccountStatus::Expired)
.sync_runtime_account_state(&refreshed.account, refreshed.sync_usage_window)
.await
{
tracing::warn!(
account_id,
account_id = %refreshed.account.id,
error = %error,
"failed to persist runtime-expired account status"
"failed to persist refreshed runtime account state"
);
}
}
backend/src/upstream/accounts/store/mod.rs
+77
View File
@@ -338,6 +338,13 @@ pub trait AccountStore: Send + Sync + 'static {
cooldown_until: Option<DateTime<Utc>>,
) -> AccountStoreResult<bool>;
/// 同步运行时自然刷新出来的账号状态。
async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> AccountStoreResult<bool>;
/// 同步账号当前 rate-limit 统计窗口。
async fn sync_rate_limit_window(
&self,
@@ -711,6 +718,66 @@ impl SqliteAccountStore {
Ok(result.rows_affected() > 0)
}
/// 同步运行时自然刷新出来的账号状态。
pub async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> SqliteAccountStoreResult<bool> {
let now = Utc::now().to_rfc3339();
let quota_limit_reached = if account.quota_limit_reached { 1 } else { 0 };
let quota_verify_required = if account.quota_verify_required { 1 } else { 0 };
let quota_cooldown_until = account.quota_cooldown_until.map(|dt| dt.to_rfc3339());
let cloudflare_cooldown_until = account.cloudflare_cooldown_until.map(|dt| dt.to_rfc3339());
let mut tx = self.pool.begin().await?;
let result = sqlx::query(SYNC_RUNTIME_ACCOUNT_STATE_SQL)
.bind(status_to_db(account.status))
.bind(quota_limit_reached)
.bind(&now)
.bind(quota_limit_reached)
.bind(quota_verify_required)
.bind(&now)
.bind(quota_verify_required)
.bind(&quota_cooldown_until)
.bind(&now)
.bind(&quota_cooldown_until)
.bind(&cloudflare_cooldown_until)
.bind(&now)
.bind(&cloudflare_cooldown_until)
.bind(&now)
.bind(&account.id)
.execute(&mut *tx)
.await?;
if result.rows_affected() == 0 {
tx.commit().await?;
return Ok(false);
}
if sync_usage_window {
sqlx::query(SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL)
.bind(&account.id)
.bind(u64_to_i64_saturating(account.window_request_count))
.bind(u64_to_i64_saturating(account.window_input_tokens))
.bind(u64_to_i64_saturating(account.window_output_tokens))
.bind(u64_to_i64_saturating(account.window_cached_tokens))
.bind(u64_to_i64_saturating(account.window_image_input_tokens))
.bind(u64_to_i64_saturating(account.window_image_output_tokens))
.bind(u64_to_i64_saturating(account.window_image_request_count))
.bind(u64_to_i64_saturating(
account.window_image_request_failed_count,
))
.bind(account.window_started_at.map(|dt| dt.to_rfc3339()))
.bind(account.window_reset_at.map(|dt| dt.to_rfc3339()))
.bind(account.limit_window_seconds.map(u64_to_i64_saturating))
.bind(&now)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
Ok(true)
}
/// 标记账号 Cloudflare 冷却期。
pub async fn set_cloudflare_cooldown_until(
&self,
@@ -1082,6 +1149,16 @@ impl AccountStore for SqliteAccountStore {
.map_err(|error| map_account_store_error(&error))
}
async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> AccountStoreResult<bool> {
SqliteAccountStore::sync_runtime_account_state(self, account, sync_usage_window)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn sync_rate_limit_window(
&self,
account_id: &str,
backend/src/upstream/accounts/store/queries.rs
+57
View File
@@ -432,6 +432,63 @@ set
updated_at = ?
where id = ?";
pub(super) const SYNC_RUNTIME_ACCOUNT_STATE_SQL: &str = r"
update accounts
set
status = case
when status in ('disabled', 'banned', 'refreshing') then status
else ?
end,
quota_limit_reached = case
when ? = 0 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_limit_reached
else ?
end,
quota_verify_required = case
when ? = 1 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_verify_required
else ?
end,
quota_cooldown_until = case
when ? is null and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_cooldown_until
else ?
end,
cloudflare_cooldown_until = case
when ? is null and cloudflare_cooldown_until is not null and cloudflare_cooldown_until > ? then cloudflare_cooldown_until
else ?
end,
updated_at = ?
where id = ?";
pub(super) const SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL: &str = r"
insert into account_usage (
account_id,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_image_input_tokens,
window_image_output_tokens,
window_image_request_count,
window_image_request_failed_count,
window_started_at,
window_reset_at,
limit_window_seconds
) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
on conflict(account_id) do update set
window_request_count = excluded.window_request_count,
window_input_tokens = excluded.window_input_tokens,
window_output_tokens = excluded.window_output_tokens,
window_cached_tokens = excluded.window_cached_tokens,
window_image_input_tokens = excluded.window_image_input_tokens,
window_image_output_tokens = excluded.window_image_output_tokens,
window_image_request_count = excluded.window_image_request_count,
window_image_request_failed_count = excluded.window_image_request_failed_count,
window_started_at = excluded.window_started_at,
window_reset_at = excluded.window_reset_at,
limit_window_seconds = excluded.limit_window_seconds
where account_usage.window_reset_at is null
or account_usage.window_reset_at <= excluded.window_reset_at
or account_usage.window_reset_at <= ?";
pub(super) const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r"
update accounts
set
backend/tests/http/trace_middleware/mod.rs
+103 -36
View File
@@ -11,10 +11,12 @@ use axum::{
http::{Request, StatusCode},
middleware::{from_fn, from_fn_with_state},
routing::get,
Router,
Extension, Router,
};
use codex_proxy_rs::http::middleware::{
request_id::{attach_request_id, attach_request_id_with_proxy_config, TrustedProxyConfig},
request_id::{
attach_request_id, attach_request_id_with_proxy_config, ClientIp, TrustedProxyConfig,
},
trace::http_trace_layer,
};
use tokio::time::{sleep, Duration};
@@ -219,16 +221,12 @@ async fn request_context_should_generate_new_request_id_for_invalid_header() {
}
#[tokio::test]
async fn request_context_should_attach_real_ip_from_connection() {
async fn request_context_should_attach_client_ip_from_connection() {
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string()
get(|Extension(client_ip): Extension<ClientIp>| async move {
client_ip.as_str().to_string()
}),
)
.layer(from_fn(attach_request_id));
@@ -247,17 +245,30 @@ async fn request_context_should_attach_real_ip_from_connection() {
}
#[tokio::test]
async fn request_context_should_replace_untrusted_forwarded_ip_headers() {
async fn request_context_should_preserve_forwarded_headers_from_untrusted_peer() {
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing")
.to_string()
}),
get(
|Extension(client_ip): Extension<ClientIp>, headers: HeaderMap| async move {
let cf_ip = headers
.get("cf-connecting-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing");
let forwarded_for = headers
.get("x-forwarded-for")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing");
let real_ip = headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing");
format!(
"{}|cf={cf_ip}|xff={forwarded_for}|real={real_ip}",
client_ip.as_str()
)
},
),
)
.layer(from_fn(attach_request_id));
let mut request = Request::builder()
@@ -274,7 +285,10 @@ async fn request_context_should_replace_untrusted_forwarded_ip_headers() {
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "127.0.0.1");
assert_eq!(
std::str::from_utf8(&bytes).unwrap(),
"127.0.0.1|cf=203.0.113.43|xff=203.0.113.42, 10.0.0.2|real=203.0.113.44"
);
}
#[tokio::test]
@@ -284,14 +298,12 @@ async fn request_context_should_accept_forwarded_ip_from_trusted_proxy() {
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
let real_ip = headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing");
let forwarded_present = headers.contains_key("x-forwarded-for");
format!("{real_ip}|forwarded={forwarded_present}")
}),
get(
|Extension(client_ip): Extension<ClientIp>, headers: HeaderMap| async move {
let forwarded_present = headers.contains_key("x-forwarded-for");
format!("{}|forwarded={forwarded_present}", client_ip.as_str())
},
),
)
.layer(from_fn_with_state(
trusted_proxies,
@@ -311,23 +323,19 @@ async fn request_context_should_accept_forwarded_ip_from_trusted_proxy() {
assert_eq!(
std::str::from_utf8(&bytes).unwrap(),
"203.0.113.42|forwarded=false"
"203.0.113.42|forwarded=true"
);
}
#[tokio::test]
async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() {
async fn request_context_should_accept_real_ip_from_trusted_proxy_when_forwarded_for_missing() {
let trusted_proxies =
TrustedProxyConfig::from_entries(&["127.0.0.1".to_string()]).expect("trusted proxy");
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing")
.to_string()
get(|Extension(client_ip): Extension<ClientIp>| async move {
client_ip.as_str().to_string()
}),
)
.layer(from_fn_with_state(
@@ -336,8 +344,7 @@ async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() {
));
let mut request = Request::builder()
.uri("/ip")
.header("cf-connecting-ip", "198.51.100.23")
.header("x-forwarded-for", "203.0.113.42, 10.0.0.2")
.header("x-real-ip", "198.51.100.23")
.body(Body::empty())
.expect("request should build");
request
@@ -350,6 +357,66 @@ async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() {
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "198.51.100.23");
}
#[tokio::test]
async fn request_context_should_ignore_forwarded_ip_from_untrusted_peer() {
let trusted_proxies =
TrustedProxyConfig::from_entries(&["10.0.0.0/8".to_string()]).expect("trusted proxy");
let app = Router::new()
.route(
"/ip",
get(|Extension(client_ip): Extension<ClientIp>| async move {
client_ip.as_str().to_string()
}),
)
.layer(from_fn_with_state(
trusted_proxies,
attach_request_id_with_proxy_config,
));
let mut request = Request::builder()
.uri("/ip")
.header("x-forwarded-for", "203.0.113.42")
.body(Body::empty())
.expect("request should build");
request
.extensions_mut()
.insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234))));
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "127.0.0.1");
}
#[tokio::test]
async fn request_context_should_not_inject_real_ip_header() {
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing")
.to_string()
}),
)
.layer(from_fn(attach_request_id));
let mut request = Request::builder()
.uri("/ip")
.header("x-forwarded-for", "203.0.113.42, 10.0.0.2")
.header("cf-connecting-ip", "203.0.113.43")
.body(Body::empty())
.expect("request should build");
request
.extensions_mut()
.insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234))));
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "missing");
}
#[test]
fn trusted_proxy_config_should_reject_invalid_entries() {
let error = TrustedProxyConfig::from_entries(&["not-an-ip".to_string()]).unwrap_err();
backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs
+86
View File
@@ -133,6 +133,92 @@ async fn responses_websocket_stream_should_synthesize_response_failed_when_close
assert!(body.contains("stream_disconnected"));
}
#[tokio::test]
async fn responses_websocket_stream_first_error_429_should_retry_fallback_account() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base_url = format!("http://{}", listener.local_addr().unwrap());
let upstream = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first_websocket =
accept_websocket_with_authorization(first_stream, "Bearer access-primary").await;
let first_payload = send_websocket_response_and_capture_payload(
&mut first_websocket,
json!({
"type": "error",
"error": {
"type": "usage_limit_reached",
"message": "The usage limit has been reached",
"retry_after_seconds": 77
},
"status_code": 429
})
.to_string(),
)
.await;
first_websocket.close(None).await.unwrap();
let second_payload = accept_successful_websocket_response_with_authorization(
&listener,
"Bearer access-secondary",
"resp_after_ws_stream_rate_limit",
)
.await;
(first_payload, second_payload)
});
let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await;
let response = app
.oneshot(responses_json_request(
&api_key,
&json!({
"model": "gpt-5.5",
"input": [],
"stream": true,
"use_websocket": true
}),
))
.await
.unwrap();
let status = response.status();
let content_type = response
.headers()
.get("content-type")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let body = response_text(response).await;
let (first_payload, second_payload) = upstream.await.unwrap();
let primary_quota_state: (i64, Option<String>) = sqlx::query_as(
"select quota_limit_reached, quota_cooldown_until from accounts where id = ?",
)
.bind("acct_primary")
.fetch_one(&pool)
.await
.unwrap();
let primary_usage: (i64,) =
sqlx::query_as("select request_count from account_usage where account_id = ?")
.bind("acct_primary")
.fetch_one(&pool)
.await
.unwrap();
let secondary_usage: (i64,) =
sqlx::query_as("select request_count from account_usage where account_id = ?")
.bind("acct_secondary")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(status, StatusCode::OK);
assert!(content_type.starts_with("text/event-stream"));
assert!(body.contains("resp_after_ws_stream_rate_limit"));
assert_eq!(first_payload["type"], "response.create");
assert_eq!(second_payload["type"], "response.create");
assert_eq!(primary_quota_state.0, 1);
assert!(primary_quota_state.1.is_some());
assert_eq!(primary_usage.0, 1);
assert_eq!(secondary_usage.0, 1);
}
#[tokio::test]
async fn responses_websocket_should_reuse_connection_for_recorded_conversation() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
backend/tests/runtime/tasks/mod.rs
+8
View File
@@ -97,6 +97,14 @@ impl AccountStore for FakeAccountStore {
Ok(false)
}
async fn sync_runtime_account_state(
&self,
_account: &Account,
_sync_usage_window: bool,
) -> AccountStoreResult<bool> {
Ok(false)
}
async fn sync_rate_limit_window(
&self,
_account_id: &str,
backend/tests/runtime/tasks/model_refresh.rs
+8
View File
@@ -210,6 +210,14 @@ impl AccountStore for PlanAccountStore {
Ok(false)
}
async fn sync_runtime_account_state(
&self,
_account: &Account,
_sync_usage_window: bool,
) -> AccountStoreResult<bool> {
Ok(false)
}
async fn sync_rate_limit_window(
&self,
_account_id: &str,
backend/tests/upstream/accounts/account_pool/selection.rs
+124
View File
@@ -210,6 +210,130 @@ async fn runtime_account_pool_should_persist_expired_status_when_jwt_expiry_is_d
assert_eq!(status.0, "expired");
}
#[tokio::test]
async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_expires() {
let now = fixed_time();
let (pool, _dir) =
crate::support::sqlite::init_test_db("runtime-pool-quota-window-reset.sqlite").await;
let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone());
store
.insert(codex_proxy_rs::upstream::accounts::store::NewAccount {
id: "acct_quota_reset".to_string(),
email: None,
account_id: Some("chatgpt-quota-reset".to_string()),
user_id: None,
label: None,
plan_type: Some("plus".to_string()),
access_token: secrecy::SecretString::new("access-token".to_string().into()),
refresh_token: None,
access_token_expires_at: Some(now + Duration::hours(1)),
status: AccountStatus::QuotaExhausted,
added_at: None,
})
.await
.unwrap();
sqlx::query(
r#"
update accounts
set
quota_limit_reached = 1,
quota_verify_required = 0,
quota_cooldown_until = ?,
cloudflare_cooldown_until = ?
where id = ?
"#,
)
.bind((now - Duration::seconds(1)).to_rfc3339())
.bind((now - Duration::seconds(1)).to_rfc3339())
.bind("acct_quota_reset")
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"
insert into account_usage (
account_id,
request_count,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_started_at,
window_reset_at,
limit_window_seconds
) values (?, 9, 7, 11, 13, 17, ?, ?, 60)
"#,
)
.bind("acct_quota_reset")
.bind((now - Duration::seconds(120)).to_rfc3339())
.bind((now - Duration::seconds(1)).to_rfc3339())
.execute(&pool)
.await
.unwrap();
let runtime_pool = codex_proxy_rs::upstream::accounts::pool::RuntimeAccountPoolService::new(
std::sync::Arc::new(store),
AccountPoolOptions::default(),
0,
);
runtime_pool.restore_from_repository().await.unwrap();
let acquired = runtime_pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.await
.unwrap();
let stored: (
String,
i64,
i64,
Option<String>,
Option<String>,
i64,
i64,
i64,
i64,
Option<String>,
Option<i64>,
) = sqlx::query_as(
r#"
select
a.status,
a.quota_limit_reached,
a.quota_verify_required,
a.quota_cooldown_until,
a.cloudflare_cooldown_until,
au.window_request_count,
au.window_input_tokens,
au.window_output_tokens,
au.window_cached_tokens,
au.window_reset_at,
au.limit_window_seconds
from accounts a
left join account_usage au on au.account_id = a.id
where a.id = ?
"#,
)
.bind("acct_quota_reset")
.fetch_one(&pool)
.await
.unwrap();
let persisted_reset_at = stored
.9
.as_deref()
.and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok())
.map(|value| value.with_timezone(&Utc));
assert_eq!(acquired.account.id, "acct_quota_reset");
assert!(acquired.account.quota_verify_required);
assert_eq!(stored.0, "active");
assert_eq!(stored.1, 0);
assert_eq!(stored.2, 1);
assert!(stored.3.is_none());
assert!(stored.4.is_none());
assert_eq!((stored.5, stored.6, stored.7, stored.8), (0, 0, 0, 0));
assert!(persisted_reset_at.is_some_and(|reset_at| reset_at > now));
assert_eq!(stored.10, Some(60));
}
#[test]
fn account_pool_should_prefer_configured_tier_priority() {
let mut pool = AccountPool::with_options(AccountPoolOptions {
backend/tests/upstream/accounts/account_repository.rs
+93
View File
@@ -518,6 +518,99 @@ async fn account_repository_should_map_quota_snapshot_to_account_status() {
assert!(restored_status.2.is_none());
}
#[tokio::test]
async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quota_cooldown() {
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 42).await;
let repo = SqliteAccountStore::new(pool.clone());
seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await;
let future_cooldown = Utc::now() + Duration::hours(1);
sqlx::query(
"update accounts set quota_limit_reached = 1, quota_verify_required = 0, quota_cooldown_until = ? where id = ?",
)
.bind(future_cooldown.to_rfc3339())
.bind("acct_a")
.execute(&pool)
.await
.unwrap();
let mut stale_account = AccountStore::list_pool_accounts(&repo)
.await
.unwrap()
.into_iter()
.find(|account| account.id == "acct_a")
.unwrap();
stale_account.quota_limit_reached = false;
stale_account.quota_verify_required = true;
stale_account.quota_cooldown_until = None;
repo.sync_runtime_account_state(&stale_account, false)
.await
.unwrap();
let stored: (i64, i64, Option<String>) = sqlx::query_as(
"select quota_limit_reached, quota_verify_required, quota_cooldown_until from accounts where id = ?",
)
.bind("acct_a")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(stored.0, 1);
assert_eq!(stored.1, 0);
assert_eq!(stored.2, Some(future_cooldown.to_rfc3339()));
}
#[tokio::test]
async fn account_repository_runtime_state_sync_should_not_regress_newer_usage_window() {
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 43).await;
let repo = SqliteAccountStore::new(pool.clone());
seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await;
let now = Utc::now();
let future_reset = now + Duration::hours(1);
let stale_reset = now + Duration::minutes(1);
sqlx::query(
r#"
insert into account_usage (
account_id,
window_request_count,
window_input_tokens,
window_reset_at,
limit_window_seconds
) values (?, 5, 17, ?, 3600)
"#,
)
.bind("acct_a")
.bind(future_reset.to_rfc3339())
.execute(&pool)
.await
.unwrap();
let mut stale_account = AccountStore::list_pool_accounts(&repo)
.await
.unwrap()
.into_iter()
.find(|account| account.id == "acct_a")
.unwrap();
stale_account.window_request_count = 0;
stale_account.window_input_tokens = 0;
stale_account.window_reset_at = Some(stale_reset);
stale_account.window_started_at = Some(now);
stale_account.limit_window_seconds = Some(60);
repo.sync_runtime_account_state(&stale_account, true)
.await
.unwrap();
let stored: (i64, i64, Option<String>, Option<i64>) = sqlx::query_as(
"select window_request_count, window_input_tokens, window_reset_at, limit_window_seconds from account_usage where account_id = ?",
)
.bind("acct_a")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(stored.0, 5);
assert_eq!(stored.1, 17);
assert_eq!(stored.2, Some(future_reset.to_rfc3339()));
assert_eq!(stored.3, Some(3600));
}
#[tokio::test]
async fn account_repository_should_not_override_disabled_status_from_quota_snapshot() {
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 38).await;