mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-08-04 19:45:34 +08:00
feat(pi): add native catalog and gateway data plane
This commit is contained in:
@@ -6,17 +6,23 @@
|
||||
|
||||
use crate::error::AppError;
|
||||
use indexmap::IndexMap;
|
||||
use jsonc_parser::cst::{CstContainerNode, CstNode, CstObject, CstObjectProp, CstRootNode};
|
||||
use jsonc_parser::cst::{
|
||||
CstArray, CstContainerNode, CstInputValue, CstLeafNode, CstNode, CstObject, CstObjectProp,
|
||||
CstRootNode,
|
||||
};
|
||||
use jsonc_parser::ParseOptions;
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::collections::HashMap;
|
||||
use std::fs::{self, File, Metadata, OpenOptions};
|
||||
use std::io::{Read, Take};
|
||||
use std::path::Path;
|
||||
use std::sync::LazyLock;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, LazyLock, Mutex, MutexGuard};
|
||||
|
||||
const MAX_PI_MODELS_BYTES: u64 = 8 * 1024 * 1024;
|
||||
const EMPTY_MODELS_DOCUMENT: &str = "{\"providers\":{}}";
|
||||
const MAX_MUTATION_ATTEMPTS: usize = 3;
|
||||
|
||||
static PI_JSON_LINE_COMMENTS: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r#""(?:\\.|[^"\\])*"|//[^\n]*"#).expect("Pi JSON line-comment regex must compile")
|
||||
@@ -25,6 +31,23 @@ static PI_JSON_TRAILING_COMMAS: LazyLock<Regex> = LazyLock::new(|| {
|
||||
Regex::new(r#""(?:\\.|[^"\\])*"|,(\s*[}\]])"#)
|
||||
.expect("Pi JSON trailing-comma regex must compile")
|
||||
});
|
||||
static PATH_LOCKS: LazyLock<Mutex<HashMap<PathBuf, Arc<Mutex<()>>>>> =
|
||||
LazyLock::new(|| Mutex::new(HashMap::new()));
|
||||
|
||||
fn path_lock(path: &Path) -> Result<Arc<Mutex<()>>, AppError> {
|
||||
let mut locks = PATH_LOCKS
|
||||
.lock()
|
||||
.map_err(|error| AppError::Config(format!("Pi path-lock registry is poisoned: {error}")))?;
|
||||
Ok(locks
|
||||
.entry(path.to_path_buf())
|
||||
.or_insert_with(|| Arc::new(Mutex::new(())))
|
||||
.clone())
|
||||
}
|
||||
|
||||
fn lock_path(lock: &Mutex<()>) -> Result<MutexGuard<'_, ()>, AppError> {
|
||||
lock.lock()
|
||||
.map_err(|error| AppError::Config(format!("Pi config path lock is poisoned: {error}")))
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct PiRawProviderEntry {
|
||||
@@ -111,6 +134,117 @@ fn cst_object(node: CstNode, path: &Path, label: &str) -> Result<CstObject, AppE
|
||||
}
|
||||
}
|
||||
|
||||
fn cst_input(value: &Value) -> CstInputValue {
|
||||
match value {
|
||||
Value::Null => CstInputValue::Null,
|
||||
Value::Bool(value) => CstInputValue::Bool(*value),
|
||||
Value::Number(value) => CstInputValue::Number(value.to_string()),
|
||||
Value::String(value) => CstInputValue::String(value.clone()),
|
||||
Value::Array(values) => CstInputValue::Array(values.iter().map(cst_input).collect()),
|
||||
Value::Object(values) => CstInputValue::Object(
|
||||
values
|
||||
.iter()
|
||||
.map(|(key, value)| (key.clone(), cst_input(value)))
|
||||
.collect(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn replace_cst_node(node: CstNode, replacement: &Value) -> Result<(), AppError> {
|
||||
let replacement = cst_input(replacement);
|
||||
let replaced = match node {
|
||||
CstNode::Container(CstContainerNode::Array(node)) => node.replace_with(replacement),
|
||||
CstNode::Container(CstContainerNode::Object(node)) => node.replace_with(replacement),
|
||||
CstNode::Leaf(CstLeafNode::BooleanLit(node)) => node.replace_with(replacement),
|
||||
CstNode::Leaf(CstLeafNode::NullKeyword(node)) => node.replace_with(replacement),
|
||||
CstNode::Leaf(CstLeafNode::NumberLit(node)) => node.replace_with(replacement),
|
||||
CstNode::Leaf(CstLeafNode::StringLit(node)) => node.replace_with(replacement),
|
||||
CstNode::Leaf(CstLeafNode::WordLit(node)) => node.replace_with(replacement),
|
||||
CstNode::Container(CstContainerNode::Root(_))
|
||||
| CstNode::Container(CstContainerNode::ObjectProp(_))
|
||||
| CstNode::Leaf(CstLeafNode::Token(_))
|
||||
| CstNode::Leaf(CstLeafNode::Whitespace(_))
|
||||
| CstNode::Leaf(CstLeafNode::Newline(_))
|
||||
| CstNode::Leaf(CstLeafNode::Comment(_)) => None,
|
||||
};
|
||||
replaced.map(|_| ()).ok_or_else(|| {
|
||||
AppError::Config("Pi models.json CST became disconnected during update".to_string())
|
||||
})
|
||||
}
|
||||
|
||||
fn patch_cst_object(
|
||||
object: &CstObject,
|
||||
before: &serde_json::Map<String, Value>,
|
||||
after: &serde_json::Map<String, Value>,
|
||||
) -> Result<(), AppError> {
|
||||
for key in before.keys().filter(|key| !after.contains_key(*key)) {
|
||||
let matching = object
|
||||
.properties()
|
||||
.into_iter()
|
||||
.filter(|property| cst_property_name(property).as_deref() == Some(key.as_str()))
|
||||
.collect::<Vec<_>>();
|
||||
for property in matching.into_iter().rev() {
|
||||
property.remove();
|
||||
}
|
||||
}
|
||||
|
||||
for (key, after_value) in after {
|
||||
if let Some(before_value) = before.get(key) {
|
||||
let property = last_cst_property(object, key).ok_or_else(|| {
|
||||
AppError::Config(format!(
|
||||
"Pi models.json CST is missing existing property '{key}'"
|
||||
))
|
||||
})?;
|
||||
let value = property.value().ok_or_else(|| {
|
||||
AppError::Config(format!("Pi models.json CST property '{key}' has no value"))
|
||||
})?;
|
||||
patch_cst_node(value, before_value, after_value)?;
|
||||
} else {
|
||||
object.append(key, cst_input(after_value));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn patch_cst_array(array: &CstArray, before: &[Value], after: &[Value]) -> Result<(), AppError> {
|
||||
let elements = array.elements();
|
||||
if elements.len() != before.len() {
|
||||
return Err(AppError::Config(
|
||||
"Pi models.json CST array does not match its parsed value".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
for (index, (before_value, after_value)) in before.iter().zip(after).enumerate() {
|
||||
patch_cst_node(elements[index].clone(), before_value, after_value)?;
|
||||
}
|
||||
for element in elements.into_iter().skip(after.len()).rev() {
|
||||
element.remove();
|
||||
}
|
||||
for value in after.iter().skip(before.len()) {
|
||||
array.append(cst_input(value));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn patch_cst_node(node: CstNode, before: &Value, after: &Value) -> Result<(), AppError> {
|
||||
if before == after {
|
||||
return Ok(());
|
||||
}
|
||||
match (&node, before, after) {
|
||||
(
|
||||
CstNode::Container(CstContainerNode::Object(object)),
|
||||
Value::Object(before),
|
||||
Value::Object(after),
|
||||
) => patch_cst_object(object, before, after),
|
||||
(
|
||||
CstNode::Container(CstContainerNode::Array(array)),
|
||||
Value::Array(before),
|
||||
Value::Array(after),
|
||||
) => patch_cst_array(array, before, after),
|
||||
_ => replace_cst_node(node, after),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_models_source(path: &Path, source: &str) -> Result<PiModelsDocument, AppError> {
|
||||
let document: Value = serde_json::from_str(&strip_pi_json_comments(source))
|
||||
.map_err(|error| AppError::json(path, error))?;
|
||||
@@ -256,6 +390,116 @@ pub(super) fn read_pi_models_document(path: &Path) -> Result<PiModelsDocument, A
|
||||
parse_models_source(path, source)
|
||||
}
|
||||
|
||||
fn fingerprint(bytes: Option<&[u8]>) -> Option<[u8; 32]> {
|
||||
bytes.map(|bytes| Sha256::digest(bytes).into())
|
||||
}
|
||||
|
||||
fn serialize_models_mutation(
|
||||
path: &Path,
|
||||
before: Option<&[u8]>,
|
||||
mutator: &impl Fn(&mut Value) -> Result<(), AppError>,
|
||||
) -> Result<Vec<u8>, AppError> {
|
||||
if let Some(bytes) = before {
|
||||
let source = std::str::from_utf8(bytes)
|
||||
.map_err(|error| jsonc_error(path, format!("file is not UTF-8: {error}")))?;
|
||||
let mut document: Value = serde_json::from_str(&strip_pi_json_comments(source))
|
||||
.map_err(|error| AppError::json(path, error))?;
|
||||
let root = CstRootNode::parse(source, &pi_models_parse_options())
|
||||
.map_err(|error| jsonc_error(path, error))?;
|
||||
if root.to_serde_value().as_ref() != Some(&document) {
|
||||
return Err(jsonc_error(path, "CST does not match Pi's parsed document"));
|
||||
}
|
||||
let original = document.clone();
|
||||
mutator(&mut document)?;
|
||||
let root_value = root
|
||||
.value()
|
||||
.ok_or_else(|| jsonc_error(path, "document must contain a JSON value"))?;
|
||||
patch_cst_node(root_value, &original, &document)?;
|
||||
if root.to_serde_value().as_ref() != Some(&document) {
|
||||
return Err(AppError::Config(
|
||||
"Pi models.json CST update did not produce the requested document".to_string(),
|
||||
));
|
||||
}
|
||||
return Ok(root.to_string().into_bytes());
|
||||
}
|
||||
|
||||
let mut document: Value = serde_json::from_str(EMPTY_MODELS_DOCUMENT)
|
||||
.expect("empty Pi models document is valid JSON");
|
||||
mutator(&mut document)?;
|
||||
let mut serialized = serde_json::to_vec_pretty(&document)
|
||||
.map_err(|source| AppError::JsonSerialize { source })?;
|
||||
serialized.push(b'\n');
|
||||
Ok(serialized)
|
||||
}
|
||||
|
||||
/// Patch only the explicitly named provider keys in Pi's shared models.json.
|
||||
///
|
||||
/// Unknown root fields, unowned provider entries, comments, and formatting are
|
||||
/// preserved by the CST patch. An optimistic fingerprint check prevents a
|
||||
/// Pi/user write observed before replacement from being silently overwritten.
|
||||
pub(crate) fn apply_pi_provider_patch(
|
||||
path: &Path,
|
||||
patch: &IndexMap<String, Option<Value>>,
|
||||
) -> Result<(), AppError> {
|
||||
let lock = path_lock(path)?;
|
||||
let _guard = lock_path(&lock)?;
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent).map_err(|error| AppError::io(parent, error))?;
|
||||
}
|
||||
|
||||
for _ in 0..MAX_MUTATION_ATTEMPTS {
|
||||
let before = read_models_bytes(path)?;
|
||||
let observed = fingerprint(before.as_deref());
|
||||
let serialized = serialize_models_mutation(path, before.as_deref(), &|document| {
|
||||
let providers = document
|
||||
.as_object_mut()
|
||||
.and_then(|root| root.get_mut("providers"))
|
||||
.and_then(Value::as_object_mut)
|
||||
.ok_or_else(|| jsonc_error(path, "root must contain a providers object"))?;
|
||||
for (provider_key, replacement) in patch {
|
||||
match replacement {
|
||||
Some(value) => {
|
||||
providers.insert(provider_key.clone(), value.clone());
|
||||
}
|
||||
None => {
|
||||
providers.remove(provider_key);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
})?;
|
||||
|
||||
let current = read_models_bytes(path)?;
|
||||
if fingerprint(current.as_deref()) != observed {
|
||||
continue;
|
||||
}
|
||||
crate::config::atomic_write(path, &serialized)?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
Err(AppError::Conflict(format!(
|
||||
"Pi models file changed concurrently too many times: {}",
|
||||
path.display()
|
||||
)))
|
||||
}
|
||||
|
||||
pub(crate) fn current_pi_provider_values(
|
||||
path: &Path,
|
||||
provider_keys: impl IntoIterator<Item = String>,
|
||||
) -> Result<IndexMap<String, Option<Value>>, AppError> {
|
||||
let document = read_pi_models_document(path)?;
|
||||
Ok(provider_keys
|
||||
.into_iter()
|
||||
.map(|key| {
|
||||
let value = document
|
||||
.providers()
|
||||
.get(&key)
|
||||
.map(|entry| entry.value.clone());
|
||||
(key, value)
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -288,6 +532,65 @@ mod tests {
|
||||
assert!(entry.raw_source.contains("\"model//literal\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exact_key_patch_preserves_unowned_entries_comments_and_root_fields() {
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let path = temp.path().join("models.json");
|
||||
fs::write(
|
||||
&path,
|
||||
r#"{
|
||||
"theme": "native",
|
||||
"providers": {
|
||||
// user-owned
|
||||
"native": {"models": [{"id": "native"}]},
|
||||
"managed": {"models": [{"id": "old"}]}
|
||||
}
|
||||
}
|
||||
"#,
|
||||
)
|
||||
.expect("write");
|
||||
let patch = IndexMap::from([(
|
||||
"managed".to_string(),
|
||||
Some(serde_json::json!({"models": [{"id": "new"}]})),
|
||||
)]);
|
||||
|
||||
apply_pi_provider_patch(&path, &patch).expect("patch");
|
||||
|
||||
let saved = fs::read_to_string(&path).expect("read");
|
||||
assert!(saved.contains("// user-owned"));
|
||||
assert!(saved.contains("\"theme\": \"native\""));
|
||||
let document = read_pi_models_document(&path).expect("parse");
|
||||
assert_eq!(
|
||||
document.providers()["native"].value["models"][0]["id"],
|
||||
"native"
|
||||
);
|
||||
assert_eq!(
|
||||
document.providers()["managed"].value["models"][0]["id"],
|
||||
"new"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exact_key_delete_does_not_delete_same_content_sibling() {
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let path = temp.path().join("models.json");
|
||||
fs::write(
|
||||
&path,
|
||||
r#"{"providers":{
|
||||
"managed":{"models":[{"id":"same"}]},
|
||||
"native":{"models":[{"id":"same"}]}
|
||||
}}"#,
|
||||
)
|
||||
.expect("write");
|
||||
let patch = IndexMap::from([("managed".to_string(), None)]);
|
||||
|
||||
apply_pi_provider_patch(&path, &patch).expect("delete");
|
||||
|
||||
let document = read_pi_models_document(&path).expect("parse");
|
||||
assert!(!document.providers().contains_key("managed"));
|
||||
assert!(document.providers().contains_key("native"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_javascript_overflow_without_hiding_sibling_entries() {
|
||||
let document = parse_models_source(
|
||||
|
||||
@@ -81,7 +81,7 @@ fn configured_header_class(name: &HeaderName) -> ConfiguredHeaderClass {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub(super) enum PiGatewayApiFamily {
|
||||
pub(crate) enum PiGatewayApiFamily {
|
||||
AnthropicMessages,
|
||||
OpenAiCompletions,
|
||||
OpenAiResponses,
|
||||
@@ -96,7 +96,7 @@ impl PiGatewayApiFamily {
|
||||
Self::GoogleGenerativeAi,
|
||||
];
|
||||
|
||||
pub(super) const fn as_str(self) -> &'static str {
|
||||
pub(crate) const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::AnthropicMessages => "anthropic-messages",
|
||||
Self::OpenAiCompletions => "openai-completions",
|
||||
@@ -113,14 +113,14 @@ impl PiGatewayApiFamily {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum PiGatewayCapability {
|
||||
pub(crate) enum PiGatewayCapability {
|
||||
Proxyable,
|
||||
DirectOnly,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum PiGatewayReasonCode {
|
||||
pub(crate) enum PiGatewayReasonCode {
|
||||
UnsupportedFamily,
|
||||
UnsupportedCredentialKind,
|
||||
InvalidEndpoint,
|
||||
@@ -132,13 +132,13 @@ pub(super) enum PiGatewayReasonCode {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) struct PiGatewayReason {
|
||||
pub(crate) struct PiGatewayReason {
|
||||
pub code: PiGatewayReasonCode,
|
||||
pub json_pointer: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct PiGatewayAssessment {
|
||||
pub(crate) struct PiGatewayAssessment {
|
||||
pub capability: PiGatewayCapability,
|
||||
pub reasons: Vec<PiGatewayReason>,
|
||||
pub plans: Vec<CandidateHeaderPlan>,
|
||||
@@ -175,7 +175,7 @@ impl DeferredHeaderValue {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct CandidateHeaderPlan {
|
||||
pub(crate) struct CandidateHeaderPlan {
|
||||
family: PiGatewayApiFamily,
|
||||
endpoint: Url,
|
||||
credential: DeferredHeaderValue,
|
||||
@@ -194,7 +194,7 @@ struct PlannedHeader {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(super) struct MaterializedCandidate {
|
||||
pub(crate) struct MaterializedCandidate {
|
||||
pub endpoint: Url,
|
||||
pub headers: HeaderMap,
|
||||
family: PiGatewayApiFamily,
|
||||
@@ -202,7 +202,7 @@ pub(super) struct MaterializedCandidate {
|
||||
protocol_identity_predictable: bool,
|
||||
}
|
||||
|
||||
pub(super) trait DeferredValueResolver {
|
||||
pub(crate) trait DeferredValueResolver {
|
||||
fn resolve(&self, expression: &str) -> Option<String>;
|
||||
}
|
||||
|
||||
@@ -219,6 +219,7 @@ impl CandidateHeaderPlan {
|
||||
fn build(
|
||||
model: &PiComposedNativeModel,
|
||||
model_index: usize,
|
||||
allow_anthropic_oauth: bool,
|
||||
) -> Result<Self, Vec<PiGatewayReason>> {
|
||||
let mut reasons = Vec::new();
|
||||
let Some(family) = PiGatewayApiFamily::parse(model.api.as_str()) else {
|
||||
@@ -256,6 +257,7 @@ impl CandidateHeaderPlan {
|
||||
});
|
||||
} else if family == PiGatewayApiFamily::AnthropicMessages
|
||||
&& is_anthropic_oauth_credential(credential)
|
||||
&& !allow_anthropic_oauth
|
||||
{
|
||||
reasons.push(PiGatewayReason {
|
||||
code: PiGatewayReasonCode::UnsupportedCredentialKind,
|
||||
@@ -303,6 +305,37 @@ impl CandidateHeaderPlan {
|
||||
pub(super) fn materialize(
|
||||
&self,
|
||||
resolver: &impl DeferredValueResolver,
|
||||
) -> Result<MaterializedCandidate, PiGatewayReason> {
|
||||
self.materialize_with_policy(resolver, false)
|
||||
}
|
||||
|
||||
pub(crate) fn materialize_for_runtime(
|
||||
&self,
|
||||
resolver: &impl DeferredValueResolver,
|
||||
) -> Result<MaterializedCandidate, PiGatewayReason> {
|
||||
self.materialize_with_policy(resolver, true)
|
||||
}
|
||||
|
||||
pub(crate) fn with_endpoint(&self, endpoint: &str) -> Result<Self, PiGatewayReason> {
|
||||
let endpoint = Url::parse(endpoint).map_err(|_| PiGatewayReason {
|
||||
code: PiGatewayReasonCode::InvalidEndpoint,
|
||||
json_pointer: "/customEndpoints".to_string(),
|
||||
})?;
|
||||
if !matches!(endpoint.scheme(), "http" | "https") || endpoint.host().is_none() {
|
||||
return Err(PiGatewayReason {
|
||||
code: PiGatewayReasonCode::InvalidEndpoint,
|
||||
json_pointer: "/customEndpoints".to_string(),
|
||||
});
|
||||
}
|
||||
let mut candidate = self.clone();
|
||||
candidate.endpoint = endpoint;
|
||||
Ok(candidate)
|
||||
}
|
||||
|
||||
fn materialize_with_policy(
|
||||
&self,
|
||||
resolver: &impl DeferredValueResolver,
|
||||
allow_anthropic_oauth: bool,
|
||||
) -> Result<MaterializedCandidate, PiGatewayReason> {
|
||||
// A new map is allocated for every candidate. No value from a prior
|
||||
// candidate can survive failover.
|
||||
@@ -311,6 +344,7 @@ impl CandidateHeaderPlan {
|
||||
let credential = self.credential.materialize(resolver, "/apiKey")?;
|
||||
if self.family == PiGatewayApiFamily::AnthropicMessages
|
||||
&& credential.to_str().is_ok_and(is_anthropic_oauth_credential)
|
||||
&& !allow_anthropic_oauth
|
||||
{
|
||||
return Err(PiGatewayReason {
|
||||
code: PiGatewayReasonCode::UnsupportedCredentialKind,
|
||||
@@ -318,7 +352,23 @@ impl CandidateHeaderPlan {
|
||||
});
|
||||
}
|
||||
let bearer_credential = credential.clone();
|
||||
let anthropic_oauth = self.family == PiGatewayApiFamily::AnthropicMessages
|
||||
&& credential.to_str().is_ok_and(is_anthropic_oauth_credential);
|
||||
let (auth_name, auth_value) = match self.family {
|
||||
PiGatewayApiFamily::AnthropicMessages if anthropic_oauth => {
|
||||
let credential = credential.to_str().map_err(|_| PiGatewayReason {
|
||||
code: PiGatewayReasonCode::InvalidHeaderValue,
|
||||
json_pointer: "/apiKey".to_string(),
|
||||
})?;
|
||||
let bearer =
|
||||
HeaderValue::from_str(&format!("Bearer {credential}")).map_err(|_| {
|
||||
PiGatewayReason {
|
||||
code: PiGatewayReasonCode::InvalidHeaderValue,
|
||||
json_pointer: "/apiKey".to_string(),
|
||||
}
|
||||
})?;
|
||||
(HeaderName::from_static("authorization"), bearer)
|
||||
}
|
||||
PiGatewayApiFamily::AnthropicMessages => {
|
||||
(HeaderName::from_static("x-api-key"), credential)
|
||||
}
|
||||
@@ -350,6 +400,15 @@ impl CandidateHeaderPlan {
|
||||
let value = HeaderValue::from_static("2023-06-01");
|
||||
protocol_headers.insert(name.clone(), value.clone());
|
||||
headers.insert(name, value);
|
||||
let beta = if anthropic_oauth {
|
||||
"claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14"
|
||||
} else {
|
||||
"interleaved-thinking-2025-05-14"
|
||||
};
|
||||
let name = HeaderName::from_static("anthropic-beta");
|
||||
let value = HeaderValue::from_static(beta);
|
||||
protocol_headers.insert(name.clone(), value.clone());
|
||||
headers.insert(name, value);
|
||||
}
|
||||
|
||||
// Provider headers are part of provider auth resolution. Pinned SDKs
|
||||
@@ -386,6 +445,31 @@ impl CandidateHeaderPlan {
|
||||
&mut protocol_headers,
|
||||
)?;
|
||||
|
||||
// Main-project OAuth transport is a completed policy boundary, not a
|
||||
// partial SDK-header overlay. A configured auth header may override
|
||||
// synthesized auth for ordinary credentials (matching pinned Pi), but
|
||||
// an Anthropic OAuth credential is always transported as that exact
|
||||
// Bearer and never alongside x-api-key.
|
||||
if anthropic_oauth {
|
||||
headers.remove(HeaderName::from_static("x-api-key"));
|
||||
let credential = bearer_credential.to_str().map_err(|_| PiGatewayReason {
|
||||
code: PiGatewayReasonCode::InvalidHeaderValue,
|
||||
json_pointer: "/apiKey".to_string(),
|
||||
})?;
|
||||
let bearer = HeaderValue::from_str(&format!("Bearer {credential}")).map_err(|_| {
|
||||
PiGatewayReason {
|
||||
code: PiGatewayReasonCode::InvalidHeaderValue,
|
||||
json_pointer: "/apiKey".to_string(),
|
||||
}
|
||||
})?;
|
||||
headers.insert(HeaderName::from_static("authorization"), bearer);
|
||||
let beta = HeaderValue::from_static(
|
||||
"claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14",
|
||||
);
|
||||
headers.insert(HeaderName::from_static("anthropic-beta"), beta.clone());
|
||||
protocol_headers.insert(HeaderName::from_static("anthropic-beta"), beta);
|
||||
}
|
||||
|
||||
let host = authority_header(&self.endpoint).ok_or_else(|| PiGatewayReason {
|
||||
code: PiGatewayReasonCode::InvalidEndpoint,
|
||||
json_pointer: "/baseUrl".to_string(),
|
||||
@@ -460,11 +544,33 @@ fn apply_planned_headers(
|
||||
}
|
||||
|
||||
impl MaterializedCandidate {
|
||||
pub(super) fn failover_protocol_identity(&self) -> Option<(PiGatewayApiFamily, &HeaderMap)> {
|
||||
pub(crate) fn family(&self) -> PiGatewayApiFamily {
|
||||
self.family
|
||||
}
|
||||
|
||||
pub(crate) fn failover_protocol_identity(&self) -> Option<(PiGatewayApiFamily, &HeaderMap)> {
|
||||
// Auth, tenant and arbitrary custom headers are deliberately excluded.
|
||||
self.protocol_identity_predictable
|
||||
.then_some((self.family, &self.protocol_headers))
|
||||
}
|
||||
|
||||
pub(crate) fn family_name(&self) -> &'static str {
|
||||
self.family.as_str()
|
||||
}
|
||||
}
|
||||
|
||||
impl CandidateHeaderPlan {
|
||||
pub(crate) fn family(&self) -> PiGatewayApiFamily {
|
||||
self.family
|
||||
}
|
||||
|
||||
pub(crate) fn protocol_identity_is_predictable(&self) -> bool {
|
||||
self.protocol_identity_predictable
|
||||
}
|
||||
|
||||
pub(crate) fn endpoint(&self) -> &Url {
|
||||
&self.endpoint
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn assess_composition(composition: &PiNativeComposition) -> PiGatewayAssessment {
|
||||
@@ -478,7 +584,43 @@ pub(super) fn assess_composition(composition: &PiNativeComposition) -> PiGateway
|
||||
let mut plans = Vec::with_capacity(composition.models.len());
|
||||
let mut reasons = Vec::new();
|
||||
for (index, model) in composition.models.iter().enumerate() {
|
||||
match CandidateHeaderPlan::build(model, index) {
|
||||
match CandidateHeaderPlan::build(model, index, false) {
|
||||
Ok(plan) => plans.push(plan),
|
||||
Err(mut model_reasons) => reasons.append(&mut model_reasons),
|
||||
}
|
||||
}
|
||||
if reasons.is_empty() && plans.len() == composition.models.len() {
|
||||
PiGatewayAssessment {
|
||||
capability: PiGatewayCapability::Proxyable,
|
||||
reasons,
|
||||
plans,
|
||||
}
|
||||
} else {
|
||||
PiGatewayAssessment {
|
||||
capability: PiGatewayCapability::DirectOnly,
|
||||
reasons,
|
||||
plans: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Main-project data plane assessment. The certified Pre-C assessment remains
|
||||
/// unchanged and honestly reports Anthropic OAuth as DirectOnly; this entry
|
||||
/// point becomes reachable only with the complete OAuth transport policy.
|
||||
pub(crate) fn assess_composition_for_runtime(
|
||||
composition: &PiNativeComposition,
|
||||
) -> PiGatewayAssessment {
|
||||
if composition.status != PiComposerStatus::Composed {
|
||||
return PiGatewayAssessment {
|
||||
capability: PiGatewayCapability::Unknown,
|
||||
reasons: Vec::new(),
|
||||
plans: Vec::new(),
|
||||
};
|
||||
}
|
||||
let mut plans = Vec::with_capacity(composition.models.len());
|
||||
let mut reasons = Vec::new();
|
||||
for (index, model) in composition.models.iter().enumerate() {
|
||||
match CandidateHeaderPlan::build(model, index, true) {
|
||||
Ok(plan) => plans.push(plan),
|
||||
Err(mut model_reasons) => reasons.append(&mut model_reasons),
|
||||
}
|
||||
@@ -499,10 +641,14 @@ pub(super) fn assess_composition(composition: &PiNativeComposition) -> PiGateway
|
||||
}
|
||||
|
||||
fn authority_header(url: &Url) -> Option<HeaderValue> {
|
||||
let host = url.host_str()?;
|
||||
let host = match url.host()? {
|
||||
url::Host::Domain(value) => value.to_string(),
|
||||
url::Host::Ipv4(value) => value.to_string(),
|
||||
url::Host::Ipv6(value) => format!("[{value}]"),
|
||||
};
|
||||
let authority = match url.port() {
|
||||
Some(port) => format!("{host}:{port}"),
|
||||
None => host.to_string(),
|
||||
None => host,
|
||||
};
|
||||
HeaderValue::from_str(&authority).ok()
|
||||
}
|
||||
@@ -941,6 +1087,63 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_host_header_preserves_ipv6_authority_brackets() {
|
||||
let composition = composed(json!({
|
||||
"api": "openai-responses",
|
||||
"baseUrl": "http://[::1]:8443/v1",
|
||||
"apiKey": "secret",
|
||||
"models": [{"id": "m"}]
|
||||
}));
|
||||
let materialized = assess_composition(&composition)
|
||||
.plans
|
||||
.remove(0)
|
||||
.materialize(&|_expression: &str| None)
|
||||
.expect("IPv6 endpoint");
|
||||
assert_eq!(
|
||||
materialized.headers[&HeaderName::from_static("host")],
|
||||
"[::1]:8443"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completed_runtime_oauth_policy_forces_bearer_beta_and_no_x_api_key() {
|
||||
let composition = composed(json!({
|
||||
"api": "anthropic-messages",
|
||||
"baseUrl": "https://candidate.example",
|
||||
"apiKey": "prefix-sk-ant-oat01-token-suffix",
|
||||
"headers": {
|
||||
"authorization": "Bearer configured",
|
||||
"x-api-key": "configured-api-key",
|
||||
"anthropic-beta": "configured-beta"
|
||||
},
|
||||
"models": [{"id": "m"}]
|
||||
}));
|
||||
let certified = assess_composition(&composition);
|
||||
assert_eq!(certified.capability, PiGatewayCapability::DirectOnly);
|
||||
assert_eq!(
|
||||
certified.reasons[0].code,
|
||||
PiGatewayReasonCode::UnsupportedCredentialKind
|
||||
);
|
||||
|
||||
let mut runtime = assess_composition_for_runtime(&composition);
|
||||
assert_eq!(runtime.capability, PiGatewayCapability::Proxyable);
|
||||
let materialized = runtime
|
||||
.plans
|
||||
.remove(0)
|
||||
.materialize_for_runtime(&|_expression: &str| None)
|
||||
.expect("the complete main-project OAuth policy is proxyable");
|
||||
assert_eq!(
|
||||
materialized.headers[&HeaderName::from_static("authorization")],
|
||||
"Bearer prefix-sk-ant-oat01-token-suffix"
|
||||
);
|
||||
assert!(materialized.headers.get("x-api-key").is_none());
|
||||
assert_eq!(
|
||||
materialized.headers[&HeaderName::from_static("anthropic-beta")],
|
||||
"claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deferred_values_replay_actual_pinned_pi_transport_results() {
|
||||
let oracle: Value =
|
||||
|
||||
@@ -7,14 +7,16 @@
|
||||
use indexmap::IndexMap;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
mod composer;
|
||||
mod document;
|
||||
mod gateway;
|
||||
pub(crate) mod composer;
|
||||
pub(crate) mod document;
|
||||
pub(crate) mod gateway;
|
||||
pub(crate) mod model;
|
||||
pub(crate) mod native;
|
||||
#[cfg(test)]
|
||||
mod native_inspection_certification;
|
||||
mod raw_schema;
|
||||
pub(crate) mod native_settings;
|
||||
pub(crate) mod raw_schema;
|
||||
pub(crate) mod shared_file;
|
||||
|
||||
const PI_COMPAT_NESTED_SPREAD_KEYS: [&str; 3] = [
|
||||
"openRouterRouting",
|
||||
|
||||
@@ -156,6 +156,26 @@ pub(crate) fn inspect_pi_native_entry(
|
||||
PiNativeInspectionService::inspect_entry(path, provider_key, managed_claims)
|
||||
}
|
||||
|
||||
/// Compose a database-authoritative managed provider through the same raw and
|
||||
/// composer layers used by native inspection. Runtime construction must not
|
||||
/// reimplement inheritance or field semantics.
|
||||
pub(crate) fn compose_managed_pi_provider(
|
||||
provider_key: &str,
|
||||
config: &PiManagedProviderConfig,
|
||||
) -> Result<PiNativeComposition, AppError> {
|
||||
validate_pi_managed_provider(config)
|
||||
.map_err(|error| AppError::InvalidInput(error.to_string()))?;
|
||||
let value =
|
||||
serde_json::to_value(config).map_err(|source| AppError::JsonSerialize { source })?;
|
||||
let raw = evaluate_provider_value(&value);
|
||||
let provider = raw.valid_provider.as_ref().ok_or_else(|| {
|
||||
AppError::Config(format!(
|
||||
"managed Pi provider '{provider_key}' did not pass the pinned raw schema"
|
||||
))
|
||||
})?;
|
||||
Ok(compose_explicit_custom_catalog(provider_key, provider))
|
||||
}
|
||||
|
||||
fn normalize_pi_agent_dir(value: &str, home: &Path) -> Result<PathBuf, AppError> {
|
||||
if value == "~" {
|
||||
return Ok(home.to_path_buf());
|
||||
@@ -180,7 +200,15 @@ fn normalize_pi_agent_dir(value: &str, home: &Path) -> Result<PathBuf, AppError>
|
||||
Ok(PathBuf::from(value))
|
||||
}
|
||||
|
||||
pub(crate) fn get_pi_agent_dir() -> Result<PathBuf, AppError> {
|
||||
pub(crate) fn get_pi_agent_dir_for_override(
|
||||
override_dir: Option<&str>,
|
||||
) -> Result<PathBuf, AppError> {
|
||||
if let Some(override_dir) = override_dir
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return Ok(crate::settings::resolve_override_path(override_dir));
|
||||
}
|
||||
let Some(raw) = std::env::var_os("PI_CODING_AGENT_DIR") else {
|
||||
return Ok(get_home_dir().join(".pi").join("agent"));
|
||||
};
|
||||
@@ -190,6 +218,19 @@ pub(crate) fn get_pi_agent_dir() -> Result<PathBuf, AppError> {
|
||||
normalize_pi_agent_dir(&raw.to_string_lossy(), &get_home_dir())
|
||||
}
|
||||
|
||||
pub(crate) fn get_pi_agent_dir() -> Result<PathBuf, AppError> {
|
||||
if let Some(override_dir) = crate::settings::get_pi_override_dir() {
|
||||
return Ok(override_dir);
|
||||
}
|
||||
get_pi_agent_dir_for_override(None)
|
||||
}
|
||||
|
||||
pub(crate) fn get_pi_models_path_for_override(
|
||||
override_dir: Option<&str>,
|
||||
) -> Result<PathBuf, AppError> {
|
||||
Ok(get_pi_agent_dir_for_override(override_dir)?.join("models.json"))
|
||||
}
|
||||
|
||||
pub(crate) fn get_pi_models_path() -> Result<PathBuf, AppError> {
|
||||
Ok(get_pi_agent_dir()?.join("models.json"))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,300 @@
|
||||
//! Exact-field access to Pi's shared `settings.json`.
|
||||
//!
|
||||
//! cc-switch owns only `defaultProvider` and `defaultModel`. Every other field
|
||||
//! remains Pi/user-owned and survives each mutation unchanged.
|
||||
|
||||
use crate::error::AppError;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::fs::{self, File, Metadata, OpenOptions};
|
||||
use std::io::{Read, Take};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{LazyLock, Mutex};
|
||||
|
||||
const MAX_PI_SETTINGS_BYTES: u64 = 1024 * 1024;
|
||||
const MAX_WRITE_ATTEMPTS: usize = 3;
|
||||
static SETTINGS_WRITE_LOCK: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct PiNativeDefaults {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub default_provider: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub default_model: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub session_dir: Option<String>,
|
||||
}
|
||||
|
||||
pub(crate) fn get_pi_settings_path() -> Result<PathBuf, AppError> {
|
||||
Ok(super::native::get_pi_agent_dir()?.join("settings.json"))
|
||||
}
|
||||
|
||||
pub(crate) fn read_pi_native_defaults() -> Result<PiNativeDefaults, AppError> {
|
||||
read_pi_native_defaults_at(&get_pi_settings_path()?)
|
||||
}
|
||||
|
||||
pub(crate) fn read_pi_native_defaults_at(path: &Path) -> Result<PiNativeDefaults, AppError> {
|
||||
let document = read_settings_document(path)?;
|
||||
let root = document.as_object().ok_or_else(|| {
|
||||
AppError::Config(format!(
|
||||
"Pi settings root must be an object: {}",
|
||||
path.display()
|
||||
))
|
||||
})?;
|
||||
Ok(PiNativeDefaults {
|
||||
default_provider: optional_string(root, "defaultProvider", path)?,
|
||||
default_model: optional_string(root, "defaultModel", path)?,
|
||||
session_dir: optional_string(root, "sessionDir", path)?,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn set_pi_native_default(
|
||||
provider_key: &str,
|
||||
model_id: &str,
|
||||
) -> Result<PiNativeDefaults, AppError> {
|
||||
if provider_key.trim().is_empty() || model_id.trim().is_empty() {
|
||||
return Err(AppError::InvalidInput(
|
||||
"Pi default provider and model must be non-empty".to_string(),
|
||||
));
|
||||
}
|
||||
mutate_settings_document(&get_pi_settings_path()?, |root| {
|
||||
root.insert(
|
||||
"defaultProvider".to_string(),
|
||||
Value::String(provider_key.to_string()),
|
||||
);
|
||||
root.insert(
|
||||
"defaultModel".to_string(),
|
||||
Value::String(model_id.to_string()),
|
||||
);
|
||||
Ok(())
|
||||
})?;
|
||||
read_pi_native_defaults()
|
||||
}
|
||||
|
||||
/// Restore the two fields owned by cc-switch without touching Pi-owned
|
||||
/// settings. This is intentionally narrower than replacing settings.json and
|
||||
/// is used by catalog compensation after a later authority step fails.
|
||||
pub(crate) fn replace_pi_native_defaults(
|
||||
defaults: &PiNativeDefaults,
|
||||
) -> Result<PiNativeDefaults, AppError> {
|
||||
mutate_settings_document(&get_pi_settings_path()?, |root| {
|
||||
set_optional_string(
|
||||
root,
|
||||
"defaultProvider",
|
||||
defaults.default_provider.as_deref(),
|
||||
);
|
||||
set_optional_string(root, "defaultModel", defaults.default_model.as_deref());
|
||||
Ok(())
|
||||
})?;
|
||||
read_pi_native_defaults()
|
||||
}
|
||||
|
||||
fn set_optional_string(root: &mut Map<String, Value>, key: &str, value: Option<&str>) {
|
||||
match value {
|
||||
Some(value) => {
|
||||
root.insert(key.to_string(), Value::String(value.to_string()));
|
||||
}
|
||||
None => {
|
||||
root.remove(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_string(
|
||||
root: &Map<String, Value>,
|
||||
key: &str,
|
||||
path: &Path,
|
||||
) -> Result<Option<String>, AppError> {
|
||||
match root.get(key) {
|
||||
None | Some(Value::Null) => Ok(None),
|
||||
Some(Value::String(value)) => Ok(Some(value.clone())),
|
||||
Some(_) => Err(AppError::Config(format!(
|
||||
"Pi settings field '{key}' must be a string: {}",
|
||||
path.display()
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn mutate_settings_document(
|
||||
path: &Path,
|
||||
mut mutator: impl FnMut(&mut Map<String, Value>) -> Result<(), AppError>,
|
||||
) -> Result<(), AppError> {
|
||||
let _guard = SETTINGS_WRITE_LOCK
|
||||
.lock()
|
||||
.map_err(|error| AppError::Config(format!("Pi settings lock is poisoned: {error}")))?;
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent).map_err(|error| AppError::io(parent, error))?;
|
||||
}
|
||||
|
||||
for _ in 0..MAX_WRITE_ATTEMPTS {
|
||||
let before = read_regular_bytes(path, MAX_PI_SETTINGS_BYTES)?;
|
||||
let observed = fingerprint(before.as_deref());
|
||||
let mut document = match before.as_deref() {
|
||||
Some(bytes) => {
|
||||
serde_json::from_slice(bytes).map_err(|error| AppError::json(path, error))?
|
||||
}
|
||||
None => Value::Object(Map::new()),
|
||||
};
|
||||
let root = document.as_object_mut().ok_or_else(|| {
|
||||
AppError::Config(format!(
|
||||
"Pi settings root must be an object: {}",
|
||||
path.display()
|
||||
))
|
||||
})?;
|
||||
mutator(root)?;
|
||||
let mut serialized = serde_json::to_vec_pretty(&document)
|
||||
.map_err(|source| AppError::JsonSerialize { source })?;
|
||||
serialized.push(b'\n');
|
||||
|
||||
let current = read_regular_bytes(path, MAX_PI_SETTINGS_BYTES)?;
|
||||
if fingerprint(current.as_deref()) != observed {
|
||||
continue;
|
||||
}
|
||||
crate::config::atomic_write(path, &serialized)?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
Err(AppError::Conflict(format!(
|
||||
"Pi settings changed concurrently too many times: {}",
|
||||
path.display()
|
||||
)))
|
||||
}
|
||||
|
||||
fn read_settings_document(path: &Path) -> Result<Value, AppError> {
|
||||
match read_regular_bytes(path, MAX_PI_SETTINGS_BYTES)? {
|
||||
Some(bytes) => serde_json::from_slice(&bytes).map_err(|error| AppError::json(path, error)),
|
||||
None => Ok(Value::Object(Map::new())),
|
||||
}
|
||||
}
|
||||
|
||||
fn fingerprint(bytes: Option<&[u8]>) -> Option<[u8; 32]> {
|
||||
bytes.map(|bytes| Sha256::digest(bytes).into())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn open_read_only(path: &Path) -> std::io::Result<File> {
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
OpenOptions::new()
|
||||
.read(true)
|
||||
.custom_flags(libc::O_NOFOLLOW)
|
||||
.open(path)
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn open_read_only(path: &Path) -> std::io::Result<File> {
|
||||
OpenOptions::new().read(true).open(path)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn same_file(left: &Metadata, right: &Metadata) -> bool {
|
||||
use std::os::unix::fs::MetadataExt;
|
||||
left.dev() == right.dev() && left.ino() == right.ino()
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn same_file(left: &Metadata, right: &Metadata) -> bool {
|
||||
left.len() == right.len() && left.modified().ok() == right.modified().ok()
|
||||
}
|
||||
|
||||
fn read_limited(
|
||||
mut reader: Take<&mut File>,
|
||||
path: &Path,
|
||||
max_bytes: u64,
|
||||
) -> Result<Vec<u8>, AppError> {
|
||||
let mut bytes = Vec::new();
|
||||
reader
|
||||
.read_to_end(&mut bytes)
|
||||
.map_err(|error| AppError::io(path, error))?;
|
||||
if bytes.len() as u64 > max_bytes {
|
||||
return Err(AppError::Config(format!(
|
||||
"Pi settings exceeds {max_bytes} bytes: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
fn read_regular_bytes(path: &Path, max_bytes: u64) -> Result<Option<Vec<u8>>, AppError> {
|
||||
let initial = match fs::symlink_metadata(path) {
|
||||
Ok(metadata) => metadata,
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
|
||||
Err(error) => return Err(AppError::io(path, error)),
|
||||
};
|
||||
if !initial.file_type().is_file() || initial.len() > max_bytes {
|
||||
return Err(AppError::Config(format!(
|
||||
"Pi settings must be a bounded regular file: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
let mut file = open_read_only(path).map_err(|error| AppError::io(path, error))?;
|
||||
let opened = file.metadata().map_err(|error| AppError::io(path, error))?;
|
||||
let bytes = read_limited(file.by_ref().take(max_bytes + 1), path, max_bytes)?;
|
||||
let completed = file.metadata().map_err(|error| AppError::io(path, error))?;
|
||||
let current = fs::symlink_metadata(path).map_err(|error| AppError::io(path, error))?;
|
||||
if !current.file_type().is_file()
|
||||
|| !same_file(&opened, &completed)
|
||||
|| !same_file(&completed, ¤t)
|
||||
|| opened.len() != bytes.len() as u64
|
||||
|| completed.len() != bytes.len() as u64
|
||||
|| current.len() != bytes.len() as u64
|
||||
|| opened.modified().ok() != completed.modified().ok()
|
||||
{
|
||||
return Err(AppError::Conflict(format!(
|
||||
"Pi settings changed during read: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
Ok(Some(bytes))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn default_patch_preserves_every_unowned_field() {
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let path = temp.path().join("settings.json");
|
||||
fs::write(
|
||||
&path,
|
||||
serde_json::to_vec_pretty(&json!({
|
||||
"theme": "custom",
|
||||
"packages": ["npm:foreign"],
|
||||
"sessionDir": "/tmp/pi-sessions",
|
||||
"defaultProvider": "old",
|
||||
"defaultModel": "old-model"
|
||||
}))
|
||||
.expect("serialize"),
|
||||
)
|
||||
.expect("write");
|
||||
|
||||
mutate_settings_document(&path, |root| {
|
||||
root.insert("defaultProvider".into(), json!("managed"));
|
||||
root.insert("defaultModel".into(), json!("model"));
|
||||
Ok(())
|
||||
})
|
||||
.expect("mutate");
|
||||
|
||||
let saved: Value = serde_json::from_slice(&fs::read(&path).expect("read")).expect("parse");
|
||||
assert_eq!(saved["theme"], "custom");
|
||||
assert_eq!(saved["packages"], json!(["npm:foreign"]));
|
||||
assert_eq!(saved["sessionDir"], "/tmp/pi-sessions");
|
||||
assert_eq!(saved["defaultProvider"], "managed");
|
||||
assert_eq!(saved["defaultModel"], "model");
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn settings_symlink_is_rejected() {
|
||||
use std::os::unix::fs::symlink;
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let target = temp.path().join("target.json");
|
||||
let path = temp.path().join("settings.json");
|
||||
fs::write(&target, "{}").expect("target");
|
||||
symlink(&target, &path).expect("symlink");
|
||||
assert!(read_pi_native_defaults_at(&path).is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
//! Reusable safety boundary for exact Pi-owned/shared files.
|
||||
//!
|
||||
//! Callers choose an exact path and size limit. This layer supplies bounded
|
||||
//! regular-file reads, symlink rejection, optimistic revisions, per-path
|
||||
//! process locking, durable atomic replacement, and compare-before-delete.
|
||||
|
||||
use crate::error::AppError;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::collections::HashMap;
|
||||
use std::fs::{self, File, Metadata, OpenOptions};
|
||||
use std::io::{Read, Take};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, LazyLock, Mutex};
|
||||
|
||||
static FILE_LOCKS: LazyLock<Mutex<HashMap<PathBuf, Arc<Mutex<()>>>>> =
|
||||
LazyLock::new(|| Mutex::new(HashMap::new()));
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct SharedFileSnapshot {
|
||||
pub revision: String,
|
||||
pub bytes: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
impl SharedFileSnapshot {
|
||||
pub(crate) fn exists(&self) -> bool {
|
||||
self.bytes.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn read_shared_file(
|
||||
path: &Path,
|
||||
max_bytes: u64,
|
||||
label: &str,
|
||||
) -> Result<SharedFileSnapshot, AppError> {
|
||||
let bytes = read_regular_bytes(path, max_bytes, label)?;
|
||||
Ok(SharedFileSnapshot {
|
||||
revision: revision(bytes.as_deref()),
|
||||
bytes,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn replace_shared_file(
|
||||
path: &Path,
|
||||
expected_revision: &str,
|
||||
bytes: &[u8],
|
||||
max_bytes: u64,
|
||||
new_file_mode: Option<u32>,
|
||||
label: &str,
|
||||
) -> Result<SharedFileSnapshot, AppError> {
|
||||
if bytes.len() as u64 > max_bytes {
|
||||
return Err(AppError::InvalidInput(format!(
|
||||
"{label} exceeds the {max_bytes}-byte limit"
|
||||
)));
|
||||
}
|
||||
let lock = path_lock(path)?;
|
||||
let _guard = lock
|
||||
.lock()
|
||||
.map_err(|error| AppError::Config(format!("Pi file lock is poisoned: {error}")))?;
|
||||
let current = read_shared_file(path, max_bytes, label)?;
|
||||
ensure_revision(path, expected_revision, ¤t.revision)?;
|
||||
crate::config::atomic_write_durable(path, bytes, new_file_mode)?;
|
||||
read_shared_file(path, max_bytes, label)
|
||||
}
|
||||
|
||||
pub(crate) fn delete_shared_file(
|
||||
path: &Path,
|
||||
expected_revision: &str,
|
||||
max_bytes: u64,
|
||||
label: &str,
|
||||
) -> Result<bool, AppError> {
|
||||
let lock = path_lock(path)?;
|
||||
let _guard = lock
|
||||
.lock()
|
||||
.map_err(|error| AppError::Config(format!("Pi file lock is poisoned: {error}")))?;
|
||||
let current = read_shared_file(path, max_bytes, label)?;
|
||||
ensure_revision(path, expected_revision, ¤t.revision)?;
|
||||
if !current.exists() {
|
||||
return Ok(false);
|
||||
}
|
||||
fs::remove_file(path).map_err(|error| AppError::io(path, error))?;
|
||||
#[cfg(unix)]
|
||||
if let Some(parent) = path.parent() {
|
||||
File::open(parent)
|
||||
.and_then(|directory| directory.sync_all())
|
||||
.map_err(|error| AppError::io(parent, error))?;
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
fn ensure_revision(path: &Path, expected: &str, actual: &str) -> Result<(), AppError> {
|
||||
if expected == actual {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(AppError::Conflict(format!(
|
||||
"Pi file changed since it was read: {}",
|
||||
path.display()
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
fn path_lock(path: &Path) -> Result<Arc<Mutex<()>>, AppError> {
|
||||
let mut locks = FILE_LOCKS
|
||||
.lock()
|
||||
.map_err(|error| AppError::Config(format!("Pi file-lock registry is poisoned: {error}")))?;
|
||||
Ok(locks
|
||||
.entry(path.to_path_buf())
|
||||
.or_insert_with(|| Arc::new(Mutex::new(())))
|
||||
.clone())
|
||||
}
|
||||
|
||||
fn revision(bytes: Option<&[u8]>) -> String {
|
||||
bytes.map_or_else(
|
||||
|| "missing".to_string(),
|
||||
|bytes| format!("sha256:{:x}", Sha256::digest(bytes)),
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn open_read_only(path: &Path) -> std::io::Result<File> {
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
OpenOptions::new()
|
||||
.read(true)
|
||||
.custom_flags(libc::O_NOFOLLOW)
|
||||
.open(path)
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn open_read_only(path: &Path) -> std::io::Result<File> {
|
||||
OpenOptions::new().read(true).open(path)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn same_file(left: &Metadata, right: &Metadata) -> bool {
|
||||
use std::os::unix::fs::MetadataExt;
|
||||
left.dev() == right.dev() && left.ino() == right.ino()
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn same_file(left: &Metadata, right: &Metadata) -> bool {
|
||||
left.len() == right.len() && left.modified().ok() == right.modified().ok()
|
||||
}
|
||||
|
||||
fn read_limited(
|
||||
mut reader: Take<&mut File>,
|
||||
path: &Path,
|
||||
max_bytes: u64,
|
||||
) -> Result<Vec<u8>, AppError> {
|
||||
let mut bytes = Vec::new();
|
||||
reader
|
||||
.read_to_end(&mut bytes)
|
||||
.map_err(|error| AppError::io(path, error))?;
|
||||
if bytes.len() as u64 > max_bytes {
|
||||
return Err(AppError::InvalidInput(format!(
|
||||
"Pi file exceeds the {max_bytes}-byte limit: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
fn read_regular_bytes(
|
||||
path: &Path,
|
||||
max_bytes: u64,
|
||||
label: &str,
|
||||
) -> Result<Option<Vec<u8>>, AppError> {
|
||||
let initial = match fs::symlink_metadata(path) {
|
||||
Ok(metadata) => metadata,
|
||||
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
|
||||
Err(error) => return Err(AppError::io(path, error)),
|
||||
};
|
||||
if !initial.file_type().is_file() || initial.len() > max_bytes {
|
||||
return Err(AppError::InvalidInput(format!(
|
||||
"{label} must be a bounded regular file: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
let mut file = open_read_only(path).map_err(|error| AppError::io(path, error))?;
|
||||
let opened = file.metadata().map_err(|error| AppError::io(path, error))?;
|
||||
let bytes = read_limited(file.by_ref().take(max_bytes + 1), path, max_bytes)?;
|
||||
let completed = file.metadata().map_err(|error| AppError::io(path, error))?;
|
||||
let current = fs::symlink_metadata(path).map_err(|error| AppError::io(path, error))?;
|
||||
if !current.file_type().is_file()
|
||||
|| !same_file(&opened, &completed)
|
||||
|| !same_file(&completed, ¤t)
|
||||
|| opened.len() != bytes.len() as u64
|
||||
|| completed.len() != bytes.len() as u64
|
||||
|| current.len() != bytes.len() as u64
|
||||
|| opened.modified().ok() != completed.modified().ok()
|
||||
{
|
||||
return Err(AppError::Conflict(format!(
|
||||
"{label} changed during read: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
Ok(Some(bytes))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn compare_and_replace_distinguishes_missing_and_content_revisions() {
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let path = temp.path().join("shared.md");
|
||||
let missing = read_shared_file(&path, 1024, "test").expect("missing");
|
||||
assert_eq!(missing.revision, "missing");
|
||||
let written = replace_shared_file(&path, "missing", b"one", 1024, Some(0o600), "test")
|
||||
.expect("create");
|
||||
assert!(written.revision.starts_with("sha256:"));
|
||||
assert!(replace_shared_file(&path, "missing", b"two", 1024, None, "test").is_err());
|
||||
let replaced = replace_shared_file(&path, &written.revision, b"two", 1024, None, "test")
|
||||
.expect("replace");
|
||||
assert_eq!(replaced.bytes.as_deref(), Some(b"two".as_slice()));
|
||||
assert!(delete_shared_file(&path, &written.revision, 1024, "test").is_err());
|
||||
assert!(delete_shared_file(&path, &replaced.revision, 1024, "test").expect("delete"));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn symlink_targets_fail_closed() {
|
||||
use std::os::unix::fs::symlink;
|
||||
let temp = tempfile::tempdir().expect("tempdir");
|
||||
let target = temp.path().join("target");
|
||||
let path = temp.path().join("shared");
|
||||
fs::write(&target, b"secret").expect("target");
|
||||
symlink(&target, &path).expect("symlink");
|
||||
assert!(read_shared_file(&path, 1024, "test").is_err());
|
||||
assert!(replace_shared_file(&path, "missing", b"overwrite", 1024, None, "test").is_err());
|
||||
assert_eq!(fs::read(&target).expect("target remains"), b"secret");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user