diff --git a/CHANGELOG.md b/CHANGELOG.md index ae6a8eda9cb..0ad8e212d71 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ## [Windows] Unreleased +- Kimi (China): use the current Kimi Desktop session from its Local Storage (`%APPDATA%\kimi-desktop`) for web usage and monthly membership enrichment, and check that session again before each request, so a signed-out, replaced, or expired Desktop token is never sent or recovered from an older record (upstream #4416, #3536). - Codex CLI: add opt-in `cost --identified-only` request subtotals across local and SSH (`--remote`) histories, deduplicating strong request identities and keeping partial results without estimating a combined price, plus `--group-by project` reports that show the partial-history notice even when no projects are listed (upstream #4447, #4472). - OpenCode Go: reuse unchanged local history scans (keyed by `opencode.db` and its `-wal` sidecar size and modification time, at most one 4 MiB payload kept), and in Auto overlay Console quota from a stored or automatically discovered browser session before falling back to the local estimate; cookie source Off never reads a browser (upstream #4410, #4445). - Add **Parallel** prepaid balance: reads the account balance and any pending debit from the platform.parallel.ai console using your browser session or a pasted Cookie header. A signed-out console moves on to the next browser session and reports an expired session when none remain (upstream #4479, #4508). diff --git a/rust/src/browser/leveldb/mod.rs b/rust/src/browser/leveldb/mod.rs index ea5bbefd139..074ebee70a6 100644 --- a/rust/src/browser/leveldb/mod.rs +++ b/rust/src/browser/leveldb/mod.rs @@ -16,6 +16,10 @@ //! - Files larger than [`MAX_FILE_BYTES`], scans larger than [`MAX_TOTAL_SCAN_BYTES`], and blocks //! that decode beyond [`MAX_BLOCK_BYTES`] are skipped so a corrupt profile cannot exhaust //! memory or monopolize a refresh. +//! +//! [`read_current_value_strict`] is the stricter single-key variant for credentials whose stale +//! value must never be resurrected: it fails instead of skipping any file or block, and it +//! rejects symlinked entries and a directory that changes while it is read. pub mod local_storage; mod log; @@ -66,6 +70,8 @@ pub enum LevelDbError { Io(#[from] std::io::Error), #[error("LevelDB directory exceeds safe scan limits")] ResourceLimit, + #[error("LevelDB directory could not be read completely and consistently")] + Incomplete, } /// A put (`value: Some`) or delete (`value: None`) with the sequence number it was written at. @@ -225,6 +231,115 @@ pub(crate) fn read_entries_with_budget( Ok((entries.into_entries(), scanned_bytes)) } +/// Newest live value of exactly `key` in the LevelDB directory `dir`, or `None` when the key is +/// absent or its newest record is a delete. +/// +/// A skipped file or block could hold a newer record or tombstone for `key`, so unlike +/// [`read_entries`] every problem fails the whole read with [`LevelDbError::Incomplete`]: a +/// missing `CURRENT`, a symlinked entry, an oversized or unreadable file, a malformed table or +/// data block, a torn log record, or a directory listing whose names, sizes or modification +/// times differ after the read. The `MANIFEST` is still not consulted and checksums are not +/// verified, so this narrows rather than eliminates the window for a concurrent writer. +pub(crate) fn read_current_value_strict( + dir: &Path, + key: &[u8], +) -> Result>, LevelDbError> { + let before = directory_fingerprint(dir)?; + if !before.iter().any(|file| file.name == "CURRENT") { + return Err(LevelDbError::Incomplete); + } + let mut entries = EntryAccumulator { + key_prefix: Some(key), + ..EntryAccumulator::default() + }; + let mut scanned_bytes = 0u64; + for file in &before { + let path = dir.join(&file.name); + let Some(kind) = FileKind::of(&path) else { + continue; + }; + let remaining_bytes = MAX_TOTAL_SCAN_BYTES.saturating_sub(scanned_bytes); + let data = match open_no_follow(&path) + .map_err(BoundedReadError::Io) + .and_then(|file| read_bounded_open_file(file, remaining_bytes)) + { + Ok(data) => data, + Err(BoundedReadError::BudgetExceeded) => return Err(LevelDbError::ResourceLimit), + Err(BoundedReadError::FileTooLarge | BoundedReadError::Io(_)) => { + return Err(LevelDbError::Incomplete); + } + }; + let file_bytes = u64::try_from(data.len()).map_err(|_| LevelDbError::ResourceLimit)?; + scanned_bytes = scanned_bytes + .checked_add(file_bytes) + .ok_or(LevelDbError::ResourceLimit)?; + let complete = match kind { + FileKind::Log => log::read_log_until(&data, &mut |record| entries.absorb(record)), + FileKind::Table => matches!( + table::read_table_until(&data, &mut |record| entries.absorb(record)), + Ok((0, true)) + ), + }; + if entries.resource_limit_exceeded { + return Err(LevelDbError::ResourceLimit); + } + if !complete { + return Err(LevelDbError::Incomplete); + } + } + if directory_fingerprint(dir)? != before { + return Err(LevelDbError::Incomplete); + } + Ok(entries.newest.remove(key).and_then(|(_, value)| value)) +} + +#[derive(PartialEq, Eq)] +struct FileFingerprint { + name: std::ffi::OsString, + len: u64, + modified: Option, +} + +fn directory_fingerprint(dir: &Path) -> Result, LevelDbError> { + let mut files = Vec::new(); + for (count, dir_entry) in std::fs::read_dir(dir)?.enumerate() { + if count == MAX_DIRECTORY_ENTRIES { + return Err(LevelDbError::ResourceLimit); + } + let dir_entry = dir_entry.map_err(|_| LevelDbError::Incomplete)?; + // `DirEntry::metadata` does not traverse links; junctions count as symlinks on Windows. + let metadata = dir_entry.metadata().map_err(|_| LevelDbError::Incomplete)?; + if metadata.file_type().is_symlink() { + return Err(LevelDbError::Incomplete); + } + files.push(FileFingerprint { + name: dir_entry.file_name(), + len: metadata.len(), + modified: metadata.modified().ok(), + }); + } + files.sort_by(|left, right| left.name.cmp(&right.name)); + Ok(files) +} + +/// Open `path` for reading without following a link swapped in after the directory listing. +fn open_no_follow(path: &Path) -> std::io::Result { + let mut options = std::fs::OpenOptions::new(); + options.read(true); + #[cfg(windows)] + { + use std::os::windows::fs::OpenOptionsExt; + const FILE_FLAG_OPEN_REPARSE_POINT: u32 = 0x0020_0000; + options.custom_flags(FILE_FLAG_OPEN_REPARSE_POINT); + } + let file = options.open(path)?; + let metadata = file.metadata()?; + if metadata.file_type().is_symlink() || !metadata.is_file() { + return Err(std::io::Error::other("not a regular file")); + } + Ok(file) +} + #[derive(Clone, Copy)] enum FileKind { Log, @@ -252,6 +367,13 @@ enum BoundedReadError { fn read_bounded_file(path: &Path, remaining_bytes: u64) -> Result, BoundedReadError> { let file = std::fs::File::open(path).map_err(BoundedReadError::Io)?; + read_bounded_open_file(file, remaining_bytes) +} + +fn read_bounded_open_file( + file: std::fs::File, + remaining_bytes: u64, +) -> Result, BoundedReadError> { let size = file.metadata().map_err(BoundedReadError::Io)?.len(); if size > MAX_FILE_BYTES { return Err(BoundedReadError::FileTooLarge); diff --git a/rust/src/browser/leveldb/tests.rs b/rust/src/browser/leveldb/tests.rs index 7110e9d2159..cb46b1c5166 100644 --- a/rust/src/browser/leveldb/tests.rs +++ b/rust/src/browser/leveldb/tests.rs @@ -11,7 +11,7 @@ use super::local_storage::{ }; use super::snappy::{self, SnappyError}; use super::table::{TableError, read_table}; -use super::{Entry, Record, log, read_entries}; +use super::{Entry, LevelDbError, Record, log, read_current_value_strict, read_entries}; // --------------------------------------------------------------------------------------------- // Fixture builders @@ -519,6 +519,87 @@ fn read_entries_reports_a_missing_directory() { assert!(read_entries(&dir.path().join("absent")).is_err()); } +fn strict_dir(files: &[(&str, Vec)]) -> tempfile::TempDir { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("CURRENT"), b"MANIFEST-000001\n").unwrap(); + for (name, data) in files { + std::fs::write(dir.path().join(name), data).unwrap(); + } + dir +} + +#[test] +fn strict_read_returns_the_newest_value_of_the_exact_key() { + let table: &[TableRecord] = &[(b"token", 1, 1, b"old"), (b"token-other", 2, 1, b"x")]; + let dir = strict_dir(&[ + ("000002.ldb", write_table(&[(table, Compression::None)])), + ( + "000003.log", + write_log(&[write_batch(5, &[Op::Put(b"token", b"new")])]), + ), + ]); + assert_eq!( + read_current_value_strict(dir.path(), b"token").unwrap(), + Some(b"new".to_vec()) + ); + assert_eq!( + read_current_value_strict(dir.path(), b"absent").unwrap(), + None + ); +} + +#[test] +fn strict_read_honors_a_newer_delete() { + let dir = strict_dir(&[( + "000003.log", + write_log(&[ + write_batch(1, &[Op::Put(b"token", b"old")]), + write_batch(2, &[Op::Delete(b"token")]), + ]), + )]); + assert_eq!( + read_current_value_strict(dir.path(), b"token").unwrap(), + None + ); +} + +#[test] +fn strict_read_fails_instead_of_skipping_a_file() { + let good = write_log(&[write_batch(1, &[Op::Put(b"token", b"old")])]); + let mut torn = write_log(&[write_batch(2, &[Op::Delete(b"token")])]); + torn.truncate(torn.len() - 1); + let cases = [ + vec![ + ("000003.log", good.clone()), + ("000004.ldb", b"not a table".to_vec()), + ], + vec![("000003.log", good.clone()), ("000004.log", torn)], + ]; + for files in cases { + let dir = strict_dir(&files); + assert!(matches!( + read_current_value_strict(dir.path(), b"token"), + Err(LevelDbError::Incomplete) + )); + // The best-effort reader still returns the stale value from the readable file. + assert_eq!(read_entries(dir.path()).unwrap().len(), 1); + } +} + +#[test] +fn strict_read_requires_current() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("000003.log"), + write_log(&[write_batch(1, &[Op::Put(b"token", b"v")])]), + ) + .unwrap(); + assert!(matches!( + read_current_value_strict(dir.path(), b"token"), + Err(LevelDbError::Incomplete) + )); +} + // --------------------------------------------------------------------------------------------- // Chromium Local Storage layer // --------------------------------------------------------------------------------------------- diff --git a/rust/src/providers/kimi/code_api.rs b/rust/src/providers/kimi/code_api.rs index d98b7b50844..3d5e91ac76a 100644 --- a/rust/src/providers/kimi/code_api.rs +++ b/rust/src/providers/kimi/code_api.rs @@ -2,7 +2,8 @@ //! upstream 0.48.0 enrichment (#2622): Code API / CLI usage snapshots are //! merged with the monthly membership pool + Code 7-day limit when a //! `kimi.com` web token is available (manual cookie → Kimi Desktop session → -//! browser import; gated by Cookie Source Off). +//! browser import; gated by Cookie Source Off). A Desktop session is checked +//! again before each enrichment request (upstream #4416). use reqwest::Url; use std::path::{Path, PathBuf}; @@ -85,16 +86,19 @@ pub(crate) async fn fetch_via_code_api( snapshot.login_method = Some(plan_name.unwrap_or_else(|| login_method.to_string())); // Upstream #2622: enrich Code API + CLI usage with the monthly membership - // pool from a signed-in Kimi Desktop (or browser/manual) session. - for web_token in web::web_auth_tokens(ctx.manual_cookie_header.as_deref(), region) { - match web::fetch_subscription_for_enrichment_result(&client, &web_token, region).await { + // pool from a signed-in Kimi Desktop (or browser/manual) session. A + // Desktop session that changed before dispatch ends enrichment but keeps + // the Code quota snapshot (upstream #4416). + for web_credential in web::web_auth_credentials(ctx.manual_cookie_header.as_deref(), region) { + match web::fetch_subscription_for_enrichment_result(&client, &web_credential, region).await + { Ok(subscription) => { if let Some(subscription) = subscription { snapshot = super::apply_subscription_windows(snapshot, &subscription); } if !has_plan_name && let Some(plan) = - web::fetch_subscription_plan(&client, &web_token, region).await + web::fetch_subscription_plan(&client, &web_credential, region).await { snapshot.login_method = Some(plan); } diff --git a/rust/src/providers/kimi/credential.rs b/rust/src/providers/kimi/credential.rs new file mode 100644 index 00000000000..2dbe620bff2 --- /dev/null +++ b/rust/src/providers/kimi/credential.rs @@ -0,0 +1,255 @@ +//! A `kimi.com` web token together with where it came from (upstream `KimiWebCredential`, +//! #4416 / #3536). +//! +//! A token discovered from the current Kimi Desktop session keeps a revalidation callback until +//! each request is dispatched: the Desktop session is read again right before the request is +//! sent, and a missing, replaced, expired or unreadable session rejects it. Tokens from any other +//! source keep their existing behavior. + +use std::sync::{Arc, Mutex}; + +use reqwest::header::{AUTHORIZATION, COOKIE}; + +use super::KimiRegion; +use crate::core::ProviderError; + +pub(super) const DESKTOP_SESSION_CHANGED_MESSAGE: &str = "Kimi Desktop session changed before the request was sent. \ + Confirm the intended account is signed in to Kimi Desktop, then refresh CodexBar."; + +#[derive(Clone)] +pub(super) struct KimiWebCredential { + token: String, + desktop_check: Option>, +} + +impl KimiWebCredential { + pub(super) fn plain(token: impl Into) -> Self { + Self { + token: token.into(), + desktop_check: None, + } + } + + /// A Desktop session token; `is_current` must re-read the session and report whether it + /// still holds exactly this token. + pub(super) fn desktop( + token: impl Into, + is_current: impl Fn() -> bool + Send + Sync + 'static, + ) -> Self { + Self { + token: token.into(), + desktop_check: Some(Arc::new(DesktopDispatchValidation { + valid: Mutex::new(true), + is_current: Box::new(is_current), + })), + } + } + + pub(super) fn token(&self) -> &str { + &self.token + } + + /// Gate `request` before it is handed to the HTTP client. Only Desktop credentials are + /// checked: the request must be the China web API request carrying this token, and the + /// Desktop session must still hold it. This narrows the window between discovery and + /// dispatch; it cannot recall a request that was already sent. + pub(super) fn check_dispatch( + &self, + request: &reqwest::Request, + region: KimiRegion, + ) -> Result<(), ProviderError> { + let Some(validation) = &self.desktop_check else { + return Ok(()); + }; + if region == KimiRegion::China + && self.is_bound_request(request, region) + && validation.check() + { + Ok(()) + } else { + Err(ProviderError::Other(DESKTOP_SESSION_CHANGED_MESSAGE.into())) + } + } + + fn is_bound_request(&self, request: &reqwest::Request, region: KimiRegion) -> bool { + let url = request.url(); + let expected_host = region + .web_base_url() + .strip_prefix("https://") + .unwrap_or_default(); + let header = |name| { + request + .headers() + .get(name) + .and_then(|value| value.to_str().ok()) + }; + url.scheme() == "https" + && url.host_str() == Some(expected_host) + && url.port().is_none_or(|port| port == 443) + && header(AUTHORIZATION) == Some(format!("Bearer {}", self.token).as_str()) + && header(COOKIE) == Some(format!("kimi-auth={}", self.token).as_str()) + } +} + +impl PartialEq for KimiWebCredential { + fn eq(&self, other: &Self) -> bool { + self.token == other.token && self.desktop_check.is_some() == other.desktop_check.is_some() + } +} + +impl Eq for KimiWebCredential {} + +impl std::fmt::Debug for KimiWebCredential { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("KimiWebCredential") + .field("token", &"[REDACTED]") + .field("desktop", &self.desktop_check.is_some()) + .finish() + } +} + +/// Shared by every clone of one acquired credential, so concurrent requests share invalidation: +/// once a check fails, restoring the old session value cannot revive this credential. +struct DesktopDispatchValidation { + valid: Mutex, + is_current: Box bool + Send + Sync>, +} + +impl DesktopDispatchValidation { + fn check(&self) -> bool { + let mut valid = self + .valid + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if *valid { + *valid = (self.is_current)(); + } + *valid + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + fn request(client: &reqwest::Client, url: &str, token: &str) -> reqwest::Request { + super::super::kimi_web_request(client, url, KimiRegion::China, token, serde_json::json!({})) + .unwrap() + } + + fn china_url() -> String { + KimiRegion::China.web_api_url("svc") + } + + #[test] + fn plain_credentials_are_never_gated() { + let client = reqwest::Client::new(); + let credential = KimiWebCredential::plain("browser-token"); + let other = request(&client, "https://example.com/x", "something-else"); + assert!( + credential + .check_dispatch(&other, KimiRegion::International) + .is_ok() + ); + } + + #[test] + fn current_desktop_session_is_rechecked_before_every_request() { + let client = reqwest::Client::new(); + let checks = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&checks); + let credential = KimiWebCredential::desktop("desktop-token", move || { + counter.fetch_add(1, Ordering::SeqCst); + true + }); + let request = request(&client, &china_url(), "desktop-token"); + for _ in 0..3 { + credential + .check_dispatch(&request, KimiRegion::China) + .unwrap(); + } + assert_eq!(checks.load(Ordering::SeqCst), 3); + } + + #[test] + fn changed_session_rejects_the_request_and_stays_rejected() { + let client = reqwest::Client::new(); + let current = Arc::new(AtomicBool::new(true)); + let session = Arc::clone(¤t); + let credential = + KimiWebCredential::desktop("desktop-token", move || session.load(Ordering::SeqCst)); + let concurrent = credential.clone(); + let request = request(&client, &china_url(), "desktop-token"); + + assert!( + credential + .check_dispatch(&request, KimiRegion::China) + .is_ok() + ); + current.store(false, Ordering::SeqCst); + let error = credential + .check_dispatch(&request, KimiRegion::China) + .unwrap_err(); + assert_eq!(error.to_string(), DESKTOP_SESSION_CHANGED_MESSAGE); + assert!(!error.to_string().contains("desktop-token")); + + // Restoring the old value cannot revive this credential or any clone of it. + current.store(true, Ordering::SeqCst); + assert!( + credential + .check_dispatch(&request, KimiRegion::China) + .is_err() + ); + assert!( + concurrent + .check_dispatch(&request, KimiRegion::China) + .is_err() + ); + } + + #[test] + fn desktop_token_is_bound_to_the_china_web_api_request() { + let client = reqwest::Client::new(); + let credential = KimiWebCredential::desktop("desktop-token", || true); + let cases = [ + ( + request(&client, &china_url(), "other-token"), + KimiRegion::China, + ), + ( + request(&client, "https://evil.example/apiv2/svc", "desktop-token"), + KimiRegion::China, + ), + ( + request(&client, "http://www.kimi.com/apiv2/svc", "desktop-token"), + KimiRegion::China, + ), + ( + request( + &client, + "https://www.kimi.com:8443/apiv2/svc", + "desktop-token", + ), + KimiRegion::China, + ), + ( + request(&client, &china_url(), "desktop-token"), + KimiRegion::International, + ), + ]; + for (request, region) in cases { + assert!(credential.check_dispatch(&request, region).is_err()); + } + // A rejected binding is a session error too; it does not consult or poison the session. + let bound = request(&client, &china_url(), "desktop-token"); + assert!(credential.check_dispatch(&bound, KimiRegion::China).is_ok()); + } + + #[test] + fn debug_output_redacts_the_token() { + let rendered = format!("{:?}", KimiWebCredential::desktop("secret-token", || true)); + assert!(!rendered.contains("secret-token")); + } +} diff --git a/rust/src/providers/kimi/desktop_session.rs b/rust/src/providers/kimi/desktop_session.rs new file mode 100644 index 00000000000..de594465b54 --- /dev/null +++ b/rust/src/providers/kimi/desktop_session.rs @@ -0,0 +1,503 @@ +//! Current Kimi Desktop session discovery (upstream `KimiDesktopSessionDiscovery`, #4416 / +//! #3536). +//! +//! For China, the Desktop session is the exact `https://www.kimi.com` / `access_token` record in +//! the app's primary Local Storage, `%APPDATA%\kimi-desktop\Local Storage\leveldb` on Windows +//! (upstream reads `~/Library/Application Support/kimi-desktop/...`). It is read with the strict +//! LevelDB reader, so a deleted (signed-out) or replaced session is never recovered from an older +//! record, and only a bare, unexpired access JWT for the `kimi.com` audience is accepted. China +//! never falls back to the historical Desktop `kimi-auth` cookie. International keeps its +//! existing Desktop cookie discovery; its Local Storage is not read. +//! +//! The returned credential re-reads the session before each request (see +//! [`KimiWebCredential::check_dispatch`]). Tokens are secrets and are never logged. + +use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use base64::Engine; +use base64::engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig}; + +use super::KimiRegion; +use super::credential::KimiWebCredential; +use super::desktop_token::KimiDesktopAuthToken; +use crate::browser::leveldb::read_current_value_strict; + +const DESKTOP_APP_DIR: &str = "kimi-desktop"; +const DESKTOP_ORIGIN: &str = "https://www.kimi.com"; +const ACCESS_TOKEN_KEY: &str = "access_token"; +const REQUIRED_AUDIENCE: &str = "kimi.com"; +const LATIN1_FORMAT: u8 = 1; +/// Upstream bound: a Latin-1 format byte plus at most 16 KiB of token text. +const MAX_RAW_VALUE_BYTES: usize = 16 * 1024 + 1; + +const JWT_SEGMENT: GeneralPurpose = GeneralPurpose::new( + &base64::alphabet::URL_SAFE, + GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent), +); + +/// The automatic Kimi Desktop web credential for `region`, or `None` when there is no usable +/// session. Callers gate this on the automatic cookie source. +pub(super) fn desktop_credential(region: KimiRegion) -> Option { + match region { + KimiRegion::China => DesktopSession::new(dirs::data_dir()?, unix_now).credential(), + KimiRegion::International => { + KimiDesktopAuthToken::load_for_region(region).map(KimiWebCredential::plain) + } + } +} + +#[derive(Clone)] +struct DesktopSession { + /// Roaming AppData (Electron `userData` parent). Only this host-owned root may be a link; + /// the fixed profile path below it must not be. + data_root: PathBuf, + now: fn() -> f64, +} + +impl DesktopSession { + fn new(data_root: PathBuf, now: fn() -> f64) -> Self { + Self { data_root, now } + } + + fn credential(self) -> Option { + let token = self.access_token()?; + let expected = token.clone(); + Some(KimiWebCredential::desktop(token, move || { + self.access_token().as_deref() == Some(expected.as_str()) + })) + } + + fn access_token(&self) -> Option { + let dir = self.local_storage_dir()?; + let raw = read_current_value_strict(&dir, &raw_key()) + .inspect_err(|error| { + tracing::debug!(%error, "Kimi Desktop session is not readable"); + }) + .ok()??; + access_token_from_raw(&raw, (self.now)()) + } + + fn local_storage_dir(&self) -> Option { + let mut dir = self.data_root.clone(); + for component in [DESKTOP_APP_DIR, "Local Storage", "leveldb"] { + dir.push(component); + if !is_plain_directory(&dir) { + return None; + } + } + Some(dir) + } +} + +/// `symlink_metadata` does not traverse links; on Windows junctions report as symlinks too. +fn is_plain_directory(path: &Path) -> bool { + std::fs::symlink_metadata(path) + .is_ok_and(|metadata| metadata.is_dir() && !metadata.file_type().is_symlink()) +} + +/// Chromium Local Storage key: `_\0` then a Latin-1 format byte and the item key. +fn raw_key() -> Vec { + let mut key = format!("_{DESKTOP_ORIGIN}\0").into_bytes(); + key.push(LATIN1_FORMAT); + key.extend_from_slice(ACCESS_TOKEN_KEY.as_bytes()); + key +} + +/// Upstream `accessToken(rawValue:now:)`: a Latin-1 bare JWT (no JSON wrapper) with a non-`none` +/// algorithm, `typ` `access`, a finite numeric expiry after `now`, the `kimi.com` audience, and an +/// optional not-before at or before `now`. These are structural checks only; the Kimi service +/// still authenticates the signature. +fn access_token_from_raw(raw: &[u8], now: f64) -> Option { + let (&format, body) = raw.split_first()?; + if format != LATIN1_FORMAT + || raw.len() > MAX_RAW_VALUE_BYTES + || body.is_empty() + || !now.is_finite() + { + return None; + } + if !body + .iter() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.')) + { + return None; + } + let token = std::str::from_utf8(body).ok()?; + let segments: Vec<&str> = token.split('.').collect(); + let [header, claims, signature] = segments.as_slice() else { + return None; + }; + if header.is_empty() || claims.is_empty() || signature.is_empty() { + return None; + } + + let header = decode_object(header)?; + let algorithm = header.get("alg")?.as_str()?; + if algorithm.is_empty() || algorithm.eq_ignore_ascii_case("none") { + return None; + } + + let claims = decode_object(claims)?; + if claims.get("typ")?.as_str()? != "access" { + return None; + } + if numeric_date(claims.get("exp")?)? <= now { + return None; + } + if !audiences(claims.get("aud")).contains(&REQUIRED_AUDIENCE) { + return None; + } + if let Some(not_before) = claims.get("nbf") + && numeric_date(not_before)? > now + { + return None; + } + Some(token.to_string()) +} + +fn decode_object(segment: &str) -> Option> { + let bytes = JWT_SEGMENT.decode(segment).ok()?; + match serde_json::from_slice(&bytes).ok()? { + serde_json::Value::Object(object) => Some(object), + _ => None, + } +} + +/// JSON numbers only: booleans and numeric strings are not dates. +fn numeric_date(value: &serde_json::Value) -> Option { + value.as_f64().filter(|date| date.is_finite()) +} + +/// `aud` as a string or an array made only of strings; anything else has no audience. +fn audiences(value: Option<&serde_json::Value>) -> Vec<&str> { + match value { + Some(serde_json::Value::String(audience)) => vec![audience.as_str()], + Some(serde_json::Value::Array(items)) => items + .iter() + .map(serde_json::Value::as_str) + .collect::>>() + .unwrap_or_default(), + _ => Vec::new(), + } +} + +fn unix_now() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(f64::NAN, |elapsed| elapsed.as_secs_f64()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicU64, Ordering}; + + const NOW: f64 = 1_800_000_000.0; + + fn fixed_now() -> f64 { + NOW + } + + fn segment(value: &serde_json::Value) -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(value.to_string()) + } + + /// Synthetic JWT built from readable claims; no real credential. + fn jwt(claims: serde_json::Value) -> String { + format!( + "{}.{}.signature", + segment(&serde_json::json!({ "alg": "HS256", "typ": "JWT" })), + segment(&claims) + ) + } + + fn access_claims(sub: &str) -> serde_json::Value { + serde_json::json!({ + "typ": "access", + "aud": ["kimi.com"], + "exp": NOW + 3600.0, + "nbf": NOW - 60.0, + "sub": sub, + }) + } + + fn latin1(token: &str) -> Vec { + let mut raw = vec![LATIN1_FORMAT]; + raw.extend_from_slice(token.as_bytes()); + raw + } + + fn accepted(raw: &[u8]) -> bool { + access_token_from_raw(raw, NOW).is_some() + } + + #[test] + fn current_access_jwt_is_accepted_bare() { + let token = jwt(access_claims("a")); + assert_eq!(access_token_from_raw(&latin1(&token), NOW), Some(token)); + let single_audience = jwt(serde_json::json!({ + "typ": "access", "aud": "kimi.com", "exp": NOW + 1.0 + })); + assert!(accepted(&latin1(&single_audience))); + } + + #[test] + fn wrapped_encoded_or_oversized_values_are_rejected() { + let token = jwt(access_claims("a")); + let mut utf16 = vec![0u8]; + utf16.extend(token.bytes().flat_map(|byte| [byte, 0])); + let oversized = format!("{}.{}.{}", "a".repeat(8 * 1024), "b".repeat(8 * 1024), "c"); + for raw in [ + latin1(&format!("\"{token}\"")), + latin1(&format!(" {token}")), + latin1(&format!("{token}; kimi-auth=x")), + utf16, + latin1(""), + latin1("a.b"), + latin1("a..c"), + latin1(&oversized), + Vec::new(), + ] { + assert!(!accepted(&raw)); + } + } + + #[test] + fn claims_policy_rejects_refresh_expired_and_lookalike_tokens() { + let base = access_claims("a"); + let with = |key: &str, value: serde_json::Value| { + let mut claims = base.clone(); + claims[key] = value; + latin1(&jwt(claims)) + }; + let without = |key: &str| { + let mut claims = base.clone(); + claims.as_object_mut().unwrap().remove(key); + latin1(&jwt(claims)) + }; + for raw in [ + with("typ", serde_json::json!("refresh")), + without("typ"), + with("exp", serde_json::json!(NOW)), + with("exp", serde_json::json!(NOW - 1.0)), + with("exp", serde_json::json!("1900000000")), + with("exp", serde_json::json!(true)), + without("exp"), + with("aud", serde_json::json!(["kimi.com.evil"])), + with("aud", serde_json::json!("www.kimi.com")), + with("aud", serde_json::json!(["kimi.com", 7])), + without("aud"), + with("nbf", serde_json::json!(NOW + 60.0)), + with("nbf", serde_json::json!("soon")), + ] { + assert!(!accepted(&raw)); + } + assert!(accepted(&without("nbf"))); + assert!(access_token_from_raw(&latin1(&jwt(base)), f64::NAN).is_none()); + } + + #[test] + fn unsigned_or_headerless_tokens_are_rejected() { + let claims = segment(&access_claims("a")); + for header in [ + serde_json::json!({ "alg": "none" }), + serde_json::json!({ "alg": "NONE" }), + serde_json::json!({ "alg": "" }), + serde_json::json!({ "typ": "JWT" }), + serde_json::json!(["HS256"]), + ] { + let token = format!("{}.{claims}.signature", segment(&header)); + assert!(!accepted(&latin1(&token))); + } + } + + /// One LevelDB log record holding the Desktop `access_token` put (or delete when `None`). + fn write_log(sequence: u64, token: Option<&str>) -> Vec { + write_log_for(&raw_key(), sequence, token) + } + + fn write_log_for(key: &[u8], sequence: u64, token: Option<&str>) -> Vec { + let mut batch = sequence.to_le_bytes().to_vec(); + batch.extend_from_slice(&1u32.to_le_bytes()); + batch.push(u8::from(token.is_some())); + batch.push(u8::try_from(key.len()).unwrap()); + batch.extend_from_slice(key); + if let Some(token) = token { + let value = latin1(token); + let mut length = Vec::new(); + let mut remaining = value.len(); + while remaining >= 0x80 { + length.push(u8::try_from(remaining & 0x7f).unwrap() | 0x80); + remaining >>= 7; + } + length.push(u8::try_from(remaining).unwrap()); + batch.extend_from_slice(&length); + batch.extend_from_slice(&value); + } + let mut record = vec![0, 0, 0, 0]; + record.extend_from_slice(&u16::try_from(batch.len()).unwrap().to_le_bytes()); + record.push(1); + record.extend_from_slice(&batch); + record + } + + fn profile_root() -> (tempfile::TempDir, PathBuf) { + let root = tempfile::tempdir().unwrap(); + let dir = root + .path() + .join(DESKTOP_APP_DIR) + .join("Local Storage") + .join("leveldb"); + std::fs::create_dir_all(&dir).unwrap(); + std::fs::write(dir.join("CURRENT"), b"MANIFEST-000001\n").unwrap(); + (root, dir) + } + + fn write_records(dir: &Path, records: &[(u64, Option<&str>)]) { + let log: Vec = records + .iter() + .flat_map(|(sequence, token)| write_log(*sequence, *token)) + .collect(); + std::fs::write(dir.join("000003.log"), log).unwrap(); + } + + fn session(root: &Path) -> DesktopSession { + DesktopSession::new(root.to_path_buf(), fixed_now) + } + + #[test] + fn newest_desktop_session_is_used_and_stale_tokens_are_never_recovered() { + let (root, dir) = profile_root(); + let old = jwt(access_claims("old")); + let current = jwt(access_claims("current")); + + write_records(&dir, &[(1, Some(&old)), (2, Some(¤t))]); + assert_eq!(session(root.path()).access_token(), Some(current.clone())); + + // Signed out: the deletion hides every older record. + write_records(&dir, &[(1, Some(&old)), (2, Some(¤t)), (3, None)]); + assert_eq!(session(root.path()).access_token(), None); + + // The current record is expired: no older live token is recovered either. + let expired = jwt(serde_json::json!({ + "typ": "access", "aud": "kimi.com", "exp": NOW - 1.0 + })); + write_records(&dir, &[(1, Some(&old)), (2, Some(&expired))]); + assert_eq!(session(root.path()).access_token(), None); + } + + #[test] + fn other_origins_and_incomplete_profiles_yield_no_session() { + let (root, dir) = profile_root(); + let token = jwt(access_claims("a")); + let key_for = |origin: &str, item: &str| { + let mut key = format!("_{origin}\0").into_bytes(); + key.push(LATIN1_FORMAT); + key.extend_from_slice(item.as_bytes()); + key + }; + let mut log = write_log_for( + &key_for("https://www.kimi.ai", ACCESS_TOKEN_KEY), + 1, + Some(&token), + ); + log.extend(write_log_for( + &key_for(DESKTOP_ORIGIN, "refresh_token"), + 2, + Some(&token), + )); + log.extend(write_log_for( + &key_for(DESKTOP_ORIGIN, "access_token_v2"), + 3, + Some(&token), + )); + std::fs::write(dir.join("000003.log"), log).unwrap(); + assert_eq!(session(root.path()).access_token(), None); + + // An unreadable table could hide a newer deletion, so the session is not trusted. + write_records(&dir, &[(1, Some(&token))]); + assert_eq!(session(root.path()).access_token(), Some(token)); + std::fs::write(dir.join("000004.ldb"), b"truncated table").unwrap(); + assert_eq!(session(root.path()).access_token(), None); + + let missing = tempfile::tempdir().unwrap(); + assert_eq!(session(missing.path()).access_token(), None); + } + + #[cfg(windows)] + #[test] + fn junctioned_profile_directory_is_rejected() { + let (target_root, target_dir) = profile_root(); + write_records(&target_dir, &[(1, Some(&jwt(access_claims("a"))))]); + assert!(session(target_root.path()).access_token().is_some()); + + let root = tempfile::tempdir().unwrap(); + let link = root.path().join(DESKTOP_APP_DIR); + let status = std::process::Command::new("cmd") + .args(["/C", "mklink", "/J"]) + .arg(&link) + .arg(target_root.path().join(DESKTOP_APP_DIR)) + .output() + .unwrap(); + assert!(status.status.success(), "mklink /J failed"); + assert!(session(root.path()).access_token().is_none()); + } + + static CHANGING_NOW: AtomicU64 = AtomicU64::new(0); + + fn changing_now() -> f64 { + f64::from_bits(CHANGING_NOW.load(Ordering::SeqCst)) + } + + #[test] + fn acquired_credential_revalidates_the_session_before_dispatch() { + CHANGING_NOW.store(NOW.to_bits(), Ordering::SeqCst); + let (root, dir) = profile_root(); + let current = jwt(access_claims("current")); + let replacement = jwt(access_claims("replacement")); + write_records(&dir, &[(1, Some(¤t))]); + + let client = reqwest::Client::new(); + let url = KimiRegion::China.web_api_url("svc"); + let check = |credential: &KimiWebCredential| { + let request = super::super::kimi_web_request( + &client, + &url, + KimiRegion::China, + credential.token(), + serde_json::json!({}), + ) + .unwrap(); + credential.check_dispatch(&request, KimiRegion::China) + }; + + let credential = DesktopSession::new(root.path().to_path_buf(), changing_now) + .credential() + .expect("current session"); + assert_eq!(credential.token(), current); + assert!(check(&credential).is_ok()); + + // Expired by the time of dispatch. + CHANGING_NOW.store((NOW + 7200.0).to_bits(), Ordering::SeqCst); + assert!(check(&credential).is_err()); + CHANGING_NOW.store(NOW.to_bits(), Ordering::SeqCst); + assert!(check(&credential).is_err(), "invalidation is sticky"); + + // Replaced by another account: the old credential is rejected; a fresh read uses the new one. + let credential = DesktopSession::new(root.path().to_path_buf(), changing_now) + .credential() + .unwrap(); + write_records(&dir, &[(1, Some(¤t)), (2, Some(&replacement))]); + assert!(check(&credential).is_err()); + let fresh = DesktopSession::new(root.path().to_path_buf(), changing_now) + .credential() + .unwrap(); + assert_eq!(fresh.token(), replacement); + assert!(check(&fresh).is_ok()); + + // Signed out after acquisition. + write_records(&dir, &[(1, Some(&replacement)), (2, None)]); + assert!(check(&fresh).is_err()); + } +} diff --git a/rust/src/providers/kimi/desktop_token.rs b/rust/src/providers/kimi/desktop_token.rs index ac40f793ba2..fa822de01f5 100644 --- a/rust/src/providers/kimi/desktop_token.rs +++ b/rust/src/providers/kimi/desktop_token.rs @@ -16,6 +16,10 @@ //! is preferred (upstream reads only `value`); the encrypted form falls //! back to the existing `browser::cookies` decryption helpers. //! +//! Since upstream #4416 this cookie store is read only for International; China +//! uses the current Desktop Local Storage session (`desktop_session`) and never +//! falls back to a historical China Desktop cookie. +//! //! Auth cookies are secrets: token values are never logged. use base64::Engine; diff --git a/rust/src/providers/kimi/mod.rs b/rust/src/providers/kimi/mod.rs index 12da5943c13..cd7df8ef54c 100755 --- a/rust/src/providers/kimi/mod.rs +++ b/rust/src/providers/kimi/mod.rs @@ -12,8 +12,11 @@ //! from a signed-in Kimi Desktop session (#2622). //! - [`local_storage`]: current `access_token` JWTs from Chromium local storage //! for the selected region, read through the shared LevelDB reader. +//! - [`desktop_session`]: the current China Kimi Desktop session from its +//! Local Storage, revalidated before each request through +//! [`credential::KimiWebCredential`] (upstream #4416 / #3536). //! - [`desktop_token`]: read-only, WAL-safe reader for the Kimi Desktop -//! (Electron) Chromium cookie store. +//! (Electron) Chromium cookie store, still used for International. //! - [`ratio_pool`]: zero-ratio placeholder reconciliation against matching //! legacy counters (upstream 0.63.0). //! - [`auto`]: Auto-mode order after the Code API key (CLI credential, then @@ -21,6 +24,8 @@ mod auto; mod code_api; +mod credential; +mod desktop_session; pub mod desktop_token; mod local_storage; #[cfg(test)] @@ -31,6 +36,8 @@ mod web; pub use region::KimiRegion; +use credential::KimiWebCredential; + use async_trait::async_trait; use chrono::{DateTime, Utc}; use reqwest::Client; @@ -545,9 +552,21 @@ async fn kimi_web_post( client: &Client, url: &str, region: KimiRegion, - token: &str, + credential: &KimiWebCredential, body: serde_json::Value, ) -> Result { + let request = kimi_web_request(client, url, region, credential.token(), body)?; + credential.check_dispatch(&request, region)?; + client.execute(request).await.map_err(ProviderError::from) +} + +fn kimi_web_request( + client: &Client, + url: &str, + region: KimiRegion, + token: &str, + body: serde_json::Value, +) -> Result { client .post(url) .bearer_auth(token) @@ -561,8 +580,7 @@ async fn kimi_web_post( "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", ) .json(&body) - .send() - .await + .build() .map_err(ProviderError::from) } diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index c6c1b8f59f0..585e25052dd 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -6,11 +6,15 @@ //! manual cookie header → Kimi Desktop session → browser cookie import → //! Chromium local-storage `access_token` (upstream `KimiWebEnrichmentTokenResolver` //! and #3923), used both by the web fetch itself and -//! by the Code-API/CLI monthly enrichment. +//! by the Code-API/CLI monthly enrichment. A Desktop session keeps its +//! provenance as a [`KimiWebCredential`] and is revalidated before each +//! request (upstream #4416); a changed session stops the chain instead of +//! falling back to another credential. use reqwest::Client; -use super::desktop_token::KimiDesktopAuthToken; +use super::credential::KimiWebCredential; +use super::desktop_session::desktop_credential; use super::local_storage::local_storage_tokens; use super::{ KIMI_SUBSCRIPTION_SERVICE, KIMI_SUBSCRIPTION_STATS_SERVICE, KIMI_WEB_USAGE_SERVICE, @@ -67,17 +71,20 @@ impl WebFetchFailure { /// 2. Kimi Desktop session token (automatic source only). /// 3. Browser cookie import (automatic source only). /// 4. Chromium local-storage `access_token` for the region (automatic source only). -pub(crate) fn web_auth_tokens(manual_header: Option<&str>, region: KimiRegion) -> Vec { +pub(super) fn web_auth_credentials( + manual_header: Option<&str>, + region: KimiRegion, +) -> Vec { resolve_web_tokens(&WebTokenInput { manual_header, cookie_source: &cookie_source(), region, - desktop_token: KimiDesktopAuthToken::load_for_region, + desktop_token: desktop_credential, browser_token: browser_auth_token, local_storage_tokens, }) .into_iter() - .map(|candidate| candidate.token) + .map(|candidate| candidate.credential) .collect() } @@ -85,7 +92,7 @@ struct WebTokenInput<'a> { manual_header: Option<&'a str>, cookie_source: &'a str, region: KimiRegion, - desktop_token: fn(KimiRegion) -> Option, + desktop_token: fn(KimiRegion) -> Option, browser_token: fn(KimiRegion) -> Option, local_storage_tokens: fn(KimiRegion) -> Vec, } @@ -98,28 +105,18 @@ enum WebTokenSource { LocalStorage, } -#[derive(PartialEq, Eq)] +#[derive(Debug, PartialEq, Eq)] struct WebTokenCandidate { - token: String, + credential: KimiWebCredential, source: WebTokenSource, } -impl std::fmt::Debug for WebTokenCandidate { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter - .debug_struct("WebTokenCandidate") - .field("token", &"[REDACTED]") - .field("source", &self.source) - .finish() - } -} - fn resolve_web_tokens(input: &WebTokenInput<'_>) -> Vec { if let Some(header) = input.manual_header && let Ok(token) = KimiProvider::auth_token_from_cookie_header(header) { return vec![WebTokenCandidate { - token, + credential: KimiWebCredential::plain(token), source: WebTokenSource::Manual, }]; } @@ -129,11 +126,11 @@ fn resolve_web_tokens(input: &WebTokenInput<'_>) -> Vec { let mut candidates = Vec::new(); let mut seen = std::collections::HashSet::new(); - if let Some(token) = (input.desktop_token)(input.region) - && seen.insert(token.clone()) + if let Some(credential) = (input.desktop_token)(input.region) + && seen.insert(credential.token().to_string()) { candidates.push(WebTokenCandidate { - token, + credential, source: WebTokenSource::Desktop, }); } @@ -141,14 +138,14 @@ fn resolve_web_tokens(input: &WebTokenInput<'_>) -> Vec { && seen.insert(token.clone()) { candidates.push(WebTokenCandidate { - token, + credential: KimiWebCredential::plain(token), source: WebTokenSource::Browser, }); } for token in (input.local_storage_tokens)(input.region) { if seen.insert(token.clone()) { candidates.push(WebTokenCandidate { - token, + credential: KimiWebCredential::plain(token), source: WebTokenSource::LocalStorage, }); } @@ -212,7 +209,7 @@ pub(super) async fn fetch_web_session_isolated( }); } }; - return fetch_via_web_token(&http, &token, region) + return fetch_via_web_token(&http, &KimiWebCredential::plain(token), region) .await .map_err(WebFetchFailure::after_token); } @@ -229,18 +226,18 @@ pub(super) async fn fetch_web_session( manual_header: cookie_header, cookie_source: &source, region, - desktop_token: KimiDesktopAuthToken::load_for_region, + desktop_token: desktop_credential, browser_token: browser_auth_token, - local_storage_tokens: crate::providers::kimi::local_storage::local_storage_tokens, + local_storage_tokens, }; // One HTTP client for every token attempt, built on first use. let mut shared_client: Option = None; - fetch_with_web_tokens(&input, |token| { + fetch_with_web_tokens(&input, |credential| { let http = match &shared_client { Some(http) => Ok(http.clone()), None => client().inspect(|http| shared_client = Some(http.clone())), }; - async move { fetch_via_web_token(&http?, &token, region).await } + async move { fetch_via_web_token(&http?, &credential, region).await } }) .await } @@ -248,13 +245,14 @@ pub(super) async fn fetch_web_session( /// The web fetch over the token chain. An explicit manual token is /// authoritative; otherwise (automatic source only) the Kimi Desktop session /// is tried, then browser import. Only a server rejection moves on to the -/// next automatic token. +/// next automatic token; a Desktop session that changed before dispatch is a +/// non-auth error and stops the chain. async fn fetch_with_web_tokens( input: &WebTokenInput<'_>, mut fetch: F, ) -> Result where - F: FnMut(String) -> Fut, + F: FnMut(KimiWebCredential) -> Fut, Fut: Future>, { if let Some(token) = input @@ -263,7 +261,9 @@ where { // An explicit manual credential is authoritative. A rejected manual // token must not silently switch accounts underneath the user. - return fetch(token).await.map_err(WebFetchFailure::after_token); + return fetch(KimiWebCredential::plain(token)) + .await + .map_err(WebFetchFailure::after_token); } if !browser_import_allowed(input.cookie_source) { @@ -277,18 +277,19 @@ where // read only after the server rejects this automatic session, so a healthy // desktop account never causes another credential store to be touched. let desktop_token = (input.desktop_token)(input.region); - if let Some(token) = desktop_token.clone() { - match fetch(token).await { + if let Some(credential) = desktop_token.clone() { + match fetch(credential).await { Ok(usage) => return Ok(usage), Err(ProviderError::AuthRequired) => {} Err(error) => return Err(WebFetchFailure::after_token(error)), } } - let browser_token = - (input.browser_token)(input.region).filter(|token| desktop_token.as_ref() != Some(token)); + let browser_token = (input.browser_token)(input.region).filter(|token| { + desktop_token.as_ref().map(KimiWebCredential::token) != Some(token.as_str()) + }); if let Some(token) = browser_token.clone() { - match fetch(token).await { + match fetch(KimiWebCredential::plain(token)).await { Ok(usage) => return Ok(usage), Err(ProviderError::AuthRequired) => {} Err(error) => return Err(WebFetchFailure::after_token(error)), @@ -298,16 +299,17 @@ where // Local-storage tokens are read last, only after every cookie source // was rejected. let mut seen = std::collections::HashSet::new(); - if let Some(token) = desktop_token.clone() { - seen.insert(token); + if let Some(credential) = &desktop_token { + seen.insert(credential.token().to_string()); } if let Some(token) = browser_token.clone() { seen.insert(token); } for candidate in resolve_web_tokens(input).into_iter().filter(|candidate| { - candidate.source == WebTokenSource::LocalStorage && seen.insert(candidate.token.clone()) + candidate.source == WebTokenSource::LocalStorage + && seen.insert(candidate.credential.token().to_string()) }) { - match fetch(candidate.token).await { + match fetch(candidate.credential).await { Ok(usage) => return Ok(usage), Err(ProviderError::AuthRequired) => {} Err(error) => return Err(WebFetchFailure::after_token(error)), @@ -339,7 +341,7 @@ fn client() -> Result { async fn fetch_via_web_token( client: &reqwest::Client, - token: &str, + credential: &KimiWebCredential, region: KimiRegion, ) -> Result { let usage_url = region.web_api_url(KIMI_WEB_USAGE_SERVICE); @@ -347,7 +349,7 @@ async fn fetch_via_web_token( client, &usage_url, region, - token, + credential, serde_json::json!({ "scope": ["FEATURE_CODING"] }), ) .await?; @@ -365,7 +367,7 @@ async fn fetch_via_web_token( .await .map_err(|e| ProviderError::Parse(e.to_string()))?; - let (subscription, plan_name) = fetch_subscription_details(client, token, region).await; + let (subscription, plan_name) = fetch_subscription_details(client, credential, region).await; snapshot_from_web_usage_response_with_plan(usage, subscription, plan_name) } @@ -374,18 +376,18 @@ const SUBSCRIPTION_ENRICHMENT_TIMEOUT: std::time::Duration = std::time::Duration async fn fetch_subscription_details( client: &reqwest::Client, - token: &str, + credential: &KimiWebCredential, region: KimiRegion, ) -> (Option, Option) { // The quota statistics and the optional title are independent. Keep a // completed statistics response when the plan endpoint is slow or absent. let stats = tokio::time::timeout( SUBSCRIPTION_ENRICHMENT_TIMEOUT, - fetch_subscription_for_enrichment(client, token, region), + fetch_subscription_for_enrichment(client, credential, region), ); let plan = tokio::time::timeout( SUBSCRIPTION_ENRICHMENT_TIMEOUT, - fetch_subscription_plan(client, token, region), + fetch_subscription_plan(client, credential, region), ); let (stats, plan) = tokio::join!(stats, plan); (stats.ok().flatten(), plan.ok().flatten()) @@ -430,11 +432,11 @@ fn snapshot_from_web_usage_response_with_plan( pub(super) async fn fetch_subscription_plan( client: &Client, - token: &str, + credential: &KimiWebCredential, region: KimiRegion, ) -> Option { let url = region.web_api_url(KIMI_SUBSCRIPTION_SERVICE); - match kimi_web_post(client, &url, region, token, serde_json::json!({})).await { + match kimi_web_post(client, &url, region, credential, serde_json::json!({})).await { Ok(response) if response.status().is_success() => response .json::() .await @@ -448,10 +450,10 @@ pub(super) async fn fetch_subscription_plan( // token; any failure means "no enrichment", never an error. pub(super) async fn fetch_subscription_for_enrichment( client: &Client, - token: &str, + credential: &KimiWebCredential, region: KimiRegion, ) -> Option { - fetch_subscription_for_enrichment_result(client, token, region) + fetch_subscription_for_enrichment_result(client, credential, region) .await .ok() .flatten() @@ -459,11 +461,11 @@ pub(super) async fn fetch_subscription_for_enrichment( pub(super) async fn fetch_subscription_for_enrichment_result( client: &Client, - token: &str, + credential: &KimiWebCredential, region: KimiRegion, ) -> Result, ProviderError> { let url = region.web_api_url(KIMI_SUBSCRIPTION_STATS_SERVICE); - match kimi_web_post(client, &url, region, token, serde_json::json!({})).await { + match kimi_web_post(client, &url, region, credential, serde_json::json!({})).await { Ok(response) if response.status().is_success() => response .json() .await @@ -497,18 +499,30 @@ mod tests { ); } - fn static_desktop(_: KimiRegion) -> Option { - Some("desktop-token".to_string()) + fn static_desktop(_: KimiRegion) -> Option { + Some(KimiWebCredential::plain("desktop-token")) + } + + fn changed_desktop(_: KimiRegion) -> Option { + Some(KimiWebCredential::desktop("desktop-token", || false)) } fn static_browser(_: KimiRegion) -> Option { Some("browser-token".to_string()) } + fn no_desktop(_: KimiRegion) -> Option { + None + } + fn no_token(_: KimiRegion) -> Option { None } + fn unread_desktop(_: KimiRegion) -> Option { + panic!("automatic Kimi token sources must not be read here") + } + fn unread(_: KimiRegion) -> Option { panic!("automatic Kimi token sources must not be read here") } @@ -546,21 +560,48 @@ mod tests { respond: fn(&str) -> Result, ) -> (Result, Vec) { let sent = std::cell::RefCell::new(Vec::new()); - let result = fetch_with_web_tokens(&input, |token| { - let response = respond(&token); - sent.borrow_mut().push(token); + let client = reqwest::Client::new(); + let url = KimiRegion::China.web_api_url(KIMI_WEB_USAGE_SERVICE); + let result = fetch_with_web_tokens(&input, |credential| { + let request = super::super::kimi_web_request( + &client, + &url, + KimiRegion::China, + credential.token(), + serde_json::json!({}), + ) + .unwrap(); + let response = credential + .check_dispatch(&request, KimiRegion::China) + .and_then(|()| { + sent.borrow_mut().push(credential.token().to_string()); + respond(credential.token()) + }); async move { response } }) .await; (result, sent.into_inner()) } + #[tokio::test] + async fn changed_desktop_session_is_not_sent_and_does_not_fall_back() { + let (result, sent) = + run_chain(input(None, "auto", changed_desktop, unread), accept_all).await; + let failure = result.expect_err("desktop session changed"); + assert!(failure.had_token); + assert_eq!( + failure.error.to_string(), + super::super::credential::DESKTOP_SESSION_CHANGED_MESSAGE + ); + assert!(sent.is_empty()); + } + #[tokio::test] async fn web_auth_without_a_token_reports_that_none_was_tried() { for source in ["off", "manual"] { for manual in [None, Some("not-a-token")] { let (result, sent) = - run_chain(input(manual, source, unread, unread), accept_all).await; + run_chain(input(manual, source, unread_desktop, unread), accept_all).await; let failure = result.expect_err("no web token to try"); assert!(!failure.had_token); assert_eq!( @@ -571,7 +612,7 @@ mod tests { } } - let (result, sent) = run_chain(input(None, "auto", no_token, no_token), accept_all).await; + let (result, sent) = run_chain(input(None, "auto", no_desktop, no_token), accept_all).await; let failure = result.expect_err("no automatic token found"); assert!(!failure.had_token); assert!(matches!(failure.error, ProviderError::AuthRequired)); @@ -581,7 +622,12 @@ mod tests { #[tokio::test] async fn rejected_manual_token_is_authoritative_and_counts_as_tried() { let (result, sent) = run_chain( - input(Some("kimi-auth=synthetic-web"), "auto", unread, unread), + input( + Some("kimi-auth=synthetic-web"), + "auto", + unread_desktop, + unread, + ), reject_all, ) .await; @@ -655,7 +701,7 @@ mod tests { fn input<'a>( manual_header: Option<&'a str>, cookie_source: &'a str, - desktop_token: fn(KimiRegion) -> Option, + desktop_token: fn(KimiRegion) -> Option, browser_token: fn(KimiRegion) -> Option, ) -> WebTokenInput<'a> { WebTokenInput { @@ -683,7 +729,7 @@ mod tests { assert_eq!( candidates, vec![WebTokenCandidate { - token: "manual-token".to_string(), + credential: KimiWebCredential::plain("manual-token"), source: WebTokenSource::Manual, }] ); @@ -695,9 +741,9 @@ mod tests { resolve_web_tokens(&input(None, "browser", static_desktop, static_browser)); assert_eq!(candidates.len(), 2); assert_eq!(candidates[0].source, WebTokenSource::Desktop); - assert_eq!(candidates[0].token, "desktop-token"); + assert_eq!(candidates[0].credential.token(), "desktop-token"); assert_eq!(candidates[1].source, WebTokenSource::Browser); - assert_eq!(candidates[1].token, "browser-token"); + assert_eq!(candidates[1].credential.token(), "browser-token"); } #[test] @@ -711,7 +757,7 @@ mod tests { resolve_web_tokens(&input( Some("not-a-token"), source, - no_token, + no_desktop, static_browser )), Vec::new() @@ -721,11 +767,11 @@ mod tests { #[test] fn browser_token_used_when_desktop_absent() { - let candidates = resolve_web_tokens(&input(None, "auto", no_token, static_browser)); + let candidates = resolve_web_tokens(&input(None, "auto", no_desktop, static_browser)); assert_eq!( candidates, vec![WebTokenCandidate { - token: "browser-token".to_string(), + credential: KimiWebCredential::plain("browser-token"), source: WebTokenSource::Browser, }] ); @@ -742,7 +788,7 @@ mod tests { assert_eq!( candidates, vec![WebTokenCandidate { - token: "manual-token".to_string(), + credential: KimiWebCredential::plain("manual-token"), source: WebTokenSource::Manual, }] ); @@ -755,7 +801,7 @@ mod tests { assert_eq!( candidates, vec![WebTokenCandidate { - token: "desktop-token".to_string(), + credential: KimiWebCredential::plain("desktop-token"), source: WebTokenSource::Desktop, }] ); @@ -769,7 +815,7 @@ mod tests { }); let ordered: Vec<_> = candidates .iter() - .map(|candidate| (candidate.source, candidate.token.as_str())) + .map(|candidate| (candidate.source, candidate.credential.token())) .collect(); assert_eq!( ordered, @@ -785,7 +831,7 @@ mod tests { fn local_storage_is_not_read_for_manual_credentials_or_blocked_sources() { let manual = resolve_web_tokens(&WebTokenInput { local_storage_tokens: static_local_storage, - ..input(Some("kimi-auth=manual-token"), "auto", no_token, no_token) + ..input(Some("kimi-auth=manual-token"), "auto", no_desktop, no_token) }); assert_eq!(manual.len(), 1); assert_eq!(manual[0].source, WebTokenSource::Manual); @@ -794,7 +840,7 @@ mod tests { assert!( resolve_web_tokens(&WebTokenInput { local_storage_tokens: static_local_storage, - ..input(None, source, no_token, no_token) + ..input(None, source, no_desktop, no_token) }) .is_empty() );