diff --git a/controller/src/providers/signing.rs b/controller/src/providers/signing.rs index c1e481932..f719d93b9 100644 --- a/controller/src/providers/signing.rs +++ b/controller/src/providers/signing.rs @@ -168,6 +168,7 @@ pub fn sha256_hex(bytes: &[u8]) -> String { use std::fmt::Write; let _ = write!(out, "{b:02x}"); } + out } diff --git a/controller/src/reconciler/governed_services.rs b/controller/src/reconciler/governed_services.rs new file mode 100644 index 000000000..e2d4ff22c --- /dev/null +++ b/controller/src/reconciler/governed_services.rs @@ -0,0 +1,223 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Narrow router service identity/credential projection; no approvals or grants. + +use crate::{crd::KarsSandbox, kars_task::KarsTask}; +use k8s_openapi::api::{apps::v1::Deployment, core::v1::Namespace}; +use kube::{Api, Client, ResourceExt}; +use serde_json::{Value, json}; + +#[cfg(test)] +mod continuity_tests; +#[cfg(test)] +mod credential_tests; +mod credentials; + +const SECRET: &str = "router-services-admin"; +const SOURCE_UID: &str = "kars.azure.com/sandbox-uid"; +const NAMESPACE_UID: &str = "kars.azure.com/namespace-uid"; +pub(super) use credentials::quarantine_on_privacy_loss; + +pub struct Projection { + pub identity: Value, + credential: credentials::Projection, +} + +impl Projection { + pub fn decorate(&self, deployment: &mut Deployment) { + self.credential.decorate(deployment); + } + + pub async fn consumers_current( + &self, + client: &Client, + namespace: &str, + name: &str, + ) -> Result { + self.credential + .consumers_current(client, namespace, name) + .await + } +} + +fn api_error(error: kube::Error) -> String { + match error { + kube::Error::Api(status) => { + format!("Governed service credential API status {}", status.code) + } + _ => "Governed service credential API/transport failure".into(), + } +} + +fn authorized_task( + task: &KarsTask, + workspace: &str, + sandbox: &str, + name: &str, + uid: &str, +) -> Option { + let status = task.status.as_ref()?; + let authorization = task.spec.authorization_digest(); + (task.metadata.namespace.as_deref() == Some(workspace) + && task.metadata.name.as_deref() == Some(name) + && task.metadata.uid.as_deref() == Some(uid) + && task + .metadata + .generation + .is_some_and(|generation| generation > 0) + && task.metadata.deletion_timestamp.is_none() + && status.observed_generation == task.metadata.generation + && status.phase.as_deref() == Some("Ready") + && status + .conditions + .iter() + .flatten() + .any(|condition| condition.type_ == "Ready" && condition.status == "True") + && status.envelope_digest.as_deref() == Some(authorization.as_str()) + && status + .sandbox_ref + .as_ref() + .map(|reference| reference.name.as_str()) + == Some(sandbox) + && task + .spec + .execution + .as_ref() + .is_some_and(|execution| execution.launch) + && crate::kars_task::validate_execution_contract(&task.spec).is_ok()) + .then_some(authorization) +} + +pub async fn ensure( + client: &Client, + sandbox: &KarsSandbox, + namespace: &Namespace, +) -> Result { + // Reuse the authoritative namespace claim path, not labels or a caller's + // requested namespace. Recreated CRs/namespaces cannot inherit this token. + let (live, owned) = super::namespace_ownership::ensure(client, sandbox) + .await + .map_err(|_| "Governed service namespace authority could not be verified")?; + let owned = owned.ok_or("Governed service namespace is absent")?; + if live.metadata.uid != sandbox.metadata.uid || owned.metadata.uid != namespace.metadata.uid { + return Err("Governed service namespace or Sandbox incarnation changed".into()); + } + let sandbox_uid = live + .metadata + .uid + .as_deref() + .filter(|uid| !uid.is_empty()) + .ok_or("Sandbox UID missing")?; + let namespace_uid = owned + .metadata + .uid + .as_deref() + .filter(|uid| !uid.is_empty()) + .ok_or("Namespace UID missing")?; + let workspace = live.namespace().ok_or("Sandbox workspace missing")?; + let name = live.name_any(); + let mut task_identity = Value::Null; + let mut task_authorization = None; + let mut task_generation = None; + if let Some(owner) = live.metadata.owner_references.as_ref().and_then(|owners| { + owners.iter().find(|owner| { + owner.kind == "KarsTask" + && owner.api_version == "kars.azure.com/v1alpha1" + && owner.controller == Some(true) + }) + }) { + let tasks: Api = Api::namespaced(client.clone(), &workspace); + let task = tasks.get(&owner.name).await.map_err(api_error)?; + let authorization = authorized_task(&task, &workspace, &name, &owner.name, &owner.uid) + .ok_or("Task UID/effective authorization does not bind this Sandbox")?; + task_identity = json!({"namespace":workspace,"name":owner.name,"uid":owner.uid}); + task_authorization = Some(authorization); + task_generation = task.metadata.generation; + } + let credential = credentials::ensure(client, &live, &owned).await?; + Ok(Projection { + identity: json!({"sandbox":{"namespace":workspace,"name":name,"uid":sandbox_uid}, + "namespace_uid":namespace_uid,"task":task_identity,"task_authorization":task_authorization, + "task_generation":task_generation,"managed":true}), + credential, + }) +} + +pub fn mount(pod: &mut Value) { + pod["volumes"] + .as_array_mut() + .expect("pod volumes array") + .push(json!({ + "name":"governed-services-control","secret":{"secretName":SECRET, + "items":[{"key":"control-token","path":"control-token"}]}})); + for container in pod["containers"] + .as_array_mut() + .expect("pod containers array") + { + if container["name"] == "inference-router" { + container["volumeMounts"].as_array_mut().expect("router mounts array").push(json!({ + "name":"governed-services-control","mountPath":"/etc/kars/services","readOnly":true})); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + fn task() -> KarsTask { + let mut task:KarsTask=serde_json::from_value(json!({ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsTask", + "metadata":{"name":"task","namespace":"workspace","uid":"task-uid","generation":1}, + "spec":{"objective":"inspect","envelope":{"tier":1,"authorityCeiling":1,"delegationDepth":0},"execution":{"launch":true}}, + "status":{"phase":"Ready","observedGeneration":1,"sandboxRef":{"name":"sandbox"}, + "conditions":[{"type":"Ready","status":"True","reason":"Valid","message":"validated","lastTransitionTime":"2026-09-08T00:00:00Z"}]}, + })).unwrap(); + task.status.as_mut().unwrap().envelope_digest = Some(task.spec.authorization_digest()); + task + } + #[test] + fn service_task_attribution_requires_full_authorization_and_current_uid() { + let mut task = task(); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "task-uid").is_some()); + assert!(authorized_task(&task, "other", "sandbox", "task", "task-uid").is_none()); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "old-task-uid").is_none()); + task.spec.blueprint = Some(crate::kars_task::TaskBlueprint { + instructions: Some("changed after approval".into()), + ..Default::default() + }); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "task-uid").is_none()); + } + #[test] + fn envelope_only_hash_and_unsupported_launch_budgets_never_authorize_projection() { + let mut task = task(); + task.status.as_mut().unwrap().envelope_digest = Some(task.spec.envelope.digest()); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "task-uid").is_none()); + task.spec.envelope.budget = Some(crate::kars_task::TaskBudget { + tokens: Some(10), + ..Default::default() + }); + task.status.as_mut().unwrap().envelope_digest = Some(task.spec.authorization_digest()); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "task-uid").is_none()); + } + #[test] + fn control_credential_is_never_mounted_in_agent_containers() { + let mut pod = json!({"volumes":[],"containers":[ + {"name":"openclaw","volumeMounts":[]},{"name":"agent","volumeMounts":[]}, + {"name":"inference-router","volumeMounts":[]}]}); + mount(&mut pod); + assert_eq!(pod["containers"][0]["volumeMounts"], json!([])); + assert_eq!(pod["containers"][1]["volumeMounts"], json!([])); + assert_eq!( + pod["containers"][2]["volumeMounts"][0]["mountPath"], + "/etc/kars/services" + ); + } + #[test] + fn standard_control_tokens_are_bounded_and_distinct() { + let first = crate::providers::signing::generate_service_token(); + assert_eq!(first.len(), 64); + assert!(first.bytes().all(|byte| byte.is_ascii_alphanumeric())); + assert_ne!(first, crate::providers::signing::generate_service_token()); + } +} diff --git a/controller/src/reconciler/governed_services/continuity_tests.rs b/controller/src/reconciler/governed_services/continuity_tests.rs new file mode 100644 index 000000000..8cb580eed --- /dev/null +++ b/controller/src/reconciler/governed_services/continuity_tests.rs @@ -0,0 +1,489 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::sre_registration::{ + AGENT_SECRET, CONTROL_VERSION, EPOCH, KarsSRERegistration, OWNER, PRIVATE_SECRET, ROUTER_SA, +}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use serde_json::{Value, json}; +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, +}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const REG: &str = "/apis/kars.azure.com/v1alpha1/karssreregistrations/canonical"; +const SRE: &str = "/apis/kars.azure.com/v1alpha1/namespaces/kars-system/karssandboxes/sre"; +const NORMAL: &str = "/apis/kars.azure.com/v1alpha1/namespaces/kars-system/karssandboxes/normal"; + +struct State { + objects: BTreeMap, + calls: Vec<(String, String, Value)>, + allowed: bool, +} + +fn failure(code: u16) -> ResponseTemplate { + ResponseTemplate::new(code).set_body_json(json!({"kind":"Status","apiVersion":"v1","status":"Failure", + "reason":if code==404 {"NotFound"} else {"Conflict"},"code":code,"message":"PRIVATE_SENTINEL"})) +} + +fn merge(value: &mut Value, patch: &Value) { + if let Some(object) = patch.as_object() { + if !value.is_object() { + *value = json!({}); + } + for (key, item) in object { + if item.is_null() { + value.as_object_mut().unwrap().remove(key); + } else { + merge(&mut value[key], item); + } + } + } else { + *value = patch.clone(); + } +} + +fn registration() -> KarsSRERegistration { + let mut reg: KarsSRERegistration = serde_json::from_value(json!({ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsSRERegistration", + "metadata":{"name":"canonical","uid":"registration","resourceVersion":"1","generation":1}, + "spec":{"controller":{"namespace":{"name":"kars-system","uid":"system"}, + "deployment":{"name":"kars-controller","uid":"controller"},"release":"kars"}, + "sandbox":{"namespace":"kars-system","name":"sre","uid":"source-sre"}, + "runtimeNamespace":{"name":"kars-sre","uid":"namespace-sre"},"enabled":true} + })) + .unwrap(); + reg.status = Some( + serde_json::from_value(json!({"phase":"Ready","observedGeneration":1, + "privacyEpoch":reg.epoch(),"privacyRevision":crate::sre_privacy::REVISION, + "legacySecretAccessDenied":true,"routerServiceAccountUid":"private-sa"})) + .unwrap(), + ); + reg +} + +fn initial() -> State { + let reg = registration(); + let epoch = reg.epoch(); + let mut objects = BTreeMap::new(); + objects.insert(REG.into(), serde_json::to_value(®).unwrap()); + objects.insert( + "/api/v1/namespaces/kars-system".into(), + json!({"metadata":{"name":"kars-system","uid":"system","resourceVersion":"1"}}), + ); + objects.insert("/apis/apps/v1/namespaces/kars-system/deployments/kars-controller".into(), + json!({"metadata":{"name":"kars-controller","namespace":"kars-system","uid":"controller","resourceVersion":"1"}})); + for name in ["sre", "normal"] { + let namespace = format!("kars-{name}"); + let source_uid = format!("source-{name}"); + let namespace_uid = format!("namespace-{name}"); + let source = json!({"apiVersion":"kars.azure.com/v1alpha1","kind":"KarsSandbox", + "metadata":{"name":name,"namespace":"kars-system","uid":source_uid,"resourceVersion":"1","generation":1, + "finalizers":[super::super::namespace_ownership::FINALIZER], + "labels":if name=="sre" {json!({"kars.azure.com/role":"sre"})} else {json!({})}, + "annotations":{NAMESPACE_UID:namespace_uid}}, + "spec":{"runtime":{"kind":"Hermes","hermes":{}},"sandbox":{"isolation":"standard"},"inferenceRef":{"name":"inference"}}, + "status":{"phase":"Running","namespace":namespace}}); + objects.insert( + format!("/apis/kars.azure.com/v1alpha1/namespaces/kars-system/karssandboxes/{name}"), + source, + ); + objects.insert(format!("/api/v1/namespaces/{namespace}"), json!({"metadata":{ + "name":namespace,"uid":namespace_uid,"resourceVersion":"1","annotations":{ + "kars.azure.com/namespace-claim-version":"v1","kars.azure.com/sandbox-namespace":"kars-system", + "kars.azure.com/sandbox-name":name,SOURCE_UID:source_uid}}})); + let version = format!("control-{name}:1"); + objects.insert(format!("/api/v1/namespaces/{namespace}/secrets/router-services-admin"), json!({ + "apiVersion":"v1","kind":"Secret","metadata":{"name":"router-services-admin","namespace":namespace, + "uid":format!("control-{name}"),"resourceVersion":"1","labels":{"app.kubernetes.io/managed-by":"kars-controller"}, + "annotations":{SOURCE_UID:source_uid,NAMESPACE_UID:namespace_uid,EPOCH:epoch}}, + "data":{"control-token":STANDARD.encode("x".repeat(64))}})); + objects.insert(format!("/apis/apps/v1/namespaces/{namespace}/deployments/{name}"), json!({ + "apiVersion":"apps/v1","kind":"Deployment","metadata":{"name":name,"namespace":namespace, + "uid":format!("deployment-{name}"),"resourceVersion":"1","generation":1, + "labels":{"kars.azure.com/sandbox":name}, + "managedFields":[{"manager":crate::field_managers::CLAWSANDBOX,"operation":"Apply","apiVersion":"apps/v1", + "fieldsType":"FieldsV1","fieldsV1":{"f:spec":{}}}]}, + "spec":{"replicas":1,"selector":{"matchLabels":{"kars.azure.com/sandbox":name}}, + "template":{"metadata":{"annotations":{OWNER:"registration",EPOCH:epoch,CONTROL_VERSION:version}}, + "spec":{"containers":[{"name":"inference-router","image":"router:pinned"}]}}}, + "status":{"observedGeneration":1,"updatedReplicas":1,"availableReplicas":0}})); + objects.insert(format!("/api/v1/namespaces/{namespace}/pods"), json!({ + "apiVersion":"v1","kind":"PodList","metadata":{},"items":[{"metadata":{"name":format!("pod-{name}"), + "uid":format!("pod-{name}"),"annotations":{EPOCH:epoch,CONTROL_VERSION:version}}}]})); + } + let annotations = json!({OWNER:"registration",EPOCH:epoch,SOURCE_UID:"source-sre",NAMESPACE_UID:"namespace-sre"}); + objects.insert(format!("/api/v1/namespaces/kars-sre/serviceaccounts/{ROUTER_SA}"),json!({ + "metadata":{"name":ROUTER_SA,"namespace":"kars-sre","uid":"private-sa","resourceVersion":"1","annotations":annotations}, + "automountServiceAccountToken":false})); + for name in [PRIVATE_SECRET, AGENT_SECRET] { + let mut annotations = annotations.clone(); + annotations["kars.azure.com/sre-tls-expiry"] = (chrono::Utc::now().timestamp() + 1_000_000) + .to_string() + .into(); + let values = if name == PRIVATE_SECRET { + vec![ + ("server-cert.pem", "test"), + ("server-key.pem", "test"), + ("agent-token", "opaque"), + ("agent-ca.crt", "ca"), + ] + } else { + vec![ + ("token", "opaque"), + ("ca.crt", "ca"), + ("namespace", "kars-sre"), + ] + }; + let mut data: BTreeMap<_, _> = values + .into_iter() + .map(|(key, value)| (key.to_string(), STANDARD.encode(value))) + .collect(); + if name == PRIVATE_SECRET { + data.insert( + "kube-expires-at".into(), + STANDARD.encode((chrono::Utc::now() + chrono::Duration::hours(1)).to_rfc3339()), + ); + data.insert("kube-token".into(), STANDARD.encode("kube-token")); + } + objects.insert(format!("/api/v1/namespaces/kars-sre/secrets/{name}"),json!({ + "kind":"Secret","apiVersion":"v1","metadata":{"name":name,"namespace":"kars-sre","uid":name,"resourceVersion":"1", + "annotations":annotations},"data":data})); + } + for (binding, role, rule) in [ + ( + "kars-sre-private-reader", + "kars-sre-private-diagnostics", + json!({"apiGroups":[""],"resources":["secrets","pods"],"verbs":["get","list","watch"]}), + ), + ( + "kars-sre-private-author", + "kars-sre-action-author", + json!({"apiGroups":["kars.azure.com"],"resources":["karssreactions"],"verbs":["create"]}), + ), + ( + "kars-sre-private-renew", + "kars-sre-router-renew", + json!({"apiGroups":["kars.azure.com"],"resources":["karssreregistrations"],"resourceNames":["canonical"],"verbs":["get","renew"]}), + ), + ] { + objects.insert( + format!("/apis/rbac.authorization.k8s.io/v1/clusterroles/{role}"), + json!({ + "metadata":{"name":role},"rules":[rule]}), + ); + objects.insert(format!("/apis/rbac.authorization.k8s.io/v1/clusterrolebindings/{binding}"),json!({ + "metadata":{"name":binding,"uid":binding,"resourceVersion":"1","annotations":annotations}, + "roleRef":{"apiGroup":"rbac.authorization.k8s.io","kind":"ClusterRole","name":role}, + "subjects":[{"kind":"ServiceAccount","name":ROUTER_SA,"namespace":"kars-sre"}]})); + } + objects.insert("/apis/rbac.authorization.k8s.io/v1/namespaces/kars-sre/roles/sre-api-self-renew".into(),json!({ + "metadata":{"name":"sre-api-self-renew","namespace":"kars-sre","uid":"renew-role","resourceVersion":"1","annotations":annotations}, + "rules":[{"apiGroups":[""],"resources":["serviceaccounts/token"],"resourceNames":[ROUTER_SA],"verbs":["create"]}]})); + objects.insert("/apis/rbac.authorization.k8s.io/v1/namespaces/kars-sre/rolebindings/sre-api-self-renew".into(),json!({ + "metadata":{"name":"sre-api-self-renew","namespace":"kars-sre","uid":"renew-binding","resourceVersion":"1","annotations":annotations}, + "roleRef":{"apiGroup":"rbac.authorization.k8s.io","kind":"Role","name":"sre-api-self-renew"}, + "subjects":[{"kind":"ServiceAccount","name":ROUTER_SA,"namespace":"kars-sre"}]})); + State { + objects, + calls: Vec::new(), + allowed: false, + } +} + +async fn fixture() -> (MockServer, Client, Arc>) { + let server = MockServer::start().await; + let state = Arc::new(Mutex::new(initial())); + let handler = state.clone(); + Mock::given(|_: &wiremock::Request|true).respond_with(move |request: &wiremock::Request| { + let mut state=handler.lock().unwrap(); + let path=request.url.path(); + let method=request.method.as_str(); + let body:Value=request.body_json().unwrap_or(Value::Null); + state.calls.push((method.into(),path.into(),body.clone())); + if method=="GET" { + if let Some(value)=state.objects.get(path) {return ResponseTemplate::new(200).set_body_json(value);} + if path=="/api/v1/secrets" { + return ResponseTemplate::new(200).set_body_json(json!({"apiVersion":"v1","kind":"SecretList","metadata":{}, + "items":state.objects.values().filter(|value|value["metadata"]["name"]=="router-services-admin").collect::>()})); + } + if path=="/api/v1/namespaces/kars-sre/secrets" { + return ResponseTemplate::new(200).set_body_json(json!({"apiVersion":"meta.k8s.io/v1","kind":"PartialObjectMetadataList","metadata":{},"items":[]})); + } + if path.ends_with("/clusterrolebindings") || path.ends_with("/rolebindings") { + return ResponseTemplate::new(200).set_body_json(json!({"apiVersion":"rbac.authorization.k8s.io/v1", + "kind":if path.ends_with("/clusterrolebindings") {"ClusterRoleBindingList"} else {"RoleBindingList"},"metadata":{},"items":[]})); + } + if path.contains("/validatingadmissionpolicies/") { + return ResponseTemplate::new(200).set_body_json(json!({"metadata":{"name":path.rsplit('/').next().unwrap(),"generation":1}, + "spec":{"failurePolicy":"Fail","validations":[]},"status":{"observedGeneration":1,"typeChecking":{}}})); + } + if path.contains("/validatingadmissionpolicybindings/") { + let name=path.rsplit('/').next().unwrap(); + return ResponseTemplate::new(200).set_body_json(json!({"metadata":{"name":name},"spec":{"policyName":name,"validationActions":["Deny"]}})); + } + } + if method=="POST" && path.ends_with("/subjectaccessreviews") { + return ResponseTemplate::new(201).set_body_json(json!({"apiVersion":"authorization.k8s.io/v1","kind":"SubjectAccessReview", + "spec":body["spec"],"status":{"allowed":state.allowed}})); + } + if method=="PATCH" { + let key=path.strip_suffix("/status").unwrap_or(path); + let Some(value)=state.objects.get_mut(key) else {return failure(404)}; + if !path.ends_with("/status") { + assert_eq!(body["metadata"]["uid"],value["metadata"]["uid"]); + assert_eq!(body["metadata"]["resourceVersion"],value["metadata"]["resourceVersion"]); + } + merge(value,&body); + if let Some(material)=body["stringData"]["control-token"].as_str() { + value["data"]["control-token"]=STANDARD.encode(material).into(); + value.as_object_mut().unwrap().remove("stringData"); + } + value["metadata"]["resourceVersion"]=(value["metadata"]["resourceVersion"].as_str().unwrap().parse::().unwrap()+1).to_string().into(); + return ResponseTemplate::new(200).set_body_json(value.clone()); + } + if method=="DELETE" { + let Some(value)=state.objects.get(path) else {return failure(404)}; + assert_eq!(body["preconditions"]["uid"],value["metadata"]["uid"]); + assert_eq!(body["preconditions"]["resourceVersion"],value["metadata"]["resourceVersion"]); + state.objects.remove(path); + return ResponseTemplate::new(200).set_body_json(json!({"apiVersion":"v1","kind":"Status","status":"Success"})); + } + failure(404) + }).mount(&server).await; + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let client = Client::try_from(kube::Config::new(server.uri().parse().unwrap())).unwrap(); + (server, client, state) +} + +fn context(client: Client) -> Arc { + Arc::new(super::super::Context { + client, + wi_client_id: String::new(), + inference_router_image: String::new(), + sandbox_image: String::new(), + openai_endpoint: String::new(), + foundry_endpoint: String::new(), + foundry_project_endpoint: String::new(), + foundry_deployments: String::new(), + imds_client_id: String::new(), + content_safety_endpoint: String::new(), + fedcred: None, + byo_strict: false, + dev_openai_api_key: String::new(), + dev_provider: String::new(), + dev_copilot_github_token: String::new(), + anthropic_api_key: String::new(), + anthropic_endpoint: String::new(), + ollama_endpoint: String::new(), + openai_moderation_api_key: String::new(), + openai_moderation_endpoint: String::new(), + dev_profile: false, + cluster_name: None, + cluster_uid: String::new(), + agent_id_cache: Arc::new(crate::agent_id_provisioning::ProvisionerCache::new()), + }) +} + +fn object(state: &Arc>, key: &str) -> T { + serde_json::from_value(state.lock().unwrap().objects[key].clone()).unwrap() +} + +#[tokio::test] +async fn full_sandbox_reconciler_quarantines_canonical_sre_before_the_early_privacy_return() { + let (_server, client, state) = fixture().await; + state.lock().unwrap().allowed = true; + let sandbox: KarsSandbox = object(&state, SRE); + super::super::reconcile(Arc::new(sandbox), context(client)) + .await + .unwrap(); + let state = state.lock().unwrap(); + assert_eq!( + state.objects["/apis/apps/v1/namespaces/kars-sre/deployments/sre"]["spec"]["replicas"], + 0 + ); + assert_eq!( + state.objects["/api/v1/namespaces/kars-sre/secrets/router-services-admin"]["metadata"]["annotations"] + [credentials::RETIRED], + "true" + ); + assert!( + state + .calls + .iter() + .all(|(_, _, body)| body.get("stringData").is_none()) + ); + assert!( + !state + .calls + .iter() + .any(|(method, path, _)| method == "POST" && path.ends_with("/secrets")) + ); +} + +#[tokio::test] +async fn full_sandbox_reconciler_does_not_quarantine_a_foreign_namespace() { + let (_server, client, state) = fixture().await; + { + let mut state = state.lock().unwrap(); + state.allowed = true; + state + .objects + .get_mut("/api/v1/namespaces/kars-sre") + .unwrap()["metadata"]["annotations"][SOURCE_UID] = "foreign".into(); + } + let sandbox: KarsSandbox = object(&state, SRE); + assert!( + super::super::reconcile(Arc::new(sandbox), context(client)) + .await + .is_err() + ); + let state = state.lock().unwrap(); + assert_eq!( + state.objects["/apis/apps/v1/namespaces/kars-sre/deployments/sre"]["spec"]["replicas"], + 1 + ); + assert!( + !state + .calls + .iter() + .any(|(_, path, _)| path.contains("/secrets/")) + ); +} + +#[tokio::test] +async fn interleaved_full_sre_and_sandbox_credential_reconcile_keep_qualified_epoch_during_health_delay() + { + let (_server, client, state) = fixture().await; + let sandbox: KarsSandbox = object(&state, NORMAL); + let namespace: Namespace = object(&state, "/api/v1/namespaces/kars-normal"); + for _ in 0..3 { + let reg: KarsSRERegistration = object(&state, REG); + crate::sre_authority::reconcile(&client, ®) + .await + .unwrap(); + super::ensure(&client, &sandbox, &namespace).await.unwrap(); + let state = state.lock().unwrap(); + assert_eq!(state.objects[REG]["status"]["phase"], "Ready"); + for name in ["sre", "normal"] { + assert_eq!( + state.objects[&format!("/apis/apps/v1/namespaces/kars-{name}/deployments/{name}")] + ["spec"]["replicas"], + 1 + ); + } + assert!( + state + .calls + .iter() + .all(|(_, _, body)| body.get("stringData").is_none()) + ); + } + state.lock().unwrap().allowed = true; + let reg: KarsSRERegistration = object(&state, REG); + assert!( + crate::sre_authority::reconcile(&client, ®) + .await + .is_err() + ); + assert!(super::ensure(&client, &sandbox, &namespace).await.is_err()); + let state = state.lock().unwrap(); + assert_eq!(state.objects[REG]["status"]["phase"], "Blocked"); + assert_eq!( + state.objects["/apis/apps/v1/namespaces/kars-normal/deployments/normal"]["spec"]["replicas"], + 0 + ); +} + +#[tokio::test] +async fn credential_transition_qualifies_before_authority_dependent_sre_availability() { + let (_server, client, state) = fixture().await; + { + let mut state = state.lock().unwrap(); + let mut old = state.objects["/api/v1/namespaces/kars-normal/pods"]["items"][0].clone(); + old["metadata"]["name"] = "old-cached-router".into(); + old["metadata"]["uid"] = "old-cached-router".into(); + old["metadata"]["deletionTimestamp"] = "2026-09-08T00:00:00Z".into(); + old["metadata"]["annotations"][CONTROL_VERSION] = "old-token".into(); + state + .objects + .get_mut("/api/v1/namespaces/kars-normal/pods") + .unwrap()["items"] + .as_array_mut() + .unwrap() + .push(old); + } + let reg: KarsSRERegistration = object(&state, REG); + assert!( + crate::sre_authority::reconcile(&client, ®) + .await + .is_err() + ); + let sandbox: KarsSandbox = object(&state, NORMAL); + let namespace: Namespace = object(&state, "/api/v1/namespaces/kars-normal"); + let sre: KarsSandbox = object(&state, SRE); + let sre_namespace: Namespace = object(&state, "/api/v1/namespaces/kars-sre"); + assert!(super::ensure(&client, &sandbox, &namespace).await.is_err()); + assert!( + crate::sre_authority::pod::authorize(&client, &sre, &sre_namespace) + .await + .is_err() + ); + { + let mut state = state.lock().unwrap(); + assert_eq!(state.objects[REG]["status"]["phase"], "Migrating"); + assert_eq!( + state.objects["/apis/apps/v1/namespaces/kars-normal/deployments/normal"]["spec"]["replicas"], + 1 + ); + assert!( + state + .calls + .iter() + .all(|(_, _, body)| body.get("stringData").is_none()) + ); + state + .objects + .get_mut("/api/v1/namespaces/kars-normal/pods") + .unwrap()["items"] + .as_array_mut() + .unwrap() + .retain(|pod| pod["metadata"]["uid"] != "old-cached-router"); + for name in ["sre", "normal"] { + assert_eq!( + state.objects[&format!("/apis/apps/v1/namespaces/kars-{name}/deployments/{name}")] + ["status"]["availableReplicas"], + 0 + ); + } + } + let reg: KarsSRERegistration = object(&state, REG); + crate::sre_authority::reconcile(&client, ®) + .await + .unwrap(); + super::ensure(&client, &sandbox, &namespace).await.unwrap(); + assert!( + crate::sre_authority::pod::authorize(&client, &sre, &sre_namespace) + .await + .unwrap() + .is_some() + ); + let state = state.lock().unwrap(); + assert_eq!(state.objects[REG]["status"]["phase"], "Ready"); + assert_eq!( + state.objects["/apis/apps/v1/namespaces/kars-sre/deployments/sre"]["status"]["availableReplicas"], + 0 + ); + assert_eq!( + state.objects["/apis/apps/v1/namespaces/kars-normal/deployments/normal"]["status"]["availableReplicas"], + 0 + ); + assert!( + state + .calls + .iter() + .all(|(_, _, body)| body.get("stringData").is_none()) + ); +} diff --git a/controller/src/reconciler/governed_services/credential_tests.rs b/controller/src/reconciler/governed_services/credential_tests.rs new file mode 100644 index 000000000..c85991897 --- /dev/null +++ b/controller/src/reconciler/governed_services/credential_tests.rs @@ -0,0 +1,650 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + credentials::{RETIRED, REVISION, VERSION}, + *, +}; +use crate::sre_registration::{EPOCH, KarsSRERegistration}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use serde_json::{Value, json}; +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, +}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const NS: &str = "kars-normal"; +const SECRETS: &str = "/api/v1/namespaces/kars-normal/secrets"; +const REG: &str = "/apis/kars.azure.com/v1alpha1/karssreregistrations/canonical"; +const DEPLOY: &str = "/apis/apps/v1/namespaces/kars-normal/deployments/normal"; + +fn source() -> KarsSandbox { + serde_json::from_value(json!({ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsSandbox", + "metadata":{"name":"normal","namespace":"kars-system","uid":"source","resourceVersion":"1", + "annotations":{NAMESPACE_UID:"namespace"}}, + "spec":{"runtime":{"kind":"BYO","byo":{"image":"pinned:test","contractVersion":"v1"}}, + "sandbox":{"isolation":"standard"},"inferenceRef":{"name":"inference"}} + })) + .unwrap() +} + +fn namespace() -> Namespace { + serde_json::from_value(json!({"apiVersion":"v1","kind":"Namespace", + "metadata":{"name":NS,"uid":"namespace","resourceVersion":"1","annotations":{ + "kars.azure.com/namespace-claim-version":"v1","kars.azure.com/sandbox-namespace":"kars-system", + "kars.azure.com/sandbox-name":"normal",SOURCE_UID:"source"}}})).unwrap() +} + +fn secret(epoch: Option<&str>, stamped: bool) -> Value { + let mut value = json!({"apiVersion":"v1","kind":"Secret","type":"Opaque", + "metadata":{"name":SECRET,"namespace":NS,"uid":"secret","resourceVersion":"7", + "labels":{"app.kubernetes.io/managed-by":"kars-controller","customer":"preserve"}, + "annotations":{SOURCE_UID:"source",NAMESPACE_UID:"namespace"}}, + "data":{"control-token":STANDARD.encode("x".repeat(64))}}); + if let Some(epoch) = epoch { + value["metadata"]["annotations"][EPOCH] = epoch.into(); + } + if stamped { + value["metadata"]["annotations"][REVISION] = crate::sre_privacy::REVISION.into(); + } + value +} + +fn deployment() -> Value { + json!({"apiVersion":"apps/v1","kind":"Deployment", + "metadata":{"name":"normal","namespace":NS,"uid":"deployment","resourceVersion":"3", + "labels":{"kars.azure.com/sandbox":"normal"}, + "managedFields":[{"manager":crate::field_managers::CLAWSANDBOX,"operation":"Apply", + "apiVersion":"apps/v1","fieldsType":"FieldsV1","fieldsV1":{"f:spec":{}}}]}, + "spec":{"replicas":1,"selector":{"matchLabels":{"kars.azure.com/sandbox":"normal"}}, + "template":{"metadata":{"annotations":{"customer":"preserve"}}, + "spec":{"containers":[{"name":"agent","image":"customer:pinned"}]}}}}) +} + +fn registration() -> Value { + let mut reg: KarsSRERegistration = serde_json::from_value(json!({ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsSRERegistration", + "metadata":{"name":"canonical","uid":"registration","resourceVersion":"1","generation":1}, + "spec":{"controller":{"namespace":{"name":"kars-system","uid":"system"}, + "deployment":{"name":"kars-controller","uid":"controller"},"release":"kars"}, + "sandbox":{"namespace":"kars-system","name":"sre","uid":"sre-source"}, + "runtimeNamespace":{"name":"kars-sre","uid":"sre-namespace"},"enabled":true} + })) + .unwrap(); + reg.status = Some( + serde_json::from_value(json!({ + "phase":"Ready","observedGeneration":1,"legacySecretAccessDenied":true, + "privacyRevision":crate::sre_privacy::REVISION,"privacyEpoch":reg.epoch(), + "routerServiceAccountUid":"sre-router", + })) + .unwrap(), + ); + serde_json::to_value(reg).unwrap() +} + +#[derive(Default)] +struct State { + objects: BTreeMap, + calls: Vec<(String, String, Value)>, + pods: Vec, + allow_verb: Option, + sar_error: bool, + revoke_on_consumer_read: bool, + conflict: bool, + fail_policy: bool, + alias: bool, + wrong_write_stamp: bool, +} + +fn failure(code: u16) -> ResponseTemplate { + ResponseTemplate::new(code).set_body_json(json!({"apiVersion":"v1","kind":"Status","code":code, + "reason":if code==404 {"NotFound"} else if code==409 {"Conflict"} else {"Forbidden"}, + "status":"Failure","message":"PRIVATE_RESPONSE_SENTINEL"})) +} + +fn merge(value: &mut Value, patch: &Value) { + if let Some(entries) = patch.as_object() { + if !value.is_object() { + *value = json!({}); + } + for (key, item) in entries { + if item.is_null() { + value.as_object_mut().unwrap().remove(key); + } else { + merge(&mut value[key], item); + } + } + } else { + *value = patch.clone(); + } +} + +async fn fixture() -> (MockServer, Client, Arc>) { + let server = MockServer::start().await; + let state = Arc::new(Mutex::new(State::default())); + let handler = state.clone(); + Mock::given(|_: &wiremock::Request| true).respond_with(move |request: &wiremock::Request| { + let mut state = handler.lock().unwrap(); + let path = request.url.path(); + let body: Value = request.body_json().unwrap_or(Value::Null); + let method = request.method.as_str(); + state.calls.push((method.into(), path.into(), body.clone())); + if method == "POST" && path == "/apis/authorization.k8s.io/v1/subjectaccessreviews" { + return ResponseTemplate::new(201).set_body_json(json!({ + "apiVersion":"authorization.k8s.io/v1","kind":"SubjectAccessReview", + "spec":body["spec"],"status":{ + "allowed": state.allow_verb.as_deref() == body["spec"]["resourceAttributes"]["verb"].as_str(), + "evaluationError":if state.sar_error {"PRIVATE_RESPONSE_SENTINEL"} else {""}, + }})); + } + if method == "GET" { + if path == DEPLOY && state.revoke_on_consumer_read { + state.allow_verb = Some("watch".into()); + } + if let Some(value) = state.objects.get(path) { + return ResponseTemplate::new(200).set_body_json(value); + } + if path == "/api/v1/namespaces/kars-normal/pods" { + assert!(request.url.query_pairs().any(|(key, value)| + key == "labelSelector" && value == "kars.azure.com/sandbox=normal") || !state.objects.contains_key(DEPLOY)); + return ResponseTemplate::new(200).set_body_json(json!({ + "apiVersion":"v1","kind":"PodList","metadata":{},"items":state.pods})); + } + if path.starts_with("/apis/admissionregistration.k8s.io/v1/validatingadmissionpolicies/") { + return ResponseTemplate::new(200).set_body_json(json!({ + "metadata":{"name":path.rsplit('/').next().unwrap(),"generation":1}, + "spec":{"failurePolicy":if state.fail_policy {"Ignore"} else {"Fail"},"validations":[]}, + "status":{"observedGeneration":1,"typeChecking":{}}})); + } + if path.starts_with("/apis/admissionregistration.k8s.io/v1/validatingadmissionpolicybindings/") { + let name = path.rsplit('/').next().unwrap(); + return ResponseTemplate::new(200).set_body_json(json!({ + "metadata":{"name":name},"spec":{"policyName":name,"validationActions":["Deny"]}})); + } + if path == "/api/v1/namespaces/kars-sre/secrets" { + assert!(request.headers.get("accept").unwrap().to_str().unwrap().contains("PartialObjectMetadataList")); + return ResponseTemplate::new(200).set_body_json(json!({ + "apiVersion":"meta.k8s.io/v1","kind":"PartialObjectMetadataList","metadata":{}, + "items":if state.alias { vec![json!({"metadata":{"name":"alias","uid":"alias","resourceVersion":"1", + "annotations":{"kubernetes.io/service-account.name":"sre-api-router"}}})] } else {vec![]}})); + } + } + if (method == "POST" && path == SECRETS) + || (method == "PATCH" && path == format!("{SECRETS}/{SECRET}")) + { + if state.conflict { + return failure(409); + } + let key = format!("{SECRETS}/{SECRET}"); + let mut value = if method == "PATCH" { + let existing = state.objects.get(&key).unwrap().clone(); + assert_eq!(body["metadata"]["uid"], existing["metadata"]["uid"]); + assert_eq!(body["metadata"]["resourceVersion"], existing["metadata"]["resourceVersion"]); + existing + } else { + json!({"apiVersion":"v1","kind":"Secret","metadata":{"uid":"new-secret","resourceVersion":"0"}}) + }; + let version = if method == "POST" {8} else { + value["metadata"]["resourceVersion"].as_str().unwrap().parse::().unwrap()+1 + }; + merge(&mut value, &body); + value["metadata"]["resourceVersion"] = version.to_string().into(); + if let Some(material) = body["stringData"]["control-token"].as_str() { + value["data"]["control-token"] = STANDARD.encode(material).into(); + } + value.as_object_mut().unwrap().remove("stringData"); + if state.wrong_write_stamp { + value["metadata"]["annotations"][REVISION] = "wrong".into(); + value["metadata"]["annotations"].as_object_mut().unwrap().remove(EPOCH); + } + state.objects.insert(key, value.clone()); + return ResponseTemplate::new(if method == "POST" {201} else {200}).set_body_json(value); + } + if method == "PATCH" && path == DEPLOY { + let value = state.objects.get_mut(path).unwrap(); + assert_eq!(body["metadata"]["uid"], value["metadata"]["uid"]); + assert_eq!(body["metadata"]["resourceVersion"], value["metadata"]["resourceVersion"]); + merge(value, &body); + value["metadata"]["resourceVersion"] = "4".into(); + return ResponseTemplate::new(200).set_body_json(value.clone()); + } + failure(404) + }).mount(&server).await; + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let client = Client::try_from(kube::Config::new(server.uri().parse().unwrap())).unwrap(); + (server, client, state) +} + +fn enroll(state: &mut State) -> String { + let reg = registration(); + let epoch = reg["status"]["privacyEpoch"].as_str().unwrap().to_string(); + state.objects.insert(REG.into(), reg); + state.objects.insert( + "/api/v1/namespaces/kars-system".into(), + json!({"metadata":{"name":"kars-system","uid":"system","resourceVersion":"1"}}), + ); + state.objects.insert("/apis/apps/v1/namespaces/kars-system/deployments/kars-controller".into(), + json!({"metadata":{"name":"kars-controller","namespace":"kars-system","uid":"controller","resourceVersion":"1"}})); + state.objects.insert( + "/apis/kars.azure.com/v1alpha1/namespaces/kars-system/karssandboxes/sre".into(), + json!({ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsSandbox", + "metadata":{"name":"sre","namespace":"kars-system","uid":"sre-source","resourceVersion":"1", + "labels":{"kars.azure.com/role":"sre"},"annotations":{NAMESPACE_UID:"sre-namespace"}}, + "spec":{"runtime":{"kind":"Hermes","hermes":{}},"inferenceRef":{"name":"sre-inference"}}}), + ); + state.objects.insert("/api/v1/namespaces/kars-sre".into(), json!({ + "metadata":{"name":"kars-sre","uid":"sre-namespace","resourceVersion":"1","annotations":{ + "kars.azure.com/namespace-claim-version":"v1","kars.azure.com/sandbox-namespace":"kars-system", + "kars.azure.com/sandbox-name":"sre",SOURCE_UID:"sre-source"}}})); + epoch +} + +fn secret_writes(state: &State) -> usize { + state + .calls + .iter() + .filter(|(method, path, _)| method != "GET" && path.starts_with(SECRETS)) + .count() +} + +fn token_issuances(state: &State) -> usize { + state + .calls + .iter() + .filter(|(method, path, body)| { + method != "GET" + && path.starts_with(SECRETS) + && body["stringData"]["control-token"].is_string() + }) + .count() +} + +#[tokio::test] +async fn standalone_control_issuance_requires_real_get_list_watch_denial_without_registration() { + let (_server, client, state) = fixture().await; + let projection = credentials::ensure(&client, &source(), &namespace()) + .await + .unwrap(); + let state = state.lock().unwrap(); + let secret = &state.objects[&format!("{SECRETS}/{SECRET}")]; + assert_eq!( + secret["metadata"]["annotations"][REVISION], + crate::sre_privacy::REVISION + ); + assert!(secret["metadata"]["annotations"].get(EPOCH).is_none()); + assert_eq!(projection.version, "new-secret:8"); + assert_eq!(secret_writes(&state), 1); + for verb in ["get", "list", "watch"] { + assert!( + state + .calls + .iter() + .any(|(_, _, body)| body["spec"]["resourceAttributes"]["verb"] == verb) + ); + } +} + +#[tokio::test] +async fn legacy_secret_authority_or_indeterminate_reviews_prevent_any_control_issuance() { + for verb in ["get", "list", "watch", "evaluation-error"] { + let (_server, client, state) = fixture().await; + if verb == "evaluation-error" { + state.lock().unwrap().sar_error = true; + } else { + state.lock().unwrap().allow_verb = Some(verb.into()); + } + let error = credentials::ensure(&client, &source(), &namespace()) + .await + .err() + .unwrap(); + assert!(!error.contains("PRIVATE_RESPONSE_SENTINEL")); + assert_eq!(secret_writes(&state.lock().unwrap()), 0); + } +} + +#[tokio::test] +async fn current_ready_epoch_is_required_and_recorded_before_control_token_creation() { + let (_server, client, state) = fixture().await; + let epoch = enroll(&mut state.lock().unwrap()); + credentials::ensure(&client, &source(), &namespace()) + .await + .unwrap(); + let state = state.lock().unwrap(); + assert_eq!( + state.objects[&format!("{SECRETS}/{SECRET}")]["metadata"]["annotations"][EPOCH], + epoch + ); + assert_eq!( + state + .calls + .iter() + .filter(|(_, path, _)| path.contains("/validatingadmissionpolicies/")) + .count(), + 14 + ); +} + +#[tokio::test] +async fn stale_migrating_alias_or_unenforced_sre_authority_cannot_issue_or_reuse_controls() { + for failure in [ + "Migrating", + "revision", + "source-uid", + "generation", + "alias", + "policy", + ] { + let (_server, client, state) = fixture().await; + { + let mut state = state.lock().unwrap(); + let epoch = enroll(&mut state); + state + .objects + .insert(format!("{SECRETS}/{SECRET}"), secret(Some(&epoch), true)); + match failure { + "Migrating" => { + state.objects.get_mut(REG).unwrap()["status"]["phase"] = "Migrating".into() + } + "revision" => { + state.objects.get_mut(REG).unwrap()["status"]["privacyRevision"] = "old".into() + } + "source-uid" => { + state.objects.get_mut(REG).unwrap()["spec"]["sandbox"]["uid"] = + "replaced".into() + } + "generation" => { + state.objects.get_mut(REG).unwrap()["status"]["observedGeneration"] = 0.into() + } + "alias" => state.alias = true, + _ => state.fail_policy = true, + } + } + assert!( + credentials::ensure(&client, &source(), &namespace()) + .await + .is_err(), + "{failure}" + ); + let state = state.lock().unwrap(); + assert_eq!(token_issuances(&state), 0); + if failure == "Migrating" { + assert_eq!(secret_writes(&state), 0); + assert!( + state.objects[&format!("{SECRETS}/{SECRET}")]["metadata"]["annotations"] + .get(RETIRED) + .is_none() + ); + } else { + assert_eq!( + state.objects[&format!("{SECRETS}/{SECRET}")]["metadata"]["annotations"][RETIRED], + "true" + ); + } + } +} + +#[tokio::test] +async fn retired_authority_requires_current_denial_revision_before_standalone_issuance() { + for current in [false, true] { + let (_server, client, state) = fixture().await; + { + let mut state = state.lock().unwrap(); + enroll(&mut state); + let reg = state.objects.get_mut(REG).unwrap(); + reg["spec"]["enabled"] = false.into(); + reg["status"]["phase"] = "Retired".into(); + if !current { + reg["status"]["privacyRevision"] = "old".into(); + } + } + assert_eq!( + credentials::ensure(&client, &source(), &namespace()) + .await + .is_ok(), + current + ); + assert_eq!(secret_writes(&state.lock().unwrap()), usize::from(current)); + } +} + +#[tokio::test] +async fn unqualified_owned_control_token_rotates_with_cas_and_preserves_customer_metadata() { + for enrolled in [false, true] { + let (_server, client, state) = fixture().await; + let epoch = { + let mut state = state.lock().unwrap(); + state + .objects + .insert(format!("{SECRETS}/{SECRET}"), secret(None, false)); + state.objects.insert(DEPLOY.into(), deployment()); + enrolled.then(|| enroll(&mut state)) + }; + let projection = credentials::ensure(&client, &source(), &namespace()) + .await + .unwrap(); + let state = state.lock().unwrap(); + let secret = &state.objects[&format!("{SECRETS}/{SECRET}")]; + assert_eq!(secret["metadata"]["uid"], "secret"); + assert_eq!(secret["metadata"]["labels"]["customer"], "preserve"); + assert_ne!( + secret["data"]["control-token"], + STANDARD.encode("x".repeat(64)) + ); + assert_eq!( + secret["metadata"]["annotations"] + .get(EPOCH) + .and_then(Value::as_str), + epoch.as_deref() + ); + assert_eq!(projection.version, "secret:8"); + assert_eq!(secret_writes(&state), 1); + } +} + +#[tokio::test] +async fn already_qualified_control_is_reused_without_reissuing_or_touching_foreign_workloads() { + let (_server, client, state) = fixture().await; + state + .lock() + .unwrap() + .objects + .insert(format!("{SECRETS}/{SECRET}"), secret(None, true)); + assert_eq!( + credentials::ensure(&client, &source(), &namespace()) + .await + .unwrap() + .version, + "secret:7" + ); + assert_eq!(secret_writes(&state.lock().unwrap()), 0); +} + +#[tokio::test] +async fn foreign_secret_identity_or_unowned_consumers_never_receive_rotation() { + for changed in [ + "label", + "source", + "namespace", + "owner", + "deployment", + "orphan-pod", + ] { + let (_server, client, state) = fixture().await; + { + let mut state = state.lock().unwrap(); + let mut secret = secret(None, false); + match changed { + "label" => { + secret["metadata"]["labels"]["app.kubernetes.io/managed-by"] = "Helm".into() + } + "source" => secret["metadata"]["annotations"][SOURCE_UID] = "replacement".into(), + "namespace" => { + secret["metadata"]["annotations"][NAMESPACE_UID] = "replacement".into() + } + "owner" => { + secret["metadata"]["ownerReferences"] = + json!([{"apiVersion":"v1","kind":"Pod","name":"foreign","uid":"foreign"}]) + } + "deployment" => { + let mut deployment = deployment(); + deployment["metadata"]["managedFields"] = json!([]); + state.objects.insert(DEPLOY.into(), deployment); + } + _ => state.pods.push(json!({"metadata":{"name":"unreviewed"}})), + } + state.objects.insert(format!("{SECRETS}/{SECRET}"), secret); + } + assert!( + credentials::ensure(&client, &source(), &namespace()) + .await + .is_err(), + "{changed}" + ); + assert_eq!(secret_writes(&state.lock().unwrap()), 0); + } +} + +#[tokio::test] +async fn privacy_revoked_during_consumer_inventory_retires_cache_without_issuing_a_token() { + let (_server, client, state) = fixture().await; + { + let mut state = state.lock().unwrap(); + state + .objects + .insert(format!("{SECRETS}/{SECRET}"), secret(None, false)); + state.objects.insert(DEPLOY.into(), deployment()); + state.revoke_on_consumer_read = true; + } + assert!( + credentials::ensure(&client, &source(), &namespace()) + .await + .is_err() + ); + let state = state.lock().unwrap(); + assert_eq!(token_issuances(&state), 0); + assert_eq!(state.objects[DEPLOY]["spec"]["replicas"], 0); + assert_eq!( + state.objects[&format!("{SECRETS}/{SECRET}")]["metadata"]["annotations"][RETIRED], + "true" + ); +} + +#[tokio::test] +async fn standalone_privacy_loss_quarantines_then_rotates_instead_of_reusing_exposed_cache() { + let (_server, client, state) = fixture().await; + { + let mut state = state.lock().unwrap(); + state + .objects + .insert(format!("{SECRETS}/{SECRET}"), secret(None, true)); + state.objects.insert(DEPLOY.into(), deployment()); + state.allow_verb = Some("watch".into()); + } + assert!( + credentials::ensure(&client, &source(), &namespace()) + .await + .is_err() + ); + { + let mut state = state.lock().unwrap(); + assert_eq!(state.objects[DEPLOY]["spec"]["replicas"], 0); + assert_eq!( + state.objects[&format!("{SECRETS}/{SECRET}")]["data"]["control-token"], + STANDARD.encode("x".repeat(64)) + ); + assert_eq!(token_issuances(&state), 0); + state.allow_verb = None; + } + let projection = credentials::ensure(&client, &source(), &namespace()) + .await + .unwrap(); + let state = state.lock().unwrap(); + let secret = &state.objects[&format!("{SECRETS}/{SECRET}")]; + assert!(secret["metadata"]["annotations"].get(RETIRED).is_none()); + assert_ne!( + secret["data"]["control-token"], + STANDARD.encode("x".repeat(64)) + ); + assert_eq!(projection.version, "secret:9"); + assert_eq!(token_issuances(&state), 1); +} + +#[tokio::test] +async fn control_cas_conflicts_and_mismatched_write_results_are_not_reported_as_applied() { + for conflict in [true, false] { + let (_server, client, state) = fixture().await; + { + let mut state = state.lock().unwrap(); + state + .objects + .insert(format!("{SECRETS}/{SECRET}"), secret(None, false)); + state.objects.insert(DEPLOY.into(), deployment()); + state.conflict = conflict; + state.wrong_write_stamp = !conflict; + } + let error = credentials::ensure(&client, &source(), &namespace()) + .await + .err() + .unwrap(); + assert!(!error.contains("PRIVATE_RESPONSE_SENTINEL")); + assert_eq!(secret_writes(&state.lock().unwrap()), 1); + } +} + +#[tokio::test] +async fn cached_old_and_terminating_consumers_block_completion_until_they_are_gone() { + let (_server, client, state) = fixture().await; + let projection = credentials::ensure(&client, &source(), &namespace()) + .await + .unwrap(); + let mut deployment: Deployment = serde_json::from_value(deployment()).unwrap(); + projection.decorate(&mut deployment); + assert_eq!( + deployment + .spec + .as_ref() + .unwrap() + .template + .metadata + .as_ref() + .unwrap() + .annotations + .as_ref() + .unwrap()["customer"], + "preserve" + ); + assert_eq!( + deployment + .spec + .as_ref() + .unwrap() + .template + .spec + .as_ref() + .unwrap() + .containers[0] + .image + .as_deref(), + Some("customer:pinned") + ); + let old = json!({"metadata":{"name":"old","deletionTimestamp":"2026-09-08T00:00:00Z", + "annotations":{VERSION:"old"}}}); + let current = json!({"metadata":{"name":"current","annotations":{VERSION:projection.version}}}); + state.lock().unwrap().pods = vec![old, current.clone()]; + assert!( + !projection + .consumers_current(&client, NS, "normal") + .await + .unwrap() + ); + state.lock().unwrap().pods = vec![current]; + assert!( + projection + .consumers_current(&client, NS, "normal") + .await + .unwrap() + ); +} diff --git a/controller/src/reconciler/governed_services/credentials.rs b/controller/src/reconciler/governed_services/credentials.rs new file mode 100644 index 000000000..5f09fb384 --- /dev/null +++ b/controller/src/reconciler/governed_services/credentials.rs @@ -0,0 +1,361 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{NAMESPACE_UID, SECRET, SOURCE_UID, api_error}; +use crate::{crd::KarsSandbox, sre_registration::EPOCH}; +use k8s_openapi::api::{ + apps::v1::Deployment, + core::v1::{Namespace, Pod, Secret}, +}; +use kube::{ + Api, Client, ResourceExt, + api::{ListParams, Patch, PatchParams, PostParams}, +}; +use serde_json::json; + +pub(super) const REVISION: &str = "kars.azure.com/services-privacy-revision"; +pub(super) const VERSION: &str = crate::sre_registration::CONTROL_VERSION; +pub(super) const RETIRED: &str = "kars.azure.com/services-credential-retired"; + +pub(super) struct Projection { + pub(super) version: String, +} + +impl Projection { + pub(super) fn decorate(&self, deployment: &mut Deployment) { + deployment + .spec + .as_mut() + .expect("controller Deployment spec") + .template + .metadata + .get_or_insert_with(Default::default) + .annotations + .get_or_insert_with(Default::default) + .insert(VERSION.into(), self.version.clone()); + } + + pub(super) async fn consumers_current( + &self, + client: &Client, + namespace: &str, + name: &str, + ) -> Result { + let pods = Api::::namespaced(client.clone(), namespace) + .list(&ListParams::default().labels(&format!("kars.azure.com/sandbox={name}"))) + .await + .map_err(api_error)?; + // Include terminating Pods: Deployment availability alone can hide a + // still-running router that accepts its old startup-cached token. + Ok(pods.iter().all(|pod| { + pod.metadata + .annotations + .as_ref() + .and_then(|annotations| annotations.get(VERSION)) + == Some(&self.version) + })) + } +} + +fn validate(secret: &Secret, source_uid: &str, namespace: &Namespace) -> Result<(), String> { + let annotations = secret.metadata.annotations.as_ref(); + let matches = |key, value: &str| { + annotations + .and_then(|annotations| annotations.get(key)) + .map(String::as_str) + == Some(value) + }; + if secret.metadata.uid.as_deref().is_none_or(str::is_empty) + || secret + .metadata + .resource_version + .as_deref() + .is_none_or(str::is_empty) + || secret.metadata.deletion_timestamp.is_some() + || secret.metadata.name.as_deref() != Some(SECRET) + || secret.metadata.namespace != namespace.metadata.name + || secret + .metadata + .owner_references + .as_ref() + .is_some_and(|owners| !owners.is_empty()) + || secret + .metadata + .labels + .as_ref() + .and_then(|labels| labels.get("app.kubernetes.io/managed-by")) + .map(String::as_str) + != Some("kars-controller") + || !matches(SOURCE_UID, source_uid) + || !matches( + NAMESPACE_UID, + namespace.metadata.uid.as_deref().unwrap_or_default(), + ) + || secret.type_.as_deref().is_some_and(|kind| kind != "Opaque") + || secret + .data + .as_ref() + .and_then(|data| data.get("control-token")) + .is_none_or(|value| { + value.0.len() != 64 || value.0.iter().any(|byte| !byte.is_ascii_graphic()) + }) + { + return Err( + "Existing governed service credential has conflicting ownership or invalid data".into(), + ); + } + Ok(()) +} + +fn current(secret: &Secret, epoch: Option<&str>) -> bool { + let annotations = secret.metadata.annotations.as_ref(); + if annotations.is_some_and(|annotations| annotations.contains_key(RETIRED)) { + return false; + } + match epoch { + // A Ready v2 registration has already rotated its owned credentials + // and retired the old cached consumers before publishing this epoch. + Some(epoch) => annotations.and_then(|a| a.get(EPOCH)).map(String::as_str) == Some(epoch), + None => { + annotations + .and_then(|a| a.get(REVISION)) + .map(String::as_str) + == Some(crate::sre_privacy::REVISION) + } + } +} + +async fn review_consumer( + client: &Client, + namespace: &str, + name: &str, +) -> Result, String> { + let deployment = Api::::namespaced(client.clone(), namespace) + .get_opt(name) + .await + .map_err(api_error)?; + if let Some(deployment) = deployment.as_ref() { + if deployment.metadata.name.as_deref() != Some(name) + || deployment.metadata.namespace.as_deref() != Some(namespace) + || deployment.metadata.uid.as_deref().is_none_or(str::is_empty) + || deployment + .metadata + .resource_version + .as_deref() + .is_none_or(str::is_empty) + || deployment.metadata.deletion_timestamp.is_some() + || !crate::sre_authority::controller_managed(deployment, name) + { + return Err( + "Governed service credential consumer is not a live controller-owned Deployment" + .into(), + ); + } + } else if !Api::::namespaced(client.clone(), namespace) + .list(&ListParams::default()) + .await + .map_err(api_error)? + .items + .is_empty() + { + return Err( + "Governed service credential has unreviewed consumers without its Deployment".into(), + ); + } + Ok(deployment) +} + +async fn quarantine( + client: &Client, + namespace: &str, + name: &str, + secret: &Secret, +) -> Result<(), String> { + if secret + .metadata + .annotations + .as_ref() + .and_then(|a| a.get(RETIRED)) + .map(String::as_str) + != Some("true") + { + Api::::namespaced(client.clone(), namespace) + .patch(SECRET, &PatchParams::default(), &Patch::Merge(json!({ + "metadata":{"uid":secret.metadata.uid,"resourceVersion":secret.metadata.resource_version, + "annotations":{RETIRED:"true",REVISION:null,EPOCH:null}}, + }))).await.map_err(api_error)?; + } + let consumer = review_consumer(client, namespace, name).await?; + if let Some(deployment) = consumer + && deployment + .spec + .as_ref() + .and_then(|spec| spec.replicas) + .unwrap_or(1) + != 0 + { + Api::::namespaced(client.clone(), namespace) + .patch(name, &PatchParams::default(), &Patch::Merge(json!({ + "metadata":{"uid":deployment.metadata.uid,"resourceVersion":deployment.metadata.resource_version}, + "spec":{"replicas":0}, + }))).await.map_err(api_error)?; + } + Ok(()) +} + +async fn checked_epoch( + client: &Client, + namespace: &str, + name: &str, + existing: Option<&Secret>, +) -> Result, String> { + match crate::sre_authority::privacy_readiness(client, namespace).await { + Ok(crate::sre_authority::PrivacyReadiness::Qualified(epoch)) => Ok(epoch), + Ok(crate::sre_authority::PrivacyReadiness::Pending) => { + Err("SRE privacy qualification is still pending; no credential issued or reused".into()) + } + Err(error) => { + if let Some(secret) = existing { + quarantine(client, namespace, name, secret) + .await + .map_err(|failure| { + format!("{error}; owned control credential quarantine failed: {failure}") + })?; + } + Err(error) + } + } +} + +pub(in crate::reconciler) async fn quarantine_on_privacy_loss( + client: &Client, + sandbox: &KarsSandbox, + namespace: &Namespace, +) -> Result<(), String> { + let workspace = sandbox.namespace().ok_or("Sandbox workspace missing")?; + let live = Api::::namespaced(client.clone(), &workspace) + .get(&sandbox.name_any()) + .await + .map_err(api_error)?; + if live.uid() != sandbox.uid() || live.metadata.deletion_timestamp.is_some() { + return Err("Sandbox incarnation changed before credential quarantine".into()); + } + let namespace = super::super::namespace_ownership::recheck(client, &live, namespace) + .await + .map_err(|_| "Namespace authority changed before credential quarantine")?; + if crate::sre_authority::privacy_readiness(client, &namespace.name_any()) + .await + .is_ok() + { + return Ok(()); + } + let secret = Api::::namespaced(client.clone(), &namespace.name_any()) + .get_opt(SECRET) + .await + .map_err(api_error)?; + if let Some(secret) = secret { + validate( + &secret, + live.metadata.uid.as_deref().ok_or("Sandbox UID missing")?, + &namespace, + )?; + quarantine(client, &namespace.name_any(), &live.name_any(), &secret).await?; + } + Ok(()) +} + +pub(super) async fn ensure( + client: &Client, + sandbox: &KarsSandbox, + namespace: &Namespace, +) -> Result { + if !super::super::namespace_ownership::claimed(namespace, sandbox) + .map_err(|_| "Governed service credential namespace claim is invalid")? + || sandbox.metadata.deletion_timestamp.is_some() + || namespace.metadata.deletion_timestamp.is_some() + { + return Err("Governed service credential requires a live exact namespace claim".into()); + } + let source_uid = sandbox + .metadata + .uid + .as_deref() + .ok_or("Sandbox UID missing")?; + let namespace_name = namespace.name_any(); + let secrets: Api = Api::namespaced(client.clone(), &namespace_name); + let existing = secrets.get_opt(SECRET).await.map_err(api_error)?; + if let Some(secret) = existing.as_ref() { + validate(secret, source_uid, namespace)?; + } + let mut epoch = checked_epoch( + client, + &namespace_name, + &sandbox.name_any(), + existing.as_ref(), + ) + .await?; + let secret = if let Some(secret) = existing + .as_ref() + .filter(|secret| current(secret, epoch.as_deref())) + { + secret.clone() + } else { + if existing.is_some() { + review_consumer(client, &namespace_name, &sandbox.name_any()).await?; + // The ownership inventory awaited API calls. Recheck privacy at + // the actual mint/write boundary, not only before that inventory. + epoch = checked_epoch( + client, + &namespace_name, + &sandbox.name_any(), + existing.as_ref(), + ) + .await?; + } + let mut annotations = json!({ + SOURCE_UID: source_uid, NAMESPACE_UID: namespace.metadata.uid, + REVISION: crate::sre_privacy::REVISION, + }); + if let Some(epoch) = epoch.as_ref() { + annotations[EPOCH] = json!(epoch); + } + let material = crate::providers::signing::generate_service_token(); + if let Some(secret) = existing { + annotations[RETIRED] = serde_json::Value::Null; + if epoch.is_none() { + annotations[EPOCH] = serde_json::Value::Null; + } + secrets.patch(SECRET, &PatchParams::default(), &Patch::Merge(json!({ + "metadata": {"uid": secret.metadata.uid, "resourceVersion": secret.metadata.resource_version, + "annotations": annotations}, + "stringData": {"control-token": material}, + }))).await.map_err(api_error)? + } else { + let definition: Secret = serde_json::from_value(json!({ + "apiVersion": "v1", "kind": "Secret", "type": "Opaque", + "metadata": {"name": SECRET, "namespace": namespace_name, + "labels": {"app.kubernetes.io/managed-by": "kars-controller"}, + "annotations": annotations}, + "stringData": {"control-token": material}, + })) + .map_err(|_| "Governed service credential serialization failed")?; + secrets + .create(&PostParams::default(), &definition) + .await + .map_err(api_error)? + } + }; + validate(&secret, source_uid, namespace)?; + if !current(&secret, epoch.as_deref()) { + return Err( + "Governed service credential privacy stamp did not match the verified write".into(), + ); + } + Ok(Projection { + version: format!( + "{}:{}", + secret.metadata.uid.unwrap(), + secret.metadata.resource_version.unwrap() + ), + }) +} diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index 75d1b7d66..9aeba3e9c 100644 --- a/controller/src/reconciler/mod.rs +++ b/controller/src/reconciler/mod.rs @@ -38,6 +38,7 @@ pub(crate) mod byo_contract; mod credential_sources; mod dev_env; pub(crate) mod governance_mounts; +mod governed_services; mod inference; mod mcp_egress; pub(crate) mod namespace_ownership; @@ -309,6 +310,18 @@ async fn reconcile(sandbox: Arc, ctx: Arc) -> Result projection, Err(detail) => { + let detail = match governed_services::quarantine_on_privacy_loss( + client, + &sandbox, + owned_namespace.as_ref().ok_or_else(|| { + ReconcileError::Configuration("Sandbox namespace is absent".into()) + })?, + ) + .await + { + Ok(()) => detail, + Err(error) => format!("{detail}; owned control quarantine failed: {error}"), + }; crate::status::stamp_degraded(client, &sandbox, &name, "SREAuthorityNotReady", &detail) .await; return Ok(Action::requeue(Duration::from_secs(20))); @@ -1284,6 +1297,7 @@ async fn reconcile(sandbox: Arc, ctx: Arc) -> Result, ctx: Arc) -> Result projection, + Err(detail) => degrade!("GovernedServicePrivacyNotReady", detail), + }; let mut router_env = vec![ json!({"name": "AZURE_OPENAI_ENDPOINT", "value": &ctx.openai_endpoint}), json!({"name": "FOUNDRY_ENDPOINT", "value": &ctx.foundry_endpoint}), @@ -1610,6 +1636,7 @@ async fn reconcile(sandbox: Arc, ctx: Arc) -> Result, ctx: Arc) -> Result, ctx: Arc) -> Result, ctx: Arc) -> Result, ctx: Arc) -> Result Res /// Later private-credential issuers must call this immediately before issuance. /// An absent registration is safe only when the old SRE subject has no access. +#[derive(Debug, PartialEq, Eq)] +pub(crate) enum PrivacyReadiness { + Qualified(Option), + Pending, +} + +pub(crate) fn currently_qualified(reg: &KarsSRERegistration) -> bool { + reg.spec.enabled + && reg.status.as_ref().is_some_and(|status| { + status.phase == "Ready" + && status.observed_generation == reg.metadata.generation.unwrap_or_default() + && status.legacy_secret_access_denied + && status.privacy_revision.as_deref() == Some(crate::sre_privacy::REVISION) + && status.privacy_epoch.as_deref() == Some(reg.epoch().as_str()) + }) +} + pub(crate) async fn privacy_epoch( client: &Client, target_namespace: &str, ) -> Result, String> { + match privacy_readiness(client, target_namespace).await? { + PrivacyReadiness::Qualified(epoch) => Ok(epoch), + PrivacyReadiness::Pending => Err("SRE privacy migration is not Ready".into()), + } +} + +pub(crate) async fn privacy_readiness( + client: &Client, + target_namespace: &str, +) -> Result { check_secret_denial(client, target_namespace).await?; let registrations: Api = Api::all(client.clone()); let Some(reg) = registrations @@ -137,7 +164,7 @@ pub(crate) async fn privacy_epoch( .await .map_err(|e| api_error("Read SRE privacy epoch", e))? else { - return Ok(None); + return Ok(PrivacyReadiness::Qualified(None)); }; if !reg.spec.enabled && reg.status.as_ref().is_some_and(|status| { @@ -145,22 +172,32 @@ pub(crate) async fn privacy_epoch( && status.observed_generation == reg.metadata.generation.unwrap_or_default() }) { - return Ok(None); + if reg.status.as_ref().is_none_or(|status| { + !status.legacy_secret_access_denied + || status.privacy_revision.as_deref() != Some(crate::sre_privacy::REVISION) + }) { + return Err( + "Retired SRE privacy authority lacks current verified credential denial".into(), + ); + } + return Ok(PrivacyReadiness::Qualified(None)); } verify(client, ®).await?; super::admission::verify(client).await?; super::credential_guard::scan(client, ®).await?; - let status = reg - .status - .as_ref() - .ok_or("SRE authority has not been reconciled")?; - if status.phase != "Ready" - || status.observed_generation != reg.metadata.generation.unwrap_or_default() - || !status.legacy_secret_access_denied - || status.privacy_revision.as_deref() != Some(crate::sre_privacy::REVISION) - || status.privacy_epoch.as_deref() != Some(reg.epoch().as_str()) + if currently_qualified(®) { + return Ok(PrivacyReadiness::Qualified(Some(reg.epoch()))); + } + // This state is not authority to issue or reuse a credential. It only + // distinguishes ordinary migration progress from a failed live privacy + // check, so a controller cannot stop the very rollout it is waiting for. + if reg.spec.enabled + && reg.status.as_ref().is_none_or(|status| { + matches!(status.phase.as_str(), "Provisioning" | "Migrating") + && status.observed_generation == reg.metadata.generation.unwrap_or_default() + }) { - return Err("SRE privacy migration is not Ready".into()); + return Ok(PrivacyReadiness::Pending); } - Ok(status.privacy_epoch.clone()) + Err("SRE privacy authority is stale or unqualified".into()) } diff --git a/controller/src/sre_authority/migration.rs b/controller/src/sre_authority/migration.rs index a0f9cbe4d..112c48d12 100644 --- a/controller/src/sre_authority/migration.rs +++ b/controller/src/sre_authority/migration.rs @@ -4,7 +4,7 @@ use super::{api_error, check_secret_denial}; use crate::{ crd::KarsSandbox, - sre_registration::{EPOCH, KarsSRERegistration, OWNER, RUNTIME_NAMESPACE}, + sre_registration::{CONTROL_VERSION, EPOCH, KarsSRERegistration, OWNER, RUNTIME_NAMESPACE}, }; use k8s_openapi::api::{ apps::v1::Deployment, @@ -22,7 +22,7 @@ pub(super) const WAITING_FOR_CONSUMERS: &str = const WAITING_FOR_ROTATION: &str = "Waiting for owned control credential consumers to restart on the new privacy epoch"; const WAITING_FOR_ROLLOUT: &str = - "Owned control credential consumer has not completed its privacy-epoch rollout"; + "Waiting for prior-epoch or prior-control-version consumers to terminate"; pub(super) fn is_waiting(detail: &str) -> bool { matches!( @@ -170,7 +170,7 @@ pub(super) async fn stop_legacy_consumer( Ok(()) } -fn controller_managed(deployment: &Deployment, sandbox: &str) -> bool { +pub(crate) fn controller_managed(deployment: &Deployment, sandbox: &str) -> bool { deployment .metadata .labels @@ -288,41 +288,49 @@ pub(super) async fn rotate_owned_control_credentials( return Err("Control credential consumer is not controller-owned".into()); } let epoch = reg.epoch(); - if annotations.and_then(|a| a.get(EPOCH)) != Some(&epoch) { + let rotated = annotations.and_then(|a| a.get(EPOCH)) != Some(&epoch); + let secret = if rotated { let secrets: Api = Api::namespaced(client.clone(), &namespace_name); secrets.patch("router-services-admin",&PatchParams::default(),&Patch::Merge(json!({ "metadata":{"uid":secret.metadata.uid,"resourceVersion":secret.metadata.resource_version, "annotations":{EPOCH:epoch}}, "stringData":{"control-token":crate::providers::signing::generate_service_token()}, - }))).await.map_err(|e|api_error("Rotate owned control credential",e))?; - } - if deployment + }))).await.map_err(|e|api_error("Rotate owned control credential",e))? + } else { + secret + }; + let version = format!( + "{}:{}", + secret + .metadata + .uid + .as_deref() + .ok_or("Control credential UID missing")?, + secret + .metadata + .resource_version + .as_deref() + .ok_or("Control credential version missing")? + ); + let template_annotations = deployment .spec .as_ref() .and_then(|spec| spec.template.metadata.as_ref()) - .and_then(|meta| meta.annotations.as_ref()) - .and_then(|a| a.get(EPOCH)) - != Some(&epoch) + .and_then(|meta| meta.annotations.as_ref()); + let require_version = rotated + || !super::live::currently_qualified(reg) + || template_annotations + .is_some_and(|annotations| annotations.contains_key(CONTROL_VERSION)); + if template_annotations.and_then(|a| a.get(EPOCH)) != Some(&epoch) + || (require_version + && template_annotations.and_then(|a| a.get(CONTROL_VERSION)) != Some(&version)) { deployments.patch(name,&PatchParams::default(),&Patch::Merge(json!({ "metadata":{"uid":deployment.metadata.uid,"resourceVersion":deployment.metadata.resource_version}, - "spec":{"template":{"metadata":{"annotations":{EPOCH:epoch}}}}, + "spec":{"template":{"metadata":{"annotations":{EPOCH:epoch,CONTROL_VERSION:version}}}}, }))).await.map_err(|e|api_error("Restart owned control credential consumer",e))?; return Err(WAITING_FOR_ROTATION.into()); } - let desired = deployment - .spec - .as_ref() - .and_then(|spec| spec.replicas) - .unwrap_or(1); - let ready = deployment.status.as_ref().is_some_and(|status| { - status.observed_generation == deployment.metadata.generation - && status.updated_replicas.unwrap_or(0) == desired - && status.available_replicas.unwrap_or(0) == desired - }); - if !ready { - return Err(WAITING_FOR_ROLLOUT.into()); - } let selector = deployment .spec .as_ref() @@ -338,16 +346,19 @@ pub(super) async fn rotate_owned_control_credentials( .await .map_err(|e| api_error("Verify old control credential consumers have terminated", e))?; if pods.iter().any(|pod| { - pod.metadata - .annotations - .as_ref() - .and_then(|annotations| annotations.get(EPOCH)) - != Some(&epoch) + let annotations = pod.metadata.annotations.as_ref(); + annotations.and_then(|annotations| annotations.get(EPOCH)) != Some(&epoch) + || (require_version + && annotations.and_then(|annotations| annotations.get(CONTROL_VERSION)) + != Some(&version)) }) { // Rollout availability alone can exclude a still-terminating old // router which continues accepting its startup-cached control token. return Err(WAITING_FOR_ROLLOUT.into()); } + // Privacy completion is credential/template identity plus termination + // of old caches, not workload availability. In particular, the SRE + // readiness endpoint itself requires this authority to become Ready. } Ok(()) } diff --git a/controller/src/sre_authority/tests.rs b/controller/src/sre_authority/tests.rs index c4e015b2e..1148b5481 100644 --- a/controller/src/sre_authority/tests.rs +++ b/controller/src/sre_authority/tests.rs @@ -570,7 +570,7 @@ async fn old_control_token_consumers_must_disappear_even_after_rollout_counters_ "labels":{"kars.azure.com/sandbox":"sre"}, "managedFields":[{"manager":crate::field_managers::CLAWSANDBOX,"operation":"Apply","apiVersion":"apps/v1", "fieldsType":"FieldsV1","fieldsV1":{"f:spec":{}}}]}, - "spec":{"replicas":0,"selector":{"matchLabels":{"app":"sre"}},"template":{"metadata":{"annotations":{EPOCH:epoch}}, + "spec":{"replicas":0,"selector":{"matchLabels":{"app":"sre"}},"template":{"metadata":{"annotations":{EPOCH:epoch,CONTROL_VERSION:"control:1"}}, "spec":{"containers":[]}}}, "status":{"observedGeneration":2,"updatedReplicas":0,"availableReplicas":0}})); locked.objects.insert( diff --git a/controller/src/sre_registration.rs b/controller/src/sre_registration.rs index 13b552b2f..90f35597e 100644 --- a/controller/src/sre_registration.rs +++ b/controller/src/sre_registration.rs @@ -17,6 +17,7 @@ pub const PRIVATE_SECRET: &str = "sre-api-router-identity"; pub const AGENT_SECRET: &str = "sre-api-agent"; pub const OWNER: &str = "kars.azure.com/sre-registration-uid"; pub const EPOCH: &str = "kars.azure.com/sre-privacy-epoch"; +pub const CONTROL_VERSION: &str = "kars.azure.com/services-credential-version"; #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)] #[serde(rename_all = "camelCase")] diff --git a/docs/governed-services.md b/docs/governed-services.md new file mode 100644 index 000000000..81808dddc --- /dev/null +++ b/docs/governed-services.md @@ -0,0 +1,174 @@ +# Router governed services + +These APIs provide an in-process capability-request queue and bounded router +telemetry. They do **not** deliver assignments, run agents, create approvals, +grant capabilities, install resources, or provide a durable execution ledger. + +**Qualification gate:** the combined source includes the operator-authorized +[SRE identity/migration prerequisite](how-to/sre-authority.md). Its real-API +migration acceptance (#551) remains pending; merging its implementation locally +does not establish successful hosted qualification or make this candidate ready +for deployment. Changing mounts alone does not establish operator-only authority. + +## Identity and operator authentication + +Each router has a unique process-instance scope. On Kubernetes the controller +projects the current Sandbox name, source workspace, Sandbox UID, and runtime +namespace UID. Task-owned sandboxes additionally carry the task UID, +generation, and **effective authorization digest**, validated against current +task status and its sandbox reference. The envelope-only digest is not used as +authorization, and unsupported finite/shared launch budgets remain rejected. + +Operator endpoints require a separate bearer credential. The controller creates +`router-services-admin` in the authorized sandbox namespace and mounts its +`control-token` key **only in the inference-router**, at +`/etc/kars/services/control-token`. The legacy `router-admin-token` is also +available to some agent plugins and is deliberately not accepted here. + +Issuance and reuse require the shared live Secret GET/LIST/WATCH denial checks +for the legacy SRE principal. If an SRE registration exists, it must carry current +verified v2 authority: either Ready with its exact source/controller/namespace +identities, enforcing admission and no unsafe token aliases, or fully Retired +with current denial evidence. A Ready registration's privacy epoch is recorded +on newly issued control credentials. No registration is needed for ordinary +standalone installations where the legacy principal is actually denied. + +Existing unqualified **controller-owned** credentials rotate with Secret UID/RV +preconditions; matching names or partial ownership labels never authorize +adoption. The Deployment's credential-version annotation restarts cached-token +consumers without changing their image pins, selectors or unrelated annotations. +The controller does not report its credential transition complete while old +version Pods remain, including terminating Pods. Loss of the checked privacy +proof quarantines the owned credential and scales only a proven owned consumer +to zero. Restoring authority requires a new token, not reuse of the potentially +exposed cache. Foreign consumers remain untouched and block automatic recovery. +Initial qualification still pending is not permission to issue or reuse a +credential, but it is not itself evidence of privacy loss and does not stop the +rollout being qualified. Privacy qualification completes only after fresh +denial/v2/UID/epoch/template checks and termination of old-epoch or +old-control-version Pods. Workload availability is separate: the SRE readiness +endpoint itself depends on qualified authority, so it cannot be an input to +that qualification. Canonical SRE's early authorization failure +also performs the ownership-fenced, no-issuance privacy-loss quarantine. +These checks occur during reconciliation; they are not a claim of instantaneous +cluster-wide revocation or cancellation of already accepted upstream work. + +Standalone operators can provide `KARS_SERVICES_ADMIN_TOKEN` (32–256 nonblank +ASCII characters). Without a control credential, operator APIs return 503; +there is no fallback to the legacy admin token or localhost exemption. +`ROUTER_ADMIN_ALLOW_IPS`, when configured, also restricts authenticated operator +origins. Agent APIs accept only the same-router loopback client or an +authenticated allowed operator. They never accept a caller-selected Sandbox. + +## Request lifecycle + +| Endpoint | Caller | Purpose | +|---|---|---| +| `GET /v1/access-requests` | Agent | Discover the current scope and inspect own requests | +| `POST /v1/access-request` | Agent | Queue a capability request, never a grant | +| `POST /v1/access-requests/{id}/cancel` | Agent | Cancel its pending/approved wait | +| `GET /v1/access-requests/{id}/wait?scope_id=…&timeout_ms=…` | Agent | Wait for a decision or lifecycle transition | +| `GET /internal/access-requests` | Operator | Inspect scope and the `entries` queue | +| `POST /internal/access-requests/decision` | Operator | Record one immutable approval/denial | +| `POST /internal/access-requests/reset` | Operator | Start a fresh scope and cancel old waits | + +A request body contains `scope_id`, `kind`, `target`, and an optional `reason`. +Kinds are `egress`, `tool`, `skill`, `mcp`, `command`, `permission`, +`clarification`, and `tier`. Targets are bounded machine identifiers; egress +uses a bare DNS hostname plus a separate `port` (default 443). A tier request +uses `tier: 1..5` and an empty target. Reasons are explicitly submitted +untrusted text, limited to 512 bytes; they are not copied from proxied API +bodies or written to diagnostic logs. Never put credentials in a reason. + +The service generates `request_id`. Decisions require that exact ID and its +`scope_id`, with `verdict: approved|denied`. There is no decision-by-host +fallback. Deduplication includes port and tier, preserves the original payload +and expiry, and never inherits a decision for a different capability. + +Requests live for 15 minutes. The default queue is 64 entries and creation is +limited to 32 attempts/minute, including duplicates. Requests cannot evict +other pending requests when full. Stale scopes, expired requests, and repeated +terminal decisions are rejected. Cancelled requests retain any original +operator decision separately; cancelling a wait does not revoke policy. + +Reset requires the current `scope_id` as a compare-and-swap condition and can +include an `assignment_id` correlation label. It returns a fresh scope with +the same trusted identity, clears request/telemetry state, and cancels old +waiters. Reset returns conflict while an approved dispatch is already claimed. +An assignment label is not evidence that a runtime executed anything. +Reset neither changes policy nor resets existing token-budget counters. + +## Policy-gated egress waits + +Ordinary denied `/egress/fetch` requests still return immediately. A caller can +opt in by supplying `scope_id` and `wait_for_approval_ms` (at most 300,000). +At most 16 service waits can be active; dropping a wait, cancellation, reset, +timeout, or shutdown releases its slot. Operator control routes have separate +concurrency from agent waits and model traffic. + +An approved decision **does not open the network**. The fetch resumes only +after the ordinary egress policy also allows that exact operation. The +existing signed-allowlist/controller path remains responsible for enforcement; +private-address checks, normal egress filters, and redirect behavior still +apply. Denial, cancellation, expiry, and reset never bypass those filters. + +After the asynchronous policy check, the waiter revalidates the request. An +atomic dispatch claim then coordinates with cancellation and reset immediately +before sending. Cancellation acknowledged before that claim prevents dispatch; +after the claim it returns conflict rather than promising prevention. The claim +stays active through response handling and releases on completion or caller +cancellation. It does not assert that the upstream accepted the request. + +Transparent proxy traffic records bounded denied requests and policy +observations, but keeps its existing immediate-denial/retry behavior. It does +not acquire a new implicit human-approval wait. + +## Telemetry + +| Endpoint | Purpose | +|---|---| +| `GET /telemetry/cursor` | Current `scope_id` and sequence | +| `GET /telemetry/trace?since=N` | Scoped observations; send `X-Kars-Service-Scope` | +| `POST /telemetry/tool` | Complete a scoped, AGT-correlated harness report | +| `GET /telemetry/budget` | Existing per-sandbox counters/limits, not a task budget | + +The trace ring retains at most 1,024 events and pages at most 256 at a time. +Follow `has_more` using the returned cursor; retention loss is exposed through +`dropped_events`. Native/model correlation state is separately bounded. + +Model `round` events describe **forwarding attempts**, not autonomous agent +turns. Fallback attempts retain their actual configured provider identity. +Missing usage remains null; incomplete/cancelled streams and transport errors +are not successes. Metadata parsing is bounded, and `partial_observation` +identifies incomplete observations. Streams are forwarded unchanged and are +never replayed by telemetry after acceptance. + +HTTP acceptance is distinct from semantic completion: buffered Responses with +`failed` or `incomplete` status and OpenAI error frames followed by `[DONE]` do +not become completed generations. Their accepted status, original response +bytes and available usage remain intact, without retaining provider error text. + +Only model identifiers, tool names/IDs, status, timing, and reported usage are +retained. Prompts, assistant text, URLs, headers, arguments, and result bodies +are not retained in telemetry. A model-proposed tool has no successful outcome +until an outcome is observed. Harness reports are explicitly labelled +`harness-reported`, not authoritative router execution. A native completion +requires a unique `tool_call_id` registered by an allowed `/agt/evaluate` +request whose context contains the current `scope_id`. + +MCP `isError`, JSON-RPC failures, and missing results are distinguished from +HTTP transport success. Telemetry makes no receipt-completeness, durable +budget, shared broker, or billing claim. + +## Integration contract for later task delivery + +The existing planned reset path remains +`POST /internal/access-requests/reset`, but its caller must use the private +service-control credential and the scoped compare-and-swap body. Do not +restore the old shared-admin-token/unscoped-reset implementation. + +A future task controller must compare the returned task UID, generation, and +effective authorization with current live authority before acting on a +request or decision. It must revoke stale assignment-scoped grants through +the real policy controller before reusing a sandbox. Reset is not that +revocation, and this slice intentionally provides no autonomous grant worker. diff --git a/docs/security-audits/2026-09-08-governed-router-services.md b/docs/security-audits/2026-09-08-governed-router-services.md new file mode 100644 index 000000000..7609cd672 --- /dev/null +++ b/docs/security-audits/2026-09-08-governed-router-services.md @@ -0,0 +1,105 @@ +# Capability audit — Scoped router governed services + +Date: 2026-09-08 +Status: Implementation candidate; independent review and sign-offs pending. + +## Known blocking finding + +The supported legacy SRE installation grants its agent-held Kubernetes +credential cluster-wide Secret reads. It can therefore obtain the purportedly +operator-only service token through the API. Router-only mounts do not close +that path. The HIGH remains unresolved in this candidate; a separately reviewed +operator-controlled UID registration/migration and SRE credential boundary is +required before deployment or merge. Legacy grants cannot be retired merely by +trusting names, labels or ownership-looking annotations. + +## Scope + +Bounded access-request services, operator-only reset/inspection/decisions, +explicit policy-gated waits, and metadata-only inference/MCP/governance +observations. Narrow controller changes project qualified identity and a +separate router-only service-control credential. + +No task delivery, runtime worker, structural execution plan, GitHub write +service, managed MCP deployment, skill-package installation, memory-mount +suite, provider implementation, or agent transport is introduced. + +## Authority and lifecycle + +- Operator service controls do not trust localhost or the agent-visible legacy + admin token. Their separate credential is never mounted in agent containers. +- Scope is server-owned and includes UID-qualified Sandbox/namespace identity; + task attribution requires current task UID, generation, full effective + authorization, launch contract, and sandbox binding. +- Reset is compare-and-swap fenced, atomically coordinates request/telemetry + scope changes, and cancels old waits. Old IDs cannot become new grants. +- Requests, records, IDs, reasons, lifetimes, rates, waits, and telemetry + correlation/storage are bounded. +- Decisions do not modify permissions. Egress waits require both approval and + the existing normal policy check; resets do not modify budgets or policies. +- Existing finite/shared/monetary launch-budget restrictions remain intact. + +## Observation integrity and privacy + +- No prompts, tool arguments/results, request URLs/headers, provider error + bodies, or credentials are retained in the new telemetry. +- Model proposals have no fabricated success. Harness outcomes are labelled + as reported; MCP `isError` and dispatch failures remain failures at HTTP 200. +- Missing usage, incomplete streams, cancellations, and retention loss are + explicit. Accepted traffic is never replayed by the observer. +- Existing provider provenance, typed failure classification and routing + recovery are unchanged except for attaching optional observations. +- This infrastructure is not a durable execution/receipt ledger or shared + budget broker. Existing legacy operator APIs are not redefined by this slice. + +## Verification + +Local qualification completed before independent review: + +- 1,067 router unit tests passed, including a real TCP forward-proxy denial + that queues a scoped request without introducing an implicit wait. +- 24 new HTTP/lifecycle/telemetry integration tests passed. These include a + live loopback HTTP server, strict credential separation, scope fencing, + rate/body limits, reset/cancel/timeout, policy-gated wakeup, provider identity, + fragmented streams, missing usage, MCP `isError`, and transport failures. +- 73 existing qualified provider/governance/egress/guardrail integration tests + passed unchanged apart from construction of new optional observer state. +- Four controller projection tests passed, including current full task + authorization/UID checks and continued rejection of unsupported launch budgets. +- Strict affected-crate Clippy and formatting checks passed. +- Static A2A-isolation, null-provider, and tracked-source copyright checks + passed; new candidate files are also checked for headers and module limits. + +Two MEDIUM review findings have been repaired and source-reviewed separately: +the cancellation/expiry race across an asynchronous policy check, and semantic +model errors reported as completed generations. Dispatch now has a mutex- +coordinated claim boundary; cancellation/reset cannot acknowledge prevention +after that claim. HTTP acceptance remains distinct from semantic failure or +incompleteness, with no response rewriting or replay. + +Repair qualification passed 1,071 router unit tests and 29 governed-service +integration tests, including deterministic cancellation/reset/expiry/shutdown +interleavings and accepted failed/incomplete/error-plus-DONE responses. Strict +router all-target Clippy passed; the focused 12 telemetry integrations passed +again after a type-safe initializer adjustment. These results do not resolve +the separate HIGH SRE credential-access issue. + +The existing disposable Kind harness now includes the actual deployed router: +it checks that the control Secret is mounted only in the router, rejects the +agent-visible token even through loopback port-forwarding, compares returned +Sandbox/namespace UIDs with live objects, exercises private decisions/reset, +and rejects stale scopes after reset. Private fixture credentials travel via +curl's stdin rather than command arguments or logs. The port-forward process +and named scratch files are cleaned up on exit. Shell syntax passed locally; +hosted execution of this added case remains pending. + +No active customer cluster, Azure mutation, deployment, or image publication +was performed during preparation. Dependencies were not changed or installed. The root Cargo target +was reused offline/locked with incremental compilation disabled and a guarded +disk reserve. Committed-diff publication gates and independent audit sign-offs +remain the parent publication process's responsibility. + +## Verdict + +Pending. No reviewer approval or signature is claimed. Genuine author and +independent reviewer sign-offs are required by the publication process. diff --git a/inference-router/src/access_request.rs b/inference-router/src/access_request.rs new file mode 100644 index 000000000..23b28582c --- /dev/null +++ b/inference-router/src/access_request.rs @@ -0,0 +1,453 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Scope-fenced capability requests. Decisions are observations, never grants. + +use serde::{Deserialize, Serialize}; +use std::{ + collections::VecDeque, + sync::Mutex, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; + +pub mod scope; +pub use scope::{Identity, Scope}; + +#[cfg(test)] +#[path = "access_request/dispatch_tests.rs"] +mod dispatch_tests; + +const CAPACITY: usize = 64; +const REQUESTS_PER_MINUTE: u32 = 32; +const TTL: Duration = Duration::from_secs(15 * 60); + +#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum Status { + Pending, + Approved, + Denied, + Cancelled, + Expired, + DispatchClaimed, +} + +#[derive(Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Request { + pub scope_id: String, + pub kind: String, + #[serde(default)] + pub target: String, + #[serde(default)] + pub reason: String, + pub tier: Option, + pub port: Option, +} + +#[derive(Clone, Serialize)] +pub struct Entry { + pub request_id: String, + pub scope_id: String, + pub kind: String, + pub target: String, + pub reason: String, + pub tier: Option, + pub port: Option, + pub status: Status, + pub decision: Option, + pub count: u32, + pub first_seen_unix: u64, + pub expires_at_unix: u64, + pub dispatch_active: bool, + #[serde(skip)] + expires: Instant, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Error { + Invalid, + StaleScope, + Full, + RateLimited, + Missing, + Expired, + WaitTimeout, + Terminal, + Unavailable, +} + +struct Inner { + scope: Scope, + generation: u64, + sequence: u64, + requests: VecDeque, + window: Instant, + rate: u32, + cancellation: CancellationToken, + active_dispatches: usize, +} + +pub struct AccessRequestBuffer { + inner: Mutex, + changed: watch::Sender, + capacity: usize, + rate_limit: u32, + ttl: Duration, +} + +/// The linearization boundary for an outgoing operation, not evidence that an +/// upstream accepted it. Cancel/reset cannot acknowledge prevention while held. +#[must_use] +pub struct DispatchClaim<'a> { + buffer: &'a AccessRequestBuffer, + request_id: Option, +} + +impl Drop for DispatchClaim<'_> { + fn drop(&mut self) { + if let Ok(mut inner) = self.buffer.inner.lock() { + inner.active_dispatches = inner.active_dispatches.saturating_sub(1); + if let Some(id) = &self.request_id + && let Some(entry) = inner + .requests + .iter_mut() + .find(|entry| &entry.request_id == id) + { + entry.dispatch_active = false; + } + } + self.buffer.notify(); + } +} + +impl AccessRequestBuffer { + pub fn new(identity: Identity) -> Self { + Self::with_limits(identity, CAPACITY, REQUESTS_PER_MINUTE, TTL) + } + + pub fn with_limits( + identity: Identity, + capacity: usize, + rate_limit: u32, + ttl: Duration, + ) -> Self { + let scope = Scope { + id: format!("{:032x}:0", rand::random::()), + identity, + assignment_id: None, + }; + Self { + inner: Mutex::new(Inner { + scope, + generation: 0, + sequence: 0, + requests: VecDeque::new(), + window: Instant::now(), + rate: 0, + cancellation: CancellationToken::new(), + active_dispatches: 0, + }), + changed: watch::channel(0).0, + capacity: capacity.clamp(1, 256), + rate_limit: rate_limit.clamp(1, 128), + ttl: ttl.min(Duration::from_secs(3600)), + } + } + + fn notify(&self) { + self.changed + .send_modify(|value| *value = value.saturating_add(1)); + } + pub fn subscribe(&self) -> watch::Receiver { + self.changed.subscribe() + } + pub fn scope(&self) -> Result { + Ok(self + .inner + .lock() + .map_err(|_| Error::Unavailable)? + .scope + .clone()) + } + + fn check(inner: &Inner, scope: &str) -> Result<(), Error> { + if inner.scope.id == scope { + Ok(()) + } else { + Err(Error::StaleScope) + } + } + fn expire(inner: &mut Inner) { + for entry in &mut inner.requests { + if entry.expires <= Instant::now() + && matches!(entry.status, Status::Pending | Status::Approved) + { + entry.status = Status::Expired; + } + } + } + + pub fn record(&self, mut request: Request) -> Result<(Entry, bool), Error> { + validate(&mut request)?; + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, &request.scope_id)?; + if inner.window.elapsed() >= Duration::from_secs(60) { + inner.window = Instant::now(); + inner.rate = 0; + } + if inner.rate >= self.rate_limit { + return Err(Error::RateLimited); + } + inner.rate += 1; + Self::expire(&mut inner); + if let Some(entry) = inner.requests.iter_mut().find(|entry| { + entry.kind == request.kind + && entry.target == request.target + && entry.port == request.port + && entry.tier == request.tier + && entry.expires > Instant::now() + }) { + // An agent cannot rewrite the reviewed payload or extend its lifetime. + entry.count = entry.count.saturating_add(1); + return Ok((entry.clone(), false)); + } + if inner.requests.len() >= self.capacity { + if let Some(index) = inner + .requests + .iter() + .position(|entry| entry.status != Status::Pending && !entry.dispatch_active) + { + inner.requests.remove(index); + } else { + return Err(Error::Full); + } + } + inner.sequence = inner.sequence.checked_add(1).ok_or(Error::Unavailable)?; + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + let entry = Entry { + request_id: format!("{}:{}", inner.scope.id, inner.sequence), + scope_id: inner.scope.id.clone(), + kind: request.kind, + target: request.target, + reason: request.reason, + tier: request.tier, + port: request.port, + status: Status::Pending, + decision: None, + count: 1, + first_seen_unix: now, + expires_at_unix: now.saturating_add(self.ttl.as_secs()), + dispatch_active: false, + expires: Instant::now() + self.ttl, + }; + inner.requests.push_back(entry.clone()); + drop(inner); + self.notify(); + Ok((entry, true)) + } + + pub fn snapshot(&self, scope: &str) -> Result, Error> { + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, scope)?; + Self::expire(&mut inner); + Ok(inner.requests.iter().cloned().collect()) + } + + pub fn transition(&self, scope: &str, id: &str, status: Status) -> Result { + if !matches!( + status, + Status::Approved | Status::Denied | Status::Cancelled + ) { + return Err(Error::Invalid); + } + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, scope)?; + Self::expire(&mut inner); + let entry = inner + .requests + .iter_mut() + .find(|entry| entry.request_id == id) + .ok_or(Error::Missing)?; + if entry.status == Status::DispatchClaimed { + return Err(Error::Terminal); + } + if entry.expires <= Instant::now() { + return Err(Error::Expired); + } + if entry.status != Status::Pending + && !(status == Status::Cancelled && entry.status == Status::Approved) + { + return Err(Error::Terminal); + } + if matches!(status, Status::Approved | Status::Denied) { + entry.decision = Some(status); + } + entry.status = status; + let entry = entry.clone(); + drop(inner); + self.notify(); + Ok(entry) + } + + pub fn reset( + &self, + expected_scope: &str, + assignment_id: Option, + ) -> Result<(Scope, usize), Error> { + if assignment_id + .as_deref() + .is_some_and(|id| !scope::identifier(id, 253)) + { + return Err(Error::Invalid); + } + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, expected_scope)?; + if inner.active_dispatches > 0 { + return Err(Error::Terminal); + } + inner.generation = inner.generation.checked_add(1).ok_or(Error::Unavailable)?; + let instance = inner.scope.id.split(':').next().ok_or(Error::Unavailable)?; + inner.scope.id = format!("{instance}:{}", inner.generation); + inner.scope.assignment_id = assignment_id; + let cleared = inner.requests.len(); + inner.requests.clear(); + inner.cancellation.cancel(); + inner.cancellation = CancellationToken::new(); + let scope = inner.scope.clone(); + drop(inner); + self.notify(); + Ok((scope, cleared)) + } + + pub fn cancellation(&self, scope: &str) -> Result { + let inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, scope)?; + Ok(inner.cancellation.clone()) + } + + pub fn entry(&self, scope: &str, id: &str) -> Result { + self.snapshot(scope)? + .into_iter() + .find(|entry| entry.request_id == id) + .ok_or(Error::Missing) + } + + fn dispatch_ready( + inner: &mut Inner, + scope: &str, + id: Option<&str>, + shutdown: &CancellationToken, + ) -> Result<(), Error> { + Self::check(inner, scope)?; + if shutdown.is_cancelled() { + return Err(Error::Unavailable); + } + Self::expire(inner); + if let Some(id) = id { + let entry = inner + .requests + .iter() + .find(|entry| entry.request_id == id) + .ok_or(Error::Missing)?; + if entry.status == Status::Expired { + return Err(Error::Expired); + } + if entry.status != Status::Approved { + return Err(Error::Terminal); + } + } + Ok(()) + } + + /// Recheck after an asynchronous policy lookup. This is deliberately not a + /// dispatch permit: callers must still claim at the actual send boundary. + pub fn validate_dispatch( + &self, + scope: &str, + id: &str, + shutdown: &CancellationToken, + ) -> Result<(), Error> { + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::dispatch_ready(&mut inner, scope, Some(id), shutdown) + } + + pub fn claim_dispatch( + &self, + scope: &str, + id: Option<&str>, + shutdown: &CancellationToken, + ) -> Result, Error> { + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::dispatch_ready(&mut inner, scope, id, shutdown)?; + let active = inner.active_dispatches.checked_add(1).ok_or(Error::Full)?; + if let Some(id) = id { + let entry = inner + .requests + .iter_mut() + .find(|entry| entry.request_id == id) + .ok_or(Error::Missing)?; + entry.status = Status::DispatchClaimed; + entry.dispatch_active = true; + } + inner.active_dispatches = active; + drop(inner); + self.notify(); + Ok(DispatchClaim { + buffer: self, + request_id: id.map(str::to_string), + }) + } +} + +fn validate(request: &mut Request) -> Result<(), Error> { + request.kind = request.kind.trim().to_ascii_lowercase(); + request.target = request.target.trim().to_string(); + if ![ + "egress", + "tool", + "skill", + "mcp", + "command", + "permission", + "clarification", + "tier", + ] + .contains(&request.kind.as_str()) + || request.reason.len() > 512 + || request.reason.chars().any(|c| c.is_control() && c != '\n') + { + return Err(Error::Invalid); + } + if request.kind == "tier" { + if !request.tier.is_some_and(|tier| (1..=5).contains(&tier)) + || !request.target.is_empty() + || request.port.is_some() + { + return Err(Error::Invalid); + } + } else { + if !scope::identifier(&request.target, 253) || request.tier.is_some() { + return Err(Error::Invalid); + } + if request.kind == "egress" { + if request.target.contains(['/', ':']) { + return Err(Error::Invalid); + } + request.target = + crate::egress_blocked::normalize_host(&request.target).ok_or(Error::Invalid)?; + request.port = Some(request.port.unwrap_or(443)); + if request.port == Some(0) { + return Err(Error::Invalid); + } + } else if request.port.is_some() { + return Err(Error::Invalid); + } + } + Ok(()) +} diff --git a/inference-router/src/access_request/dispatch_tests.rs b/inference-router/src/access_request/dispatch_tests.rs new file mode 100644 index 000000000..bb73fff89 --- /dev/null +++ b/inference-router/src/access_request/dispatch_tests.rs @@ -0,0 +1,225 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::governed_services::GovernedServices; +use std::sync::Arc; +use tokio::sync::{RwLock, oneshot}; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + +fn approved(port: u16) -> (Arc, String, String) { + let services = Arc::new(GovernedServices::new(Identity::standalone("test"), None)); + let scope = services.requests.scope().unwrap().id; + let (entry, _) = services + .requests + .record(Request { + scope_id: scope.clone(), + kind: "egress".into(), + target: "dispatch.example".into(), + reason: String::new(), + tier: None, + port: Some(port), + }) + .unwrap(); + services + .requests + .transition(&scope, &entry.request_id, Status::Approved) + .unwrap(); + (services, scope, entry.request_id) +} + +#[tokio::test] +async fn cancellation_expiry_reset_and_shutdown_during_policy_await_prevent_post_dispatch() { + for action in ["cancel", "expire", "reset", "shutdown"] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let address = reqwest::Url::parse(&server.uri()) + .unwrap() + .socket_addrs(|| None) + .unwrap()[0]; + let target = format!("http://dispatch.example:{}/side-effect", address.port()); + let client = reqwest::Client::builder() + .no_proxy() + .resolve("dispatch.example", address) + .build() + .unwrap(); + let (services, scope, id) = approved(address.port()); + let policy = Arc::new(RwLock::new(false)); + let mut held = policy.write().await; + let (entered, observing) = oneshot::channel(); + let entered = Arc::new(Mutex::new(Some(entered))); + let worker = services.clone(); + let worker_scope = scope.clone(); + let worker_id = id.clone(); + let policy_for_worker = policy.clone(); + let task = tokio::spawn(async move { + worker + .wait_for_egress_check( + &worker_scope, + &worker_id, + &target, + "test", + Duration::from_secs(5), + || { + let policy = policy_for_worker.clone(); + let entered = entered.clone(); + async move { + if let Some(sender) = entered.lock().unwrap().take() { + let _ = sender.send(()); + } + if *policy.read().await { + Ok(()) + } else { + Err("blocked".to_string()) + } + } + }, + ) + .await?; + let _claim = worker.claim_egress_dispatch(&worker_scope, Some(&worker_id))?; + client + .post(&target) + .body("side-effect") + .send() + .await + .map_err(|_| Error::Unavailable)?; + Ok::<(), Error>(()) + }); + observing.await.unwrap(); + let expected = match action { + "cancel" => { + let cancelled = services + .requests + .transition(&scope, &id, Status::Cancelled) + .unwrap(); + assert_eq!(cancelled.status, Status::Cancelled); + Error::Terminal + } + "expire" => { + // Fixture mutation advances expiry only after the approved + // waiter is demonstrably blocked in its asynchronous check. + services + .requests + .inner + .lock() + .unwrap() + .requests + .iter_mut() + .find(|entry| entry.request_id == id) + .unwrap() + .expires = Instant::now(); + Error::Expired + } + "reset" => { + services + .reset(&scope, Some("next-assignment".into())) + .unwrap(); + Error::StaleScope + } + _ => { + services.shutdown.cancel(); + Error::Unavailable + } + }; + *held = true; + drop(held); + assert_eq!(task.await.unwrap(), Err(expected), "{action}"); + assert!( + server.received_requests().await.unwrap().is_empty(), + "{action}" + ); + } +} + +#[test] +fn acknowledged_cancel_or_reset_before_claim_wins_and_claims_are_single_use() { + let (services, scope, id) = approved(443); + services + .requests + .validate_dispatch(&scope, &id, &services.shutdown) + .unwrap(); + services + .requests + .transition(&scope, &id, Status::Cancelled) + .unwrap(); + assert!(matches!( + services.claim_egress_dispatch(&scope, Some(&id)), + Err(Error::Terminal) + )); + + let (services, scope, id) = approved(443); + services.reset(&scope, Some("next".into())).unwrap(); + assert!(matches!( + services.claim_egress_dispatch(&scope, Some(&id)), + Err(Error::StaleScope) + )); + + let (services, scope, id) = approved(443); + let claim = services.claim_egress_dispatch(&scope, Some(&id)).unwrap(); + assert_eq!( + services.requests.entry(&scope, &id).unwrap().status, + Status::DispatchClaimed + ); + assert!( + services + .requests + .entry(&scope, &id) + .unwrap() + .dispatch_active + ); + assert!(matches!( + services.requests.transition(&scope, &id, Status::Cancelled), + Err(Error::Terminal) + )); + assert!(matches!( + services.claim_egress_dispatch(&scope, Some(&id)), + Err(Error::Terminal) + )); + assert!(matches!(services.reset(&scope, None), Err(Error::Terminal))); + drop(claim); + assert!( + !services + .requests + .entry(&scope, &id) + .unwrap() + .dispatch_active + ); + assert!(matches!( + services.requests.transition(&scope, &id, Status::Cancelled), + Err(Error::Terminal) + )); + services.reset(&scope, None).unwrap(); +} + +#[test] +fn expiry_and_shutdown_between_readiness_and_claim_still_prevent_dispatch() { + for expired in [true, false] { + let (services, scope, id) = approved(443); + services + .requests + .validate_dispatch(&scope, &id, &services.shutdown) + .unwrap(); + if expired { + services.requests.inner.lock().unwrap().requests[0].expires = Instant::now(); + } else { + services.shutdown.cancel(); + } + assert!(matches!( + services.claim_egress_dispatch(&scope, Some(&id)), + Err(Error::Expired | Error::Unavailable) + )); + } +} + +#[test] +fn scope_only_dispatch_also_prevents_successful_reset_until_released() { + let (services, scope, _) = approved(443); + let claim = services.claim_egress_dispatch(&scope, None).unwrap(); + assert!(matches!(services.reset(&scope, None), Err(Error::Terminal))); + drop(claim); + services.reset(&scope, None).unwrap(); +} diff --git a/inference-router/src/access_request/scope.rs b/inference-router/src/access_request/scope.rs new file mode 100644 index 000000000..dae8b1a32 --- /dev/null +++ b/inference-router/src/access_request/scope.rs @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct ResourceIdentity { + pub namespace: String, + pub name: String, + pub uid: String, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct Identity { + pub sandbox: ResourceIdentity, + pub namespace_uid: String, + #[serde(default)] + pub task: Option, + #[serde(default)] + pub task_authorization: Option, + #[serde(default)] + pub task_generation: Option, + pub managed: bool, +} + +#[derive(Clone, Debug, Serialize)] +pub struct Scope { + pub id: String, + pub identity: Identity, + pub assignment_id: Option, +} + +pub fn identifier(value: &str, max: usize) -> bool { + !value.is_empty() + && value.len() <= max + && value + .bytes() + .all(|c| c.is_ascii_alphanumeric() || b"-._:/".contains(&c)) +} + +impl Identity { + pub fn standalone(name: &str) -> Self { + Self { + sandbox: ResourceIdentity { + namespace: "standalone".into(), + name: name.into(), + uid: "process-local".into(), + }, + namespace_uid: "process-local".into(), + task: None, + task_authorization: None, + task_generation: None, + managed: false, + } + } + pub fn valid(&self, sandbox_name: &str) -> bool { + let valid = |identity: &ResourceIdentity| { + identifier(&identity.namespace, 253) + && identifier(&identity.name, 253) + && identifier(&identity.uid, 128) + }; + self.sandbox.name == sandbox_name + && valid(&self.sandbox) + && identifier(&self.namespace_uid, 128) + && self + .task + .as_ref() + .is_none_or(|task| valid(task) && task.namespace == self.sandbox.namespace) + && match &self.task { + Some(_) => { + self.task_generation + .is_some_and(|generation| generation > 0) + && self + .task_authorization + .as_deref() + .and_then(|digest| digest.strip_prefix("sha256:")) + .is_some_and(|digest| { + digest.len() == 64 + && digest.bytes().all(|byte| byte.is_ascii_hexdigit()) + }) + } + None => self.task_authorization.is_none() && self.task_generation.is_none(), + } + } +} diff --git a/inference-router/src/egress_blocked.rs b/inference-router/src/egress_blocked.rs index 72646c1f1..7ca40d4f6 100644 --- a/inference-router/src/egress_blocked.rs +++ b/inference-router/src/egress_blocked.rs @@ -9,7 +9,7 @@ //! paths, headers, query strings, or payload data are ever stored. use std::collections::{HashMap, VecDeque}; -use std::sync::Mutex; +use std::sync::{Arc, Mutex, Weak}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; /// Default capacity when the operator does not override. @@ -51,6 +51,7 @@ pub enum RecordOutcome { /// Bounded ring buffer of blocked-attempt records. pub struct BlockedBuffer { + services: Mutex>>, inner: Mutex, capacity: usize, rate_limit_window: Duration, @@ -71,6 +72,7 @@ struct RateState { impl BlockedBuffer { pub fn new(capacity: usize, rate_limit_window: Duration, rate_limit_per_source: u32) -> Self { Self { + services: Mutex::new(None), inner: Mutex::new(Inner { by_key: HashMap::new(), order: VecDeque::new(), @@ -96,7 +98,42 @@ impl BlockedBuffer { /// possible results. Hostname-only — never store paths, headers, or /// payload data. pub fn record(&self, source_sandbox: &str, host: &str, port: u16) -> RecordOutcome { - self.record_at(Instant::now(), unix_now(), source_sandbox, host, port) + let result = self.record_at(Instant::now(), unix_now(), source_sandbox, host, port); + if matches!( + result, + RecordOutcome::Recorded | RecordOutcome::Deduplicated + ) && let Some(services) = self + .services + .lock() + .ok() + .and_then(|services| services.as_ref().and_then(Weak::upgrade)) + && let Ok(scope) = services.requests.scope() + && scope.identity.sandbox.name == source_sandbox + { + let suffix = format!(":{port}"); + let host = host.strip_suffix(&suffix).unwrap_or(host); + let _ = services.record_egress_in_scope(&scope.id, host, port); + services.telemetry.record_policy(&scope.id, "egress", false); + } + result + } + + pub fn bind_services(&self, services: &Arc) { + if let Ok(mut target) = self.services.lock() { + *target = Some(Arc::downgrade(services)); + } + } + pub fn observe_allowed(&self, source_sandbox: &str) { + if let Some(services) = self + .services + .lock() + .ok() + .and_then(|services| services.as_ref().and_then(Weak::upgrade)) + && let Ok(scope) = services.requests.scope() + && scope.identity.sandbox.name == source_sandbox + { + services.telemetry.record_policy(&scope.id, "egress", true); + } } /// Testable record entry point — accepts an injected clock pair. @@ -269,7 +306,7 @@ fn unix_now() -> u64 { } /// Lowercase ASCII, strip a single trailing dot, reject empty / IP literal. -fn normalize_host(host: &str) -> Option { +pub(crate) fn normalize_host(host: &str) -> Option { let trimmed = host.trim(); if trimmed.is_empty() { return None; diff --git a/inference-router/src/failover.rs b/inference-router/src/failover.rs index 1bf916150..ecb522045 100644 --- a/inference-router/src/failover.rs +++ b/inference-router/src/failover.rs @@ -466,6 +466,7 @@ mod tests { fn upstream(dep: &str) -> UpstreamConfig { UpstreamConfig { + telemetry: None, endpoint: "https://example.openai.azure.com".into(), deployment: dep.to_string(), sandbox_name: "sbx".into(), diff --git a/inference-router/src/forward_proxy.rs b/inference-router/src/forward_proxy.rs index f5acae588..7e0541dd1 100644 --- a/inference-router/src/forward_proxy.rs +++ b/inference-router/src/forward_proxy.rs @@ -21,6 +21,10 @@ use tokio_util::sync::CancellationToken; use crate::blocklist::Blocklist; use crate::egress_blocked::BlockedBuffer; +#[cfg(test)] +#[path = "forward_proxy/governed_tests.rs"] +mod governed_tests; + /// Maximum concurrent tunnel connections (prevents resource exhaustion). const MAX_CONCURRENT_TUNNELS: usize = 256; @@ -324,6 +328,8 @@ async fn handle_connect( return Ok(()); } + blocked_egress.observe_allowed(sandbox); + // Resolve DNS immediately after policy check and validate against private IPs let resolved = match resolve_and_validate(&domain, port, sandbox, blocked_egress).await { Ok(addr) => addr, @@ -416,6 +422,8 @@ async fn handle_http( return Ok(()); } + blocked_egress.observe_allowed(sandbox); + // Resolve + validate (prevents DNS rebinding to private IPs) let (host, port) = parse_host_port(&domain, 80); let resolved = match resolve_and_validate(&host, port, sandbox, blocked_egress).await { @@ -495,6 +503,8 @@ async fn handle_tls_redirect( return Ok(()); } + blocked_egress.observe_allowed(sandbox); + // Resolve + validate (prevents DNS rebinding to private IPs) let resolved = match resolve_and_validate(&domain, 443, sandbox, blocked_egress).await { Ok(addr) => addr, diff --git a/inference-router/src/forward_proxy/governed_tests.rs b/inference-router/src/forward_proxy/governed_tests.rs new file mode 100644 index 000000000..3e6fff692 --- /dev/null +++ b/inference-router/src/forward_proxy/governed_tests.rs @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::{access_request::Identity, governed_services::GovernedServices}; + +#[tokio::test] +async fn denied_connect_records_a_scoped_request_without_an_implicit_wait() { + let services = Arc::new(GovernedServices::new(Identity::standalone("test"), None)); + let scope = services.requests.scope().unwrap().id; + let blocked = Arc::new(BlockedBuffer::with_defaults()); + blocked.bind_services(&services); + let blocklist = Blocklist::new(None).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + handle_connection(socket, &blocklist, "test", &blocked) + .await + .unwrap(); + }); + let mut client = TcpStream::connect(address).await.unwrap(); + client + .write_all(b"CONNECT blocked.example:443 HTTP/1.1\r\nHost: blocked.example\r\n\r\n") + .await + .unwrap(); + let mut response = Vec::new(); + tokio::time::timeout( + std::time::Duration::from_secs(1), + client.read_to_end(&mut response), + ) + .await + .unwrap() + .unwrap(); + server.await.unwrap(); + assert!(String::from_utf8_lossy(&response).starts_with("HTTP/1.1 403")); + let entries = services.requests.snapshot(&scope).unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].target, "blocked.example"); + assert_eq!(entries[0].port, Some(443)); + assert_eq!(services.wait_slots(), 16); +} diff --git a/inference-router/src/governed_services.rs b/inference-router/src/governed_services.rs new file mode 100644 index 000000000..a723b0056 --- /dev/null +++ b/inference-router/src/governed_services.rs @@ -0,0 +1,259 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::{ + access_request::{AccessRequestBuffer, Error, Identity, Request, Scope, Status}, + blocklist::Blocklist, + task_telemetry::TaskTelemetry, +}; +use std::{net::IpAddr, sync::Arc, time::Duration}; +use tokio::sync::Semaphore; +use tokio_util::sync::CancellationToken; + +pub const CONTROL_TOKEN_PATH: &str = "/etc/kars/services/control-token"; +pub const MAX_WAIT: Duration = Duration::from_secs(300); +const WAIT_CAPACITY: usize = 16; + +pub struct GovernedServices { + pub requests: AccessRequestBuffer, + pub telemetry: Arc, + control_token: Option, + pub allow_ips: Option>, + pub identity_valid: bool, + pub shutdown: CancellationToken, + waits: Arc, + reset_gate: std::sync::Mutex<()>, +} + +impl Default for GovernedServices { + fn default() -> Self { + Self::new(Identity::standalone("unknown"), None) + } +} + +impl GovernedServices { + pub fn new(identity: Identity, control_token: Option) -> Self { + let control_token = control_token.filter(|token| { + (32..=256).contains(&token.len()) + && token.is_ascii() + && !token.chars().any(|c| c.is_whitespace() || c.is_control()) + }); + let requests = AccessRequestBuffer::new(identity); + let scope = requests + .scope() + .expect("fresh service state is not poisoned"); + Self { + requests, + telemetry: Arc::new(TaskTelemetry::new(scope.id)), + control_token, + allow_ips: None, + identity_valid: true, + shutdown: CancellationToken::new(), + waits: Arc::new(Semaphore::new(WAIT_CAPACITY)), + reset_gate: std::sync::Mutex::new(()), + } + } + pub fn from_env(sandbox: &str) -> Self { + let identity = std::env::var("KARS_SERVICE_IDENTITY_JSON") + .ok() + .map(|raw| serde_json::from_str::(&raw).ok()); + let valid = identity.as_ref().is_none_or(|identity| { + identity + .as_ref() + .is_some_and(|identity| identity.valid(sandbox)) + }); + let identity = identity + .flatten() + .unwrap_or_else(|| Identity::standalone(sandbox)); + let token = std::env::var("KARS_SERVICES_ADMIN_TOKEN") + .ok() + .or_else(|| std::fs::read_to_string(CONTROL_TOKEN_PATH).ok()) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + let mut services = Self::new(identity, token); + services.identity_valid = valid; + services.allow_ips = std::env::var("ROUTER_ADMIN_ALLOW_IPS") + .ok() + .filter(|value| !value.trim().is_empty()) + .map(|value| { + value + .split(',') + .map(str::trim) + .map(str::parse::) + .collect::, _>>() + .unwrap_or_default() + }); + services + } + pub fn control_authorized(&self, provided: Option<&str>) -> bool { + self.identity_valid + && self + .control_token + .as_deref() + .zip(provided) + .is_some_and(|(expected, provided)| { + crate::handoff::constant_time_eq(expected.as_bytes(), provided.as_bytes()) + }) + } + pub fn control_configured(&self) -> bool { + self.identity_valid && self.control_token.is_some() + } + pub fn reset( + &self, + expected: &str, + assignment: Option, + ) -> Result<(Scope, usize), Error> { + let _guard = self.reset_gate.lock().map_err(|_| Error::Unavailable)?; + let result = self.requests.reset(expected, assignment)?; + self.telemetry.reset(result.0.id.clone()); + Ok(result) + } + pub fn record_egress_in_scope( + &self, + scope: &str, + host: &str, + port: u16, + ) -> Result { + self.requests + .record(Request { + scope_id: scope.into(), + kind: "egress".into(), + target: host.into(), + reason: "Egress policy blocked this destination".into(), + tier: None, + port: Some(port), + }) + .map(|(entry, _)| entry) + } + pub fn wait_slots(&self) -> usize { + self.waits.available_permits() + } + + pub async fn wait_for_decision( + &self, + scope: &str, + id: &str, + timeout: Duration, + ) -> Result { + let _permit = self + .waits + .clone() + .try_acquire_owned() + .map_err(|_| Error::Full)?; + let cancellation = self.requests.cancellation(scope)?; + let mut changed = self.requests.subscribe(); + let future = async { + loop { + let entry = self.requests.entry(scope, id)?; + if entry.status != Status::Pending { + return Ok(entry); + } + tokio::select! { + _ = cancellation.cancelled() => return Err(Error::StaleScope), + _ = self.shutdown.cancelled() => return Err(Error::Unavailable), + _ = changed.changed() => {}, + _ = tokio::time::sleep(Duration::from_millis(100)) => {}, + } + } + }; + tokio::time::timeout(timeout.min(MAX_WAIT), future) + .await + .map_err(|_| Error::WaitTimeout)? + } + + pub async fn wait_for_egress( + &self, + blocklist: &Blocklist, + scope: &str, + id: &str, + target: &str, + sandbox: &str, + timeout: Duration, + ) -> Result<(), Error> { + self.wait_for_egress_check(scope, id, target, sandbox, timeout, || { + blocklist.check_egress(target, sandbox) + }) + .await + } + + pub(crate) async fn wait_for_egress_check( + &self, + scope: &str, + id: &str, + target: &str, + sandbox: &str, + timeout: Duration, + mut check_policy: F, + ) -> Result<(), Error> + where + F: FnMut() -> Fut, + Fut: std::future::Future>, + { + let url = reqwest::Url::parse(target).map_err(|_| Error::Invalid)?; + let host = url.host_str().ok_or(Error::Invalid)?; + let port = url.port_or_known_default().ok_or(Error::Invalid)?; + let owner = self.requests.scope()?; + if owner.id != scope || owner.identity.sandbox.name != sandbox { + return Err(Error::StaleScope); + } + let _permit = self + .waits + .clone() + .try_acquire_owned() + .map_err(|_| Error::Full)?; + let cancellation = self.requests.cancellation(scope)?; + let mut changed = self.requests.subscribe(); + let future = async { + loop { + let entry = self.requests.entry(scope, id)?; + if entry.kind != "egress" + || !entry.target.eq_ignore_ascii_case(host) + || entry.port != Some(port) + { + return Err(Error::Invalid); + } + match entry.status { + Status::Denied | Status::Cancelled | Status::DispatchClaimed => { + return Err(Error::Terminal); + } + Status::Expired => return Err(Error::Expired), + Status::Approved => { + // The decision API never changes enforcement. Only the + // normal signed-policy path can make this check pass. + let allowed = tokio::select! { + result = check_policy() => result.is_ok(), + _ = cancellation.cancelled() => return Err(Error::StaleScope), + _ = self.shutdown.cancelled() => return Err(Error::Unavailable), + _ = changed.changed() => continue, + }; + if allowed { + self.requests.validate_dispatch(scope, id, &self.shutdown)?; + return Ok(()); + } + } + Status::Pending => {} + } + tokio::select! { + _ = cancellation.cancelled() => return Err(Error::StaleScope), + _ = self.shutdown.cancelled() => return Err(Error::Unavailable), + _ = changed.changed() => {}, + _ = tokio::time::sleep(Duration::from_millis(100)) => {}, + } + } + }; + tokio::time::timeout(timeout.min(MAX_WAIT), future) + .await + .map_err(|_| Error::WaitTimeout)? + } + + /// Claim immediately before creating/polling the outgoing send future. + /// This shares the request mutex with cancellation and scope reset. + pub fn claim_egress_dispatch( + &self, + scope: &str, + request_id: Option<&str>, + ) -> Result, Error> { + self.requests + .claim_dispatch(scope, request_id, &self.shutdown) + } +} diff --git a/inference-router/src/lib.rs b/inference-router/src/lib.rs index 4b5e95ee8..01a766fc8 100644 --- a/inference-router/src/lib.rs +++ b/inference-router/src/lib.rs @@ -15,6 +15,7 @@ pub mod a2a; pub mod a2a_mtls; +pub mod access_request; pub mod audit; pub mod audit_jsonl; pub mod audit_sink; @@ -32,6 +33,7 @@ pub mod errors; pub mod failover; pub mod forward_proxy; pub mod governance; +pub mod governed_services; pub mod guardrails; pub mod handoff; pub mod inference_policy_loader; @@ -52,6 +54,7 @@ pub mod spawn; #[path = "../../shared/sre_privacy.rs"] mod sre_privacy; pub mod sre_proxy; +pub mod task_telemetry; pub mod telemetry; /// Select RustCrypto for JWT signing and verification. Workspace builds also diff --git a/inference-router/src/main.rs b/inference-router/src/main.rs index a03b90377..0bec3fcae 100644 --- a/inference-router/src/main.rs +++ b/inference-router/src/main.rs @@ -331,6 +331,7 @@ async fn main() -> Result<()> { .unwrap_or_default(), ); + let services_shutdown = state.services.shutdown.clone(); let app = { // Public routes — no admin token required (health, metrics, inference, Foundry proxies, mesh) let public = Router::new() @@ -457,16 +458,19 @@ async fn main() -> Result<()> { let memory_binding_for_platform = state.memory_binding.clone(); let policy_status_for_platform = state.policy_status.clone(); + let telemetry = state.services.telemetry.clone(); + let services = routes::governed_service_routes(state.clone()).with_state(state.clone()); let merged = public .merge(protected) .merge(handoff_init) .merge(handoff_mutations) .merge(handoff_status) .with_state(state) - .merge(build_mcp_router().await) + .merge(build_mcp_router(Some(telemetry.clone())).await) .merge(build_platform_mcp_router( Some(memory_binding_for_platform), Some(policy_status_for_platform), + Some(telemetry), )); let merged = if let Some(a2a) = a2a_router_opt { @@ -483,6 +487,9 @@ async fn main() -> Result<()> { .and_then(|s| s.parse::().ok()) .unwrap_or(256), )) + // Operator controls must remain reachable while inference requests + // or bounded approval waits occupy their own concurrency limits. + .merge(services) // r6 — trace-id middleware is outermost so every request gets a // trace span before any other layer runs (concurrency limit, // connection_close, auth gates all log inside the span). @@ -569,6 +576,7 @@ async fn main() -> Result<()> { let (signal_fired_tx, signal_fired_rx) = tokio::sync::oneshot::channel::<()>(); let shutdown_fut = async move { shutdown_signal().await; + services_shutdown.cancel(); let _ = signal_fired_tx.send(()); }; @@ -631,7 +639,9 @@ async fn main() -> Result<()> { /// falling back to the unauthenticated dev route. Operators see a clear /// startup-time error instead of a route that quietly serves /// unauthenticated MCP traffic. -async fn build_mcp_router() -> Router { +async fn build_mcp_router( + telemetry: Option>, +) -> Router { use kars_inference_router::mcp::forwarder::RouterToolDispatcher; use kars_inference_router::mcp::oauth::OAuthVerifierConfig; use kars_inference_router::mcp::registry; @@ -683,6 +693,7 @@ async fn build_mcp_router() -> Router { }; let mut state = routes::McpRouteState::standard(); + state.task_telemetry = telemetry; if let Some(d) = dispatcher_arc { state = state.with_tools(d); } @@ -786,8 +797,10 @@ fn build_platform_mcp_router( policy_status: Option< std::sync::Arc, >, + telemetry: Option>, ) -> Router { - let state = routes::McpRouteState::platform(memory_binding, policy_status); + let mut state = routes::McpRouteState::platform(memory_binding, policy_status); + state.task_telemetry = telemetry; tracing::info!( "Mounting /platform/mcp (Foundry-shim discovery surface, loopback-only, no OAuth)" ); diff --git a/inference-router/src/proxy.rs b/inference-router/src/proxy.rs index 0d40e7ed7..122ff9605 100644 --- a/inference-router/src/proxy.rs +++ b/inference-router/src/proxy.rs @@ -34,6 +34,7 @@ mod authentication_tests; /// Upstream configuration for a single request. #[derive(Clone)] pub struct UpstreamConfig { + pub telemetry: Option>, pub endpoint: String, pub deployment: String, pub sandbox_name: String, @@ -51,11 +52,18 @@ pub struct UpstreamConfig { } impl UpstreamConfig { + fn telemetry_provider(&self) -> &str { + match &self.authentication { + AuthenticationProvenance::Named { provider_id } => provider_id, + AuthenticationProvenance::LegacyDefault => self.provider.as_tag(), + } + } /// The historic constructor shape: an Azure OpenAI / Foundry /// upstream authenticated via Workload Identity / API-key mode. #[must_use] pub fn azure(endpoint: String, deployment: String, sandbox_name: String) -> Self { Self { + telemetry: None, endpoint, deployment, sandbox_name, @@ -255,9 +263,22 @@ pub async fn forward( request_body: Bytes, ) -> Result<(StatusCode, HeaderMap, Bytes)> { let start = Instant::now(); + let mut observation = upstream.telemetry.as_ref().and_then(|telemetry| { + telemetry.begin( + path, + upstream.telemetry_provider(), + &upstream.deployment, + &request_body, + ) + }); - let (upstream_url, body) = build_upstream_url(auth, upstream, path, request_body) - .map_err(ForwardFailure::configuration)?; + let (upstream_url, body) = + build_upstream_url(auth, upstream, path, request_body).map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("configuration_error"); + } + ForwardFailure::configuration(error) + })?; let mode = match upstream.provider { ProviderKind::Anthropic => "anthropic", @@ -276,12 +297,22 @@ pub async fn forward( let credential = credential_for_upstream(auth, copilot, upstream) .await - .map_err(ForwardFailure::authentication)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("authentication_error"); + } + ForwardFailure::authentication(error) + })?; let headers = build_upstream_headers(request_headers, auth, &credential, &upstream.endpoint) - .map_err(ForwardFailure::configuration)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("configuration_error"); + } + ForwardFailure::configuration(error) + })?; - tracing::info!(sandbox = %upstream.sandbox_name, url = %upstream_url, body_len = body.len(), "Sending upstream request"); + tracing::info!(sandbox = %upstream.sandbox_name, body_len = body.len(), "Sending upstream request"); let retryable = is_idempotent(&method, path); let response = send_with_retry( @@ -293,11 +324,19 @@ pub async fn forward( retryable, &upstream.sandbox_name, ) - .await?; + .await + .inspect_err(|_| { + if let Some(observation) = observation.as_mut() { + observation.fail("transport_error"); + } + })?; let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY); let response_headers = response.headers().clone(); + if let Some(observation) = observation.as_mut() { + observation.headers(status.as_u16()); + } // r6 — surface Azure-side request ids so one log line carries both our // trace_id (from the outer tracing span) and Azure's correlation ids. @@ -312,13 +351,18 @@ pub async fn forward( .and_then(|v| v.to_str().ok()) .unwrap_or(""); - let response_body = response - .bytes() - .await - .map_err(|error| ForwardFailure::response_body(status, error))?; + let response_body = response.bytes().await.map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("response_body_error"); + } + ForwardFailure::response_body(status, error) + })?; let latency = start.elapsed(); record_metrics(upstream, status, latency, &response_body); + if let Some(observation) = observation.as_mut() { + observation.buffered(status.as_u16(), &response_body); + } tracing::info!( sandbox = %upstream.sandbox_name, @@ -485,6 +529,14 @@ pub async fn forward_stream( HeaderMap, futures::stream::BoxStream<'static, Result>, )> { + let mut observation = upstream.telemetry.as_ref().and_then(|telemetry| { + telemetry.begin( + path, + upstream.telemetry_provider(), + &upstream.deployment, + &request_body, + ) + }); // Inject stream_options.include_usage into request body so the final // SSE chunk contains a `usage` object with token counts. ONLY for // chat/completions — Anthropic Messages API (/v1/messages) rejects @@ -501,15 +553,30 @@ pub async fn forward_stream( inject_stream_usage(request_body) }; let (upstream_url, body) = build_upstream_url(&auth, &upstream, path, body_with_usage) - .map_err(ForwardFailure::configuration)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("configuration_error"); + } + ForwardFailure::configuration(error) + })?; tracing::info!(sandbox = %upstream.sandbox_name, model = %upstream.deployment, mode = "stream", "Forwarding SSE stream"); let credential = credential_for_upstream(&auth, copilot.as_deref(), &upstream) .await - .map_err(ForwardFailure::authentication)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("authentication_error"); + } + ForwardFailure::authentication(error) + })?; let headers = build_upstream_headers(&request_headers, &auth, &credential, &upstream.endpoint) - .map_err(ForwardFailure::configuration)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("configuration_error"); + } + ForwardFailure::configuration(error) + })?; let start = Instant::now(); @@ -520,11 +587,19 @@ pub async fn forward_stream( .timeout(INFERENCE_REQUEST_TIMEOUT) .send() .await - .map_err(ForwardFailure::transport)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("transport_error"); + } + ForwardFailure::transport(error) + })?; let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY); let response_headers = response.headers().clone(); + if let Some(observation) = observation.as_mut() { + observation.headers(status.as_u16()); + } // r6 — log Azure correlation ids for the stream path too. Emitted at // stream-start because headers arrive before any bytes. @@ -555,22 +630,24 @@ pub async fn forward_stream( .inc(); // On non-success, the upstream body is a short JSON error (not an SSE - // stream). Eagerly drain it, log the contents (capped), and forward as a + // stream). Eagerly drain it and forward as a // single chunk so callers see the actual reason. Without this we only // see "status=413" and have to guess at causes (token cap? bytes cap? - // schema break?). Cap at 4 KiB so a misbehaving upstream can't blow logs. + // schema break?). Do not put provider API bodies into logs. if !status.is_success() { - let body_bytes = response - .bytes() - .await - .map_err(|error| ForwardFailure::response_body(status, error))?; - let preview = String::from_utf8_lossy(&body_bytes); - let preview_trimmed: String = preview.chars().take(2048).collect(); + let body_bytes = response.bytes().await.map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("response_body_error"); + } + ForwardFailure::response_body(status, error) + })?; + if let Some(observation) = observation.as_mut() { + observation.buffered(status.as_u16(), &body_bytes); + } tracing::warn!( sandbox = %upstream.sandbox_name, status = %status.as_u16(), body_len = body_bytes.len(), - body = %preview_trimmed, "Upstream returned non-success status" ); let stream = futures::stream::once(async move { Ok::<_, reqwest::Error>(body_bytes) }); @@ -632,7 +709,15 @@ pub async fn forward_stream( chunk }); - Ok((status, response_headers, metered.boxed())) + let is_sse = response_headers + .get("content-type") + .and_then(|header| header.to_str().ok()) + .is_some_and(|value| value.starts_with("text/event-stream")); + Ok(( + status, + response_headers, + crate::task_telemetry::observe::wrap_stream(metered.boxed(), observation, is_sse), + )) } /// Inject `stream_options: { include_usage: true }` into the request body diff --git a/inference-router/src/routes/access_request.rs b/inference-router/src/routes/access_request.rs new file mode 100644 index 000000000..e15de9197 --- /dev/null +++ b/inference-router/src/routes/access_request.rs @@ -0,0 +1,269 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::AppState; +use crate::access_request::{Error, Request, Status}; +use axum::{ + Json, Router, + extract::{ConnectInfo, DefaultBodyLimit, Path, Query, State}, + http::{HeaderMap, StatusCode}, + middleware::{self, Next}, + response::{IntoResponse, Response}, + routing::{get, post}, +}; +use serde::Deserialize; +use serde_json::json; +use std::{net::SocketAddr, time::Duration}; + +pub fn routes(state: AppState) -> Router { + let agent = Router::new() + .route("/v1/access-request", post(create)) + .route("/v1/access-requests", get(list_agent)) + .route("/v1/access-requests/{id}/cancel", post(cancel)) + .route("/v1/access-requests/{id}/wait", get(wait)) + .merge(super::task_telemetry::routes()) + .route_layer(middleware::from_fn_with_state(state.clone(), agent_auth)) + .layer(tower::limit::ConcurrencyLimitLayer::new(24)); + let admin = Router::new() + .route("/internal/access-requests", get(inspect)) + .route("/internal/access-requests/reset", post(reset)) + .route("/internal/access-requests/decision", post(decide)) + .route_layer(middleware::from_fn_with_state(state, control_auth)) + .layer(tower::limit::ConcurrencyLimitLayer::new(8)); + agent.merge(admin).layer(DefaultBodyLimit::max(8192)) +} + +fn bearer(headers: &HeaderMap) -> Option<&str> { + headers + .get("authorization") + .and_then(|header| header.to_str().ok()) + .and_then(|header| header.strip_prefix("Bearer ")) +} + +async fn control_auth( + State(state): State, + request: axum::extract::Request, + next: Next, +) -> Response { + if !state.services.identity_valid { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Service identity is unavailable", + ) + .into_response(); + } + if !state.services.control_configured() { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Service control credential is not configured", + ) + .into_response(); + } + if !state.services.control_authorized(bearer(request.headers())) { + return ( + StatusCode::UNAUTHORIZED, + "Service control credential required; agent admin credentials are not accepted", + ) + .into_response(); + } + if let Some(allowed) = &state.services.allow_ips { + let remote = request + .extensions() + .get::>() + .map(|info| info.0.ip()); + if remote.is_none_or(|remote| !allowed.contains(&remote)) { + return (StatusCode::FORBIDDEN, "Control origin is not allowed").into_response(); + } + } + next.run(request).await +} + +async fn agent_auth( + State(state): State, + request: axum::extract::Request, + next: Next, +) -> Response { + if !state.services.identity_valid { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Service identity is unavailable", + ) + .into_response(); + } + let local = request + .extensions() + .get::>() + .is_some_and(|info| info.0.ip().is_loopback()); + if !local && !state.services.control_authorized(bearer(request.headers())) { + return ( + StatusCode::UNAUTHORIZED, + "Same-router agent or service control authentication required", + ) + .into_response(); + } + if !local && let Some(allowed) = &state.services.allow_ips { + let remote = request + .extensions() + .get::>() + .map(|info| info.0.ip()); + if remote.is_none_or(|ip| !allowed.contains(&ip)) { + return (StatusCode::FORBIDDEN, "Control origin is not allowed").into_response(); + } + } + next.run(request).await +} + +pub(super) fn error(error: Error) -> Response { + let (status, code) = match error { + Error::Invalid => (StatusCode::BAD_REQUEST, "invalid_request"), + Error::StaleScope => (StatusCode::CONFLICT, "stale_scope"), + Error::Full => (StatusCode::TOO_MANY_REQUESTS, "capacity_exhausted"), + Error::RateLimited => (StatusCode::TOO_MANY_REQUESTS, "rate_limited"), + Error::Missing => (StatusCode::NOT_FOUND, "request_not_found"), + Error::Expired => (StatusCode::GONE, "request_expired"), + Error::WaitTimeout => (StatusCode::REQUEST_TIMEOUT, "wait_timed_out"), + Error::Terminal => (StatusCode::CONFLICT, "request_terminal"), + Error::Unavailable => (StatusCode::SERVICE_UNAVAILABLE, "service_unavailable"), + }; + (status, Json(json!({"error":code}))).into_response() +} +pub(super) fn scope_header(headers: &HeaderMap) -> Result<&str, Error> { + headers + .get("x-kars-service-scope") + .and_then(|header| header.to_str().ok()) + .filter(|id| id.len() <= 128) + .ok_or(Error::StaleScope) +} + +async fn create(State(state): State, Json(request): Json) -> Response { + match state.services.requests.record(request) { + Ok((entry, new)) => ( + StatusCode::ACCEPTED, + Json(json!({"status":"queued","new":new, + "kind":entry.kind,"target":entry.target,"request":entry,"enforcement_changed":false})), + ) + .into_response(), + Err(err) => error(err), + } +} +async fn list_agent(State(state): State, headers: HeaderMap) -> Response { + let scope = match state.services.requests.scope() { + Ok(scope) => scope, + Err(err) => return error(err), + }; + if let Some(requested) = headers.get("x-kars-service-scope") + && requested.to_str().ok() != Some(scope.id.as_str()) + { + return error(Error::StaleScope); + } + match state.services.requests.snapshot(&scope.id) { + Ok(entries) => Json(json!({"scope":scope,"requests":entries,"enforcement_changed":false})) + .into_response(), + Err(err) => error(err), + } +} +async fn inspect(State(state): State, headers: HeaderMap) -> Response { + let scope = match state.services.requests.scope() { + Ok(scope) => scope, + Err(err) => return error(err), + }; + if let Some(requested) = headers.get("x-kars-service-scope") + && requested.to_str().ok() != Some(scope.id.as_str()) + { + return error(Error::StaleScope); + } + match state.services.requests.snapshot(&scope.id) { + Ok(entries) => Json( + json!({"schema_version":1,"scope":scope,"sandbox":state.sandbox_name.as_str(), + "count":entries.len(),"entries":entries,"enforcement_changed":false}), + ) + .into_response(), + Err(err) => error(err), + } +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct ScopeBody { + scope_id: String, +} +async fn cancel( + State(state): State, + Path(id): Path, + Json(body): Json, +) -> Response { + match state + .services + .requests + .transition(&body.scope_id, &id, Status::Cancelled) + { + Ok(entry) => Json(entry).into_response(), + Err(err) => error(err), + } +} +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Reset { + scope_id: String, + assignment_id: Option, +} +async fn reset(State(state): State, Json(body): Json) -> Response { + match state.services.reset(&body.scope_id, body.assignment_id) { + Ok((scope, cleared)) => { + Json(json!({"scope":scope,"cleared":cleared,"enforcement_changed":false})) + .into_response() + } + Err(err) => error(err), + } +} +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Decision { + scope_id: String, + request_id: String, + verdict: Status, +} +async fn decide(State(state): State, Json(body): Json) -> Response { + if !matches!(body.verdict, Status::Approved | Status::Denied) { + return error(Error::Invalid); + } + match state + .services + .requests + .transition(&body.scope_id, &body.request_id, body.verdict) + { + Ok(entry) => Json(json!({"request":entry,"enforcement_changed":false})).into_response(), + Err(err) => error(err), + } +} +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Wait { + scope_id: String, + #[serde(default = "default_wait")] + timeout_ms: u64, +} +fn default_wait() -> u64 { + 30_000 +} +async fn wait( + State(state): State, + Path(id): Path, + Query(query): Query, +) -> Response { + if query.timeout_ms == 0 || query.timeout_ms > 300_000 { + return error(Error::Invalid); + } + match state + .services + .wait_for_decision( + &query.scope_id, + &id, + Duration::from_millis(query.timeout_ms), + ) + .await + { + Ok(entry) => Json(entry).into_response(), + Err(err) => error(err), + } +} diff --git a/inference-router/src/routes/egress.rs b/inference-router/src/routes/egress.rs index cf5964ca5..9680dbd65 100644 --- a/inference-router/src/routes/egress.rs +++ b/inference-router/src/routes/egress.rs @@ -113,10 +113,30 @@ async fn egress_fetch( State(state): State, Json(req): Json, ) -> impl IntoResponse { + let started = std::time::Instant::now(); + let telemetry_scope = state.services.telemetry.cursor().0; let url = req.get("url").and_then(|v| v.as_str()).unwrap_or(""); let method = req.get("method").and_then(|v| v.as_str()).unwrap_or("GET"); let req_body = req.get("body").and_then(|v| v.as_str()).unwrap_or(""); let req_headers = req.get("headers").and_then(|v| v.as_object()); + let wait_ms = req + .get("wait_for_approval_ms") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + if req + .get("wait_for_approval_ms") + .is_some_and(|value| value.as_u64().is_none()) + { + return super::access_request::error(crate::access_request::Error::Invalid); + } + if wait_ms > 300_000 { + return super::access_request::error(crate::access_request::Error::Invalid); + } + if wait_ms > 0 + && req.get("scope_id").and_then(serde_json::Value::as_str) != Some(telemetry_scope.as_str()) + { + return super::access_request::error(crate::access_request::Error::StaleScope); + } if url.is_empty() { return errors::flat(StatusCode::BAD_REQUEST, "Missing 'url' field").into_response(); @@ -134,7 +154,15 @@ async fn egress_fetch( } }; if is_private { - tracing::warn!(url = %url, "Egress fetch blocked: private/internal target"); + tracing::warn!("Egress fetch blocked: private/internal target"); + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + Some(403), + 0, + true, + ); return ( StatusCode::FORBIDDEN, Json(serde_json::json!({ @@ -148,21 +176,93 @@ async fn egress_fetch( } let sandbox: &str = &state.sandbox_name; + let mut dispatch_request = None; // Check egress access: blocklist → allowlist (Strict denies the rest). if let Err(reason) = state.blocklist.check_egress(url, sandbox).await { - tracing::warn!(url = %url, reason = %reason, "Egress fetch denied"); - return (StatusCode::FORBIDDEN, Json(serde_json::json!({ + tracing::warn!("Egress fetch denied by policy"); + let target = reqwest::Url::parse(url).ok().and_then(|url| { + url.host_str() + .zip(url.port_or_known_default()) + .map(|(host, port)| (host.to_string(), port)) + }); + let entry = target.as_ref().and_then(|(host, port)| { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + state + .blocked_egress + .record_at(std::time::Instant::now(), now, sandbox, host, *port); + state + .services + .record_egress_in_scope(&telemetry_scope, host, *port) + .ok() + }); + state + .services + .telemetry + .record_policy(&telemetry_scope, "egress", false); + if wait_ms > 0 { + let Some(entry) = entry else { + return super::access_request::error(crate::access_request::Error::Full); + }; + if let Err(error) = state + .services + .wait_for_egress( + &state.blocklist, + &telemetry_scope, + &entry.request_id, + url, + sandbox, + std::time::Duration::from_millis(wait_ms), + ) + .await + { + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + Some(403), + started.elapsed().as_millis() as u64, + false, + ); + return super::access_request::error(error); + } + dispatch_request = Some(entry.request_id); + } else { + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + Some(403), + 0, + true, + ); + return (StatusCode::FORBIDDEN, Json(serde_json::json!({ "error": reason, "url": url, "action": "If legitimate, an operator can add the host to the baseline allowlist ('kars egress --approve ', which re-signs) or grant it temporarily ('kars egress allow-extra --host --ttl --reason ').", - }))).into_response(); + }))).into_response(); + } + } + if wait_ms > 0 + && state + .services + .requests + .scope() + .map(|scope| scope.id) + .ok() + .as_deref() + != Some(telemetry_scope.as_str()) + { + return super::access_request::error(crate::access_request::Error::StaleScope); } // Record in learn mode state.blocklist.record_learned(url).await; - tracing::info!(url = %url, method = %method, "Egress fetch proxied"); + tracing::info!("Egress fetch proxied"); // Build and send the request let http_method = match method.to_uppercase().as_str() { @@ -204,6 +304,27 @@ async fn egress_fetch( const MAX_RESPONSE_BYTES: usize = 2 * 1024 * 1024; // 2 MB + // No await separates this atomic claim from starting the send. A successful + // cancel/reset before the claim wins; afterwards prevention is not promised. + let _dispatch_claim = if wait_ms > 0 { + if state.blocklist.check_egress(url, sandbox).await.is_err() { + return errors::flat( + StatusCode::FORBIDDEN, + "Egress policy no longer permits this operation", + ) + .into_response(); + } + match state + .services + .claim_egress_dispatch(&telemetry_scope, dispatch_request.as_deref()) + { + Ok(claim) => Some(claim), + Err(error) => return super::access_request::error(error), + } + } else { + None + }; + match request .timeout(std::time::Duration::from_secs(30)) .send() @@ -236,32 +357,79 @@ async fn egress_fetch( }) .collect(); // Cap response body to prevent OOM - let body_bytes = resp.bytes().await.unwrap_or_default(); - let body = if body_bytes.len() > MAX_RESPONSE_BYTES { - let truncated = String::from_utf8_lossy(&body_bytes[..MAX_RESPONSE_BYTES]); - format!( - "{}... [truncated at {} bytes]", - truncated, MAX_RESPONSE_BYTES - ) - } else { - String::from_utf8_lossy(&body_bytes).into_owned() - }; + use futures::StreamExt; + let mut stream = resp.bytes_stream(); + let mut body_bytes = Vec::new(); + let mut truncated = false; + while let Some(chunk) = stream.next().await { + let chunk = match chunk { + Ok(chunk) => chunk, + Err(_) => { + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + Some(status), + started.elapsed().as_millis() as u64, + false, + ); + return errors::flat( + StatusCode::BAD_GATEWAY, + "Upstream response body failed", + ) + .into_response(); + } + }; + let available = MAX_RESPONSE_BYTES - body_bytes.len(); + body_bytes.extend_from_slice(&chunk[..chunk.len().min(available)]); + if chunk.len() > available { + truncated = true; + break; + } + } + let mut body = String::from_utf8_lossy(&body_bytes).into_owned(); + if truncated { + body.push_str(&format!("... [truncated at {MAX_RESPONSE_BYTES} bytes]")); + } + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + if status >= 400 { + Some(false) + } else if truncated { + None + } else { + Some(true) + }, + Some(status), + started.elapsed().as_millis() as u64, + !truncated, + ); ( StatusCode::OK, Json(serde_json::json!({ "status": status, "headers": resp_headers, "body": body, + "truncated": truncated, })), ) .into_response() } - Err(e) => { - tracing::warn!(url = %url, error = %e, "Egress fetch failed"); + Err(_) => { + tracing::warn!("Egress fetch failed"); + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + None, + started.elapsed().as_millis() as u64, + false, + ); ( StatusCode::BAD_GATEWAY, Json(serde_json::json!({ - "error": format!("Request failed: {}", e), + "error": "Upstream request failed", "url": url, })), ) diff --git a/inference-router/src/routes/governance.rs b/inference-router/src/routes/governance.rs index 258fea91e..20e6db04d 100644 --- a/inference-router/src/routes/governance.rs +++ b/inference-router/src/routes/governance.rs @@ -60,12 +60,17 @@ async fn agt_evaluate( .and_then(|v| v.as_str()) .unwrap_or("unknown"); let extra_context = body.get("context").cloned(); + let telemetry_scope = state.services.telemetry.cursor().0; // Per-tool sliding window rate limit (AGT McpSlidingRateLimiter) // Extract tool name from action format "tool:exec_command" or "tool:http_fetch" if let Some(tool_name) = action.strip_prefix("tool:") { let (allowed, retry_after) = state.governance.check_tool_rate(tool_name); if !allowed { + state + .services + .telemetry + .record_policy(&telemetry_scope, action, false); return Json(serde_json::json!({ "allowed": false, "reason": format!("per-tool rate limit exceeded for '{}'", tool_name), @@ -79,6 +84,41 @@ async fn agt_evaluate( let result = state .governance .evaluate(agent_id, action, extra_context.as_ref()); + if agent_id == state.sandbox_name.as_str() + && let Some(tool) = action.strip_prefix("tool:") + { + let allowed = result.get("allowed").and_then(serde_json::Value::as_bool) == Some(true); + state + .services + .telemetry + .record_policy(&telemetry_scope, action, allowed); + if allowed { + if let Some(context) = extra_context.as_ref() + && context.get("scope_id").and_then(serde_json::Value::as_str) + == Some(telemetry_scope.as_str()) + && let Some(id) = context + .get("tool_call_id") + .and_then(serde_json::Value::as_str) + { + state + .services + .telemetry + .authorize_harness_tool(&telemetry_scope, id, tool); + } + } else { + let _ = state + .services + .requests + .record(crate::access_request::Request { + scope_id: telemetry_scope, + kind: "tool".into(), + target: tool.into(), + reason: "Governance denied the requested tool".into(), + tier: None, + port: None, + }); + } + } Json(result).into_response() } diff --git a/inference-router/src/routes/mcp.rs b/inference-router/src/routes/mcp.rs index 7be0798ce..7066d1230 100644 --- a/inference-router/src/routes/mcp.rs +++ b/inference-router/src/routes/mcp.rs @@ -62,6 +62,7 @@ pub const MCP_SESSION_HEADER: &str = "Mcp-Session-Id"; /// Per-router MCP state. Cheap to clone (everything inside is `Arc`). #[derive(Clone)] pub struct McpRouteState { + pub task_telemetry: Option>, pub config: Arc, pub minter: Arc, pub tools: Arc, @@ -76,6 +77,7 @@ impl McpRouteState { /// `McpServer.spec.url`; see [`with_tools`]. pub fn standard() -> Self { Self { + task_telemetry: None, config: Arc::new(InitializeConfig::default()), minter: Arc::new(OsRngSessionMinter), tools: Arc::new(SyncToAsync::new(EchoDispatcher::standard())), @@ -123,6 +125,7 @@ impl McpRouteState { config: Arc::new(InitializeConfig::default()), minter: Arc::new(OsRngSessionMinter), tools: Arc::new(dispatcher), + task_telemetry: None, } } } @@ -188,6 +191,10 @@ async fn method_not_allowed() -> impl IntoResponse { } async fn post_mcp(State(state): State, headers: HeaderMap, body: Bytes) -> Response { + let telemetry_scope = state + .task_telemetry + .as_ref() + .map(|telemetry| telemetry.cursor().0); let accept = headers .get(header::ACCEPT) .and_then(|v| v.to_str().ok()) @@ -204,6 +211,15 @@ async fn post_mcp(State(state): State, headers: HeaderMap, body: Some(state.tools.as_ref()), ) .await; + if let (Some(telemetry), Some(scope)) = (&state.task_telemetry, telemetry_scope) { + crate::task_telemetry::mcp::record( + telemetry, + &scope, + &body, + &outcome, + started.elapsed().as_millis() as u64, + ); + } let status_label: &str = match &outcome { ProcessOutcome::JsonRpcResponse { .. } => "200", @@ -304,6 +320,7 @@ mod tests { fn test_state() -> McpRouteState { McpRouteState { + task_telemetry: None, config: Arc::new(InitializeConfig::default()), minter: Arc::new(FixedMinter("test-session-001")), tools: Arc::new(SyncToAsync::new(EchoDispatcher::standard())), @@ -493,6 +510,7 @@ mod tests { McpRouteState { config: Arc::new(InitializeConfig::default()), minter: Arc::new(FixedMinter("platform-session-001")), + task_telemetry: None, tools: Arc::new(crate::mcp::PlatformDispatcher::with_base_url( "http://127.0.0.1:1", )), @@ -679,6 +697,7 @@ mod tests { // Real platform dispatcher, pointed at the mock upstream. let state = McpRouteState { + task_telemetry: None, config: Arc::new(InitializeConfig::default()), minter: Arc::new(FixedMinter("platform-session-001")), tools: Arc::new(crate::mcp::PlatformDispatcher::with_base_url( diff --git a/inference-router/src/routes/mod.rs b/inference-router/src/routes/mod.rs index ff8d134a6..a2eb5a6ea 100644 --- a/inference-router/src/routes/mod.rs +++ b/inference-router/src/routes/mod.rs @@ -44,7 +44,10 @@ pub use governance::sensitive_agt_routes; mod mesh; pub use mesh::mesh_routes; +mod access_request; mod mesh_token; +mod task_telemetry; +pub use access_request::routes as governed_service_routes; mod model_routing; pub use mesh_token::mesh_token_routes; @@ -71,6 +74,7 @@ pub use a2a::{A2aRouteState, a2a_routes}; /// Shared application state. #[derive(Clone)] pub struct AppState { + pub services: Arc, pub auth: Arc, /// GitHub Copilot token cache. Constructed unconditionally (cheap, lazy) /// — `proxy::forward()` only consults it when the upstream endpoint is @@ -310,7 +314,13 @@ impl AppState { ) .await; + let services = Arc::new(crate::governed_services::GovernedServices::from_env( + &sandbox_name, + )); + let blocked_egress = Arc::new(BlockedBuffer::with_defaults()); + blocked_egress.bind_services(&services); Ok(Self { + services, auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(CopilotTokenCache::from_env()), client: client.clone(), @@ -321,7 +331,7 @@ impl AppState { signing_provider: Arc::clone(&governance) as Arc, governance, blocklist, - blocked_egress: Arc::new(BlockedBuffer::with_defaults()), + blocked_egress, sandbox_name: Arc::new(sandbox_name), inbox: Arc::new(MeshInbox::new()), mesh_metrics: Arc::new(MeshMetrics::new()), @@ -367,7 +377,11 @@ impl AppState { .and_then(|g| g.clone()) .unwrap_or_else(|| self.config.default_model.clone()); - UpstreamConfig::azure(endpoint, deployment, sandbox_name.to_string()) + let mut upstream = UpstreamConfig::azure(endpoint, deployment, sandbox_name.to_string()); + if self.services.identity_valid { + upstream.telemetry = Some(self.services.telemetry.clone()); + } + upstream } } diff --git a/inference-router/src/routes/model_routing.rs b/inference-router/src/routes/model_routing.rs index 61aa5d1ca..9ecd35dc0 100644 --- a/inference-router/src/routes/model_routing.rs +++ b/inference-router/src/routes/model_routing.rs @@ -376,6 +376,7 @@ mod tests { policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(crate::auth::WorkloadIdentityAuth::new()), copilot: Arc::new(crate::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/src/routes/task_telemetry.rs b/inference-router/src/routes/task_telemetry.rs new file mode 100644 index 000000000..999f81646 --- /dev/null +++ b/inference-router/src/routes/task_telemetry.rs @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + AppState, + access_request::{error, scope_header}, +}; +use crate::access_request::Error; +use axum::{ + Json, Router, + extract::{Query, State}, + http::HeaderMap, + response::{IntoResponse, Response}, + routing::{get, post}, +}; +use serde::Deserialize; +use serde_json::json; + +pub(super) fn routes() -> Router { + Router::new() + .route("/telemetry/cursor", get(cursor)) + .route("/telemetry/trace", get(trace)) + .route("/telemetry/tool", post(tool)) + .route("/telemetry/budget", get(budget)) +} +async fn cursor(State(state): State) -> Response { + let (scope_id, cursor) = state.services.telemetry.cursor(); + Json(json!({"scope_id":scope_id,"cursor":cursor,"durable":false})).into_response() +} +#[derive(Deserialize)] +struct Since { + #[serde(default)] + since: u64, +} +async fn trace( + State(state): State, + headers: HeaderMap, + Query(query): Query, +) -> Response { + let scope = match scope_header(&headers) { + Ok(scope) => scope, + Err(err) => return error(err), + }; + match state.services.telemetry.snapshot(scope, query.since) { + Some(trace) => Json(trace).into_response(), + None => error(Error::StaleScope), + } +} +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Tool { + scope_id: String, + call_id: String, + ok: bool, + latency_ms: u64, +} +async fn tool(State(state): State, Json(body): Json) -> Response { + if state.services.telemetry.complete_harness_tool( + &body.scope_id, + &body.call_id, + body.ok, + body.latency_ms, + ) { + Json(json!({"recorded":true,"source":"harness-reported","enforcement_changed":false})) + .into_response() + } else { + error(Error::Terminal) + } +} +async fn budget(State(state): State, headers: HeaderMap) -> Response { + let scope = match scope_header(&headers) { + Ok(scope) => scope, + Err(err) => return error(err), + }; + if state.services.telemetry.cursor().0 != scope { + return error(Error::StaleScope); + } + let (daily, monthly) = state.budget.get_usage(&state.sandbox_name).await; + let policy = crate::inference_policy_loader::current_snapshot(&state.inference_policy).await; + let daily_limit = policy + .daily_tokens + .unwrap_or(state.config.token_budget_daily); + Json( + json!({"scope_id":scope,"daily_observed_tokens":daily,"monthly_observed_tokens":monthly, + "daily_limit":(daily_limit>0).then_some(daily_limit), + "monthly_limit":policy.monthly_tokens.filter(|limit|*limit>0), + "shared_task_budget_enforced":false,"coverage":"existing-per-sandbox-counter-only", + "assignment_reset_resets_budget":false}), + ) + .into_response() +} diff --git a/inference-router/src/task_telemetry.rs b/inference-router/src/task_telemetry.rs new file mode 100644 index 000000000..ff5779529 --- /dev/null +++ b/inference-router/src/task_telemetry.rs @@ -0,0 +1,259 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Bounded router observations, not a durable execution or budget ledger. +//! No prompts, tool arguments/results, URLs, headers or credentials are retained. + +use serde_json::{Value, json}; +use std::{ + collections::{HashMap, HashSet, VecDeque}, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +pub mod mcp; +pub mod observe; +pub mod parse; + +const MAX_EVENTS: usize = 1024; +const MAX_PENDING: usize = 256; +const TOOL_TTL: Duration = Duration::from_secs(20 * 60); + +struct Pending { + name: String, + at: Instant, + round: u64, +} +struct Inner { + scope: String, + seq: u64, + round: u64, + dropped: u64, + events: VecDeque, + model_tools: HashMap, + seen_model: HashSet, + harness_tools: HashMap, + seen_harness: HashSet, +} + +pub struct TaskTelemetry { + inner: Mutex, +} + +impl TaskTelemetry { + pub fn new(scope: String) -> Self { + Self { + inner: Mutex::new(Inner { + scope, + seq: 0, + round: 0, + dropped: 0, + events: VecDeque::new(), + model_tools: HashMap::new(), + seen_model: HashSet::new(), + harness_tools: HashMap::new(), + seen_harness: HashSet::new(), + }), + } + } + pub fn reset(&self, scope: String) { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + inner.scope = scope; + inner.events.clear(); + inner.model_tools.clear(); + inner.seen_model.clear(); + inner.harness_tools.clear(); + inner.seen_harness.clear(); + inner.round = 0; + inner.dropped = 0; + } + fn push(inner: &mut Inner, mut event: Value) { + inner.seq = inner.seq.saturating_add(1); + event["seq"] = json!(inner.seq); + event["scope_id"] = json!(inner.scope); + inner.events.push_back(event); + if inner.events.len() > MAX_EVENTS { + inner.events.pop_front(); + inner.dropped = inner.dropped.saturating_add(1); + } + } + pub fn cursor(&self) -> (String, u64) { + let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + (inner.scope.clone(), inner.seq) + } + pub fn snapshot(&self, scope: &str, since: u64) -> Option { + let inner = self.inner.lock().ok()?; + if inner.scope != scope { + return None; + } + let events = inner + .events + .iter() + .filter(|event| event["seq"].as_u64().is_some_and(|seq| seq > since)) + .take(256) + .cloned() + .collect::>(); + let cursor = events + .last() + .and_then(|event| event["seq"].as_u64()) + .unwrap_or(inner.seq); + Some(json!({ + "scope_id": scope, "cursor": cursor, "high_water": inner.seq, + "has_more": cursor, + path: &str, + provider: &str, + model: &str, + body: &[u8], + ) -> Option { + let shape = parse::Shape::for_path(path)?; + let mut inner = self.inner.lock().ok()?; + inner + .model_tools + .retain(|_, pending| pending.at.elapsed() < TOOL_TTL); + for (id, ok) in parse::request_results(body, shape) { + if let Some(pending) = inner.model_tools.remove(&id) { + Self::push( + &mut inner, + json!({"kind":"tool_result", "call_id":id, "name":pending.name, + "round":pending.round, "ok":ok, "source":"harness-reported"}), + ); + } + } + inner.round = inner.round.saturating_add(1); + Some(observe::Observation::new( + self.clone(), + inner.scope.clone(), + inner.round, + shape, + parse::identifier(provider, 64), + parse::identifier(model, 253), + )) + } + pub(crate) fn finish( + &self, + scope: &str, + round: u64, + event: Value, + tools: Vec<(String, String)>, + ) { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return; + } + Self::push(&mut inner, event); + for (id, name) in tools { + Self::push( + &mut inner, + json!({"kind":"tool_proposed", "round":round, "call_id":id, + "name":name, "ok":null, "source":"model-proposed"}), + ); + if !id.is_empty() + && inner.seen_model.len() < MAX_PENDING + && inner.seen_model.insert(id.clone()) + { + inner.model_tools.insert( + id, + Pending { + name, + at: Instant::now(), + round, + }, + ); + } else { + inner.model_tools.remove(&id); + inner.dropped = inner.dropped.saturating_add(1); + } + } + } + pub fn record_policy(&self, scope: &str, capability: &str, allowed: bool) { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return; + } + Self::push( + &mut inner, + json!({"kind":"policy", "capability":parse::identifier(capability,128), + "allowed":allowed, "source":"router-policy"}), + ); + } + pub fn record_gap(&self, scope: &str, source: &str) { + let mut inner = self.inner.lock().unwrap_or_else(|error| error.into_inner()); + if inner.scope != scope { + return; + } + inner.dropped = inner.dropped.saturating_add(1); + Self::push( + &mut inner, + json!({"kind":"observation_gap","source":parse::identifier(source,128)}), + ); + } + pub fn record_router_tool( + &self, + scope: &str, + name: &str, + ok: Option, + status: Option, + latency_ms: u64, + complete: bool, + ) { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return; + } + Self::push( + &mut inner, + json!({"kind":"tool_result", "name":parse::identifier(name,128), + "ok":ok, "http_status":status, "ms":latency_ms, "complete":complete, "source":"router"}), + ); + } + pub fn authorize_harness_tool(&self, scope: &str, id: &str, name: &str) -> bool { + if !crate::access_request::scope::identifier(id, 128) + || parse::identifier(name, 128).is_none() + { + return false; + } + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return false; + } + if inner.seen_harness.contains(id) || inner.seen_harness.len() >= MAX_PENDING { + inner.dropped = inner.dropped.saturating_add(1); + return false; + } + inner.seen_harness.insert(id.to_string()); + let round = inner.round; + inner.harness_tools.insert( + id.to_string(), + Pending { + name: name.to_string(), + at: Instant::now(), + round, + }, + ); + true + } + pub fn complete_harness_tool(&self, scope: &str, id: &str, ok: bool, latency_ms: u64) -> bool { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return false; + } + let Some(pending) = inner.harness_tools.remove(id) else { + return false; + }; + if pending.at.elapsed() >= TOOL_TTL { + return false; + } + Self::push( + &mut inner, + json!({"kind":"tool_result", "call_id":id, "name":pending.name, + "round":pending.round, "ok":ok, "ms":latency_ms.min(3_600_000), "source":"harness-reported"}), + ); + true + } +} diff --git a/inference-router/src/task_telemetry/mcp.rs b/inference-router/src/task_telemetry/mcp.rs new file mode 100644 index 000000000..f07250ff9 --- /dev/null +++ b/inference-router/src/task_telemetry/mcp.rs @@ -0,0 +1,79 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{TaskTelemetry, parse}; +use crate::mcp::pipeline::ProcessOutcome; +use serde_json::Value; + +pub fn record( + telemetry: &TaskTelemetry, + scope: &str, + request: &[u8], + outcome: &ProcessOutcome, + latency_ms: u64, +) { + if request.len() > parse::MAX_BODY { + telemetry.record_gap(scope, "mcp_request_metadata"); + return; + } + let Ok(request) = serde_json::from_slice::(request) else { + return; + }; + let (status, reply) = match outcome { + ProcessOutcome::JsonRpcResponse { body, .. } => { + if body.len() > parse::MAX_BODY { + telemetry.record_gap(scope, "mcp_response_metadata"); + (200, None) + } else { + (200, serde_json::from_slice::(body).ok()) + } + } + ProcessOutcome::Accepted => (202, None), + ProcessOutcome::PayloadTooLarge => (413, None), + ProcessOutcome::NotAcceptable(_) => (406, None), + }; + let requests = match &request { + Value::Array(values) => values.as_slice(), + value => std::slice::from_ref(value), + }; + if requests.len() > parse::MAX_TOOLS { + telemetry.record_gap(scope, "mcp_batch_metadata"); + } + for request in requests.iter().take(parse::MAX_TOOLS) { + if request["method"] != "tools/call" { + continue; + } + let Some(name) = request["params"]["name"] + .as_str() + .and_then(|name| parse::identifier(name, 128)) + else { + continue; + }; + let id = request.get("id"); + let response = reply.as_ref().and_then(|reply| { + let replies = match reply { + Value::Array(values) => values.as_slice(), + value => std::slice::from_ref(value), + }; + id.filter(|id| id.is_number() || id.as_str().is_some_and(|id| id.len() <= 128)) + .and_then(|id| replies.iter().find(|reply| reply.get("id") == Some(id))) + }); + let ok = response.and_then(|reply| { + if reply.get("error").is_some() { + return Some(false); + } + let result = reply.get("result")?; + result + .get("isError") + .and_then(Value::as_bool) + .map(|error| !error) + .or_else(|| { + result + .get("content") + .and_then(Value::as_array) + .map(|_| true) + }) + }); + telemetry.record_router_tool(scope, &name, ok, Some(status), latency_ms, ok.is_some()); + } +} diff --git a/inference-router/src/task_telemetry/observe.rs b/inference-router/src/task_telemetry/observe.rs new file mode 100644 index 000000000..4267e89e4 --- /dev/null +++ b/inference-router/src/task_telemetry/observe.rs @@ -0,0 +1,192 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + TaskTelemetry, + parse::{self, Parsed, Shape}, +}; +use bytes::Bytes; +use futures::{Stream, stream::BoxStream}; +use serde_json::json; +use std::{ + pin::Pin, + sync::Arc, + task::{Context, Poll}, + time::Instant, +}; + +mod stream; + +pub struct Observation { + telemetry: Arc, + scope: String, + round: u64, + shape: Shape, + provider: Option, + model: Option, + started: Instant, + status: Option, + finished: bool, +} +impl Observation { + pub(super) fn new( + telemetry: Arc, + scope: String, + round: u64, + shape: Shape, + provider: Option, + model: Option, + ) -> Self { + Self { + telemetry, + scope, + round, + shape, + provider, + model, + started: Instant::now(), + status: None, + finished: false, + } + } + pub fn headers(&mut self, status: u16) { + self.status = Some(status); + } + pub fn fail(&mut self, outcome: &str) { + self.finish_parsed( + outcome, + Parsed { + partial: true, + ..Default::default() + }, + ); + } + pub fn buffered(&mut self, status: u16, body: &[u8]) { + self.headers(status); + let parsed = if body.len() <= parse::MAX_BODY { + serde_json::from_slice(body) + .ok() + .map(|value| parse::response(&value, self.shape)) + } else { + None + }; + self.finish_parsed( + if status < 400 { + "complete" + } else { + "http_error" + }, + parsed.unwrap_or(Parsed { + partial: true, + ..Default::default() + }), + ); + } + fn finish_parsed(&mut self, outcome: &str, parsed: Parsed) { + if self.finished { + return; + } + self.finished = true; + // HTTP acceptance and a complete body do not imply a successful model + // outcome (Responses may return failed/incomplete at HTTP 200). + let outcome = if outcome == "complete" { + match parsed.semantic { + Some(semantic) => semantic.label(), + None => outcome, + } + } else { + outcome + }; + let usage_state = match (parsed.usage.prompt_tokens, parsed.usage.completion_tokens) { + (Some(_), Some(_)) => "present", + (None, None) => "missing", + _ => "partial", + }; + let accepted = if matches!(outcome, "configuration_error" | "authentication_error") { + Some(false) + } else { + self.status.map(|status| (200..300).contains(&status)) + }; + let event = json!({"kind":"round", "round":self.round, "source":"router-upstream", + "provider":self.provider, "model":self.model, "http_status":self.status, + "accepted":accepted, + "outcome":outcome, "usage":parsed.usage, "usage_state":usage_state, + "finish_reason":parsed.finish, "partial_observation":parsed.partial, + "ms":self.started.elapsed().as_millis().min(u64::MAX as u128) as u64, + "tool_calls_observed":parsed.tools.len(), + }); + self.telemetry + .finish(&self.scope, self.round, event, parsed.tools); + } +} +impl Drop for Observation { + fn drop(&mut self) { + if !self.finished { + self.fail("cancelled"); + } + } +} + +pub fn wrap_stream( + inner: BoxStream<'static, Result>, + observation: Option, + is_sse: bool, +) -> BoxStream<'static, Result> { + match observation { + Some(observation) => { + let accumulator = stream::Accumulator::new(observation.shape, is_sse); + Box::pin(ObservedStream { + inner, + observation, + accumulator, + }) + } + None => inner, + } +} + +struct ObservedStream { + inner: BoxStream<'static, Result>, + observation: Observation, + accumulator: stream::Accumulator, +} +impl Drop for ObservedStream { + fn drop(&mut self) { + if !self.observation.finished { + let parsed = self.accumulator.finish(); + self.observation.finish_parsed("cancelled", parsed); + } + } +} +impl Stream for ObservedStream { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + match self.inner.as_mut().poll_next(context) { + Poll::Ready(Some(Ok(bytes))) => { + self.accumulator.feed(&bytes); + Poll::Ready(Some(Ok(bytes))) + } + Poll::Ready(Some(Err(error))) => { + let parsed = self.accumulator.finish(); + self.observation + .finish_parsed("response_body_error", parsed); + Poll::Ready(Some(Err(error))) + } + Poll::Ready(None) => { + let parsed = self.accumulator.finish(); + let outcome = if self.observation.status.is_some_and(|status| status >= 400) { + "http_error" + } else if self.accumulator.failed { + "upstream_error" + } else if self.accumulator.complete() { + "complete" + } else { + "incomplete" + }; + self.observation.finish_parsed(outcome, parsed); + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } +} diff --git a/inference-router/src/task_telemetry/observe/stream.rs b/inference-router/src/task_telemetry/observe/stream.rs new file mode 100644 index 000000000..9433309fb --- /dev/null +++ b/inference-router/src/task_telemetry/observe/stream.rs @@ -0,0 +1,222 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::task_telemetry::parse::{self, Parsed, Shape, add_tool, merge_usage}; +use serde_json::Value; +use std::collections::BTreeMap; + +const FRAME_LIMIT: usize = 64 * 1024; +pub(super) struct Accumulator { + shape: Shape, + sse: bool, + buffer: Vec, + discarding: bool, + tail: [u8; 4], + parsed: Parsed, + terminal: bool, + pub failed: bool, + incomplete: bool, + tools: BTreeMap<(u64, u64), (String, String)>, +} +impl Accumulator { + pub(super) fn new(shape: Shape, sse: bool) -> Self { + Self { + shape, + sse, + buffer: Vec::new(), + discarding: false, + tail: [0; 4], + parsed: Parsed::default(), + terminal: false, + failed: false, + incomplete: false, + tools: BTreeMap::new(), + } + } + pub(super) fn feed(&mut self, bytes: &[u8]) { + if !self.sse { + if self.buffer.len().saturating_add(bytes.len()) <= FRAME_LIMIT && !self.discarding { + self.buffer.extend_from_slice(bytes); + } else { + self.buffer.clear(); + self.discarding = true; + self.parsed.partial = true; + } + return; + } + for byte in bytes { + self.tail.rotate_left(1); + self.tail[3] = *byte; + if !self.discarding { + if self.buffer.len() < FRAME_LIMIT { + self.buffer.push(*byte); + } else { + self.buffer.clear(); + self.discarding = true; + self.parsed.partial = true; + } + } + if self.tail[2..] == *b"\n\n" || self.tail == *b"\r\n\r\n" { + if !self.discarding { + let frame = std::mem::take(&mut self.buffer); + self.frame(&frame); + } + self.discarding = false; + } + } + } + fn frame(&mut self, bytes: &[u8]) { + let Ok(frame) = std::str::from_utf8(bytes) else { + self.parsed.partial = true; + return; + }; + if frame.lines().any(|line| { + line.strip_prefix("event:") + .is_some_and(|name| name.trim() == "error") + }) { + self.failed = true; + self.terminal = true; + self.parsed.semantic = Some(parse::SemanticOutcome::Failed); + } + let data = frame + .lines() + .filter_map(|line| line.strip_prefix("data:").map(str::trim)) + .collect::>() + .join("\n"); + if data.is_empty() { + return; + } + if data == "[DONE]" { + self.terminal = true; + return; + } + let Ok(value) = serde_json::from_str::(&data) else { + self.parsed.partial = true; + return; + }; + if let Some(usage) = value.get("usage") { + merge_usage(&mut self.parsed.usage, usage); + } + match self.shape { + Shape::OpenAi => { + if parse::semantic_outcome(&value, self.shape).is_some() { + self.failed = true; + self.terminal = true; + self.parsed.semantic = Some(parse::SemanticOutcome::Failed); + return; + } + if let Some(choices) = value["choices"].as_array() { + for choice in choices.iter().take(parse::MAX_TOOLS) { + if let Some(reason) = choice["finish_reason"] + .as_str() + .and_then(|value| parse::identifier(value, 64)) + { + self.parsed.finish = Some(reason); + } + if let Some(calls) = choice["delta"]["tool_calls"].as_array() { + for call in calls.iter().take(parse::MAX_TOOLS) { + let key = ( + choice["index"].as_u64().unwrap_or(0), + call["index"].as_u64().unwrap_or(0), + ); + if !self.tools.contains_key(&key) + && self.tools.len() >= parse::MAX_TOOLS + { + self.parsed.partial = true; + continue; + } + let tool = self.tools.entry(key).or_default(); + for (slot, value) in [ + (&mut tool.0, &call["id"]), + (&mut tool.1, &call["function"]["name"]), + ] { + if let Some(value) = value.as_str() { + if slot.as_str() == value { + continue; + } + if slot.len().saturating_add(value.len()) <= 128 { + slot.push_str(value); + } else { + self.parsed.partial = true; + } + } + } + } + } + } + } + } + Shape::Anthropic => match value["type"].as_str() { + Some("message_start") => { + if let Some(usage) = value["message"].get("usage") { + merge_usage(&mut self.parsed.usage, usage); + } + } + Some("content_block_start") if value["content_block"]["type"] == "tool_use" => { + add_tool( + &mut self.parsed, + &value["content_block"]["id"], + &value["content_block"]["name"], + ) + } + Some("message_delta") => { + self.parsed.finish = value["delta"]["stop_reason"] + .as_str() + .and_then(|value| parse::identifier(value, 64)); + } + Some("message_stop") => self.terminal = true, + Some("error") => { + self.terminal = true; + self.failed = true; + } + _ => {} + }, + Shape::Responses => match value["type"].as_str() { + Some("response.completed" | "response.failed" | "response.incomplete") => { + self.terminal = true; + self.failed |= value["type"] == "response.failed"; + self.incomplete |= value["type"] == "response.incomplete"; + let parsed = parse::response(&value["response"], Shape::Responses); + self.failed |= parsed.semantic == Some(parse::SemanticOutcome::Failed); + self.incomplete |= parsed.semantic == Some(parse::SemanticOutcome::Incomplete); + self.parsed.usage = parsed.usage; + self.parsed.finish = parsed.finish; + self.parsed.tools = parsed.tools; + self.parsed.semantic = parsed.semantic; + self.parsed.partial |= parsed.partial; + } + Some("error") => { + self.terminal = true; + self.failed = true; + } + _ => {} + }, + } + } + pub(super) fn complete(&self) -> bool { + self.terminal && !self.failed && !self.incomplete + } + pub(super) fn finish(&mut self) -> Parsed { + if !self.sse && !self.discarding { + match serde_json::from_slice::(&self.buffer) { + Ok(value) => { + self.parsed = parse::response(&value, self.shape); + self.terminal = true; + } + Err(_) => self.parsed.partial = true, + } + } + if !self.terminal { + self.parsed.partial = true; + } + if self.failed { + self.parsed.semantic = Some(parse::SemanticOutcome::Failed); + } else if self.incomplete { + self.parsed.semantic = Some(parse::SemanticOutcome::Incomplete); + } + for (_, (id, name)) in std::mem::take(&mut self.tools) { + add_tool(&mut self.parsed, &Value::String(id), &Value::String(name)); + } + std::mem::take(&mut self.parsed) + } +} diff --git a/inference-router/src/task_telemetry/parse.rs b/inference-router/src/task_telemetry/parse.rs new file mode 100644 index 000000000..bbc854c45 --- /dev/null +++ b/inference-router/src/task_telemetry/parse.rs @@ -0,0 +1,219 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use serde::Serialize; +use serde_json::Value; + +pub const MAX_BODY: usize = 64 * 1024; +pub const MAX_TOOLS: usize = 32; + +#[derive(Clone, Copy, Debug)] +pub enum Shape { + OpenAi, + Responses, + Anthropic, +} +impl Shape { + pub fn for_path(path: &str) -> Option { + match path.trim_matches('/').trim_start_matches("v1/") { + "chat/completions" => Some(Self::OpenAi), + "responses" => Some(Self::Responses), + "messages" => Some(Self::Anthropic), + _ => None, + } + } +} + +pub fn identifier(value: &str, max: usize) -> Option { + crate::access_request::scope::identifier(value, max).then(|| value.to_string()) +} +#[derive(Default, Clone, Serialize)] +pub struct Usage { + pub prompt_tokens: Option, + pub completion_tokens: Option, + pub total_tokens: Option, + pub cached_tokens: Option, +} +#[derive(Default)] +pub struct Parsed { + pub usage: Usage, + pub finish: Option, + pub tools: Vec<(String, String)>, + pub partial: bool, + pub semantic: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SemanticOutcome { + Failed, + Incomplete, +} + +impl SemanticOutcome { + pub fn label(self) -> &'static str { + match self { + Self::Failed => "upstream_error", + Self::Incomplete => "incomplete", + } + } +} + +pub fn semantic_outcome(value: &Value, shape: Shape) -> Option { + if matches!(shape, Shape::Responses) { + match value.get("status").and_then(Value::as_str) { + Some("failed") => return Some(SemanticOutcome::Failed), + Some("incomplete") => return Some(SemanticOutcome::Incomplete), + _ => {} + } + } + (value.get("error").is_some_and(|error| !error.is_null()) + || value.get("type").and_then(Value::as_str) == Some("error")) + .then_some(SemanticOutcome::Failed) +} + +pub fn merge_usage(target: &mut Usage, usage: &Value) { + let get = |key: &str| usage.get(key).and_then(Value::as_u64); + if let Some(value) = get("prompt_tokens").or_else(|| get("input_tokens")) { + target.prompt_tokens = Some(value); + } + if let Some(value) = get("completion_tokens").or_else(|| get("output_tokens")) { + target.completion_tokens = Some(value); + } + if let Some(value) = get("total_tokens") { + target.total_tokens = Some(value); + } + if let Some(value) = get("cache_read_input_tokens").or_else(|| { + usage + .get("prompt_tokens_details") + .or_else(|| usage.get("input_tokens_details")) + .and_then(|value| value.get("cached_tokens")) + .and_then(Value::as_u64) + }) { + target.cached_tokens = Some(value); + } +} + +pub fn add_tool(parsed: &mut Parsed, id: &Value, name: &Value) { + let Some(name) = name.as_str().and_then(|value| identifier(value, 128)) else { + parsed.partial = true; + return; + }; + let id = id + .as_str() + .and_then(|value| identifier(value, 128)) + .unwrap_or_default(); + if parsed.tools.len() >= MAX_TOOLS { + parsed.partial = true; + return; + } + parsed.tools.push((id, name)); +} + +pub fn response(value: &Value, shape: Shape) -> Parsed { + let mut parsed = Parsed { + semantic: semantic_outcome(value, shape), + ..Default::default() + }; + if let Some(usage) = value.get("usage") { + merge_usage(&mut parsed.usage, usage); + } + match shape { + Shape::OpenAi => { + if let Some(choices) = value.get("choices").and_then(Value::as_array) { + parsed.partial |= choices.len() > MAX_TOOLS; + for choice in choices.iter().take(MAX_TOOLS) { + parsed.finish = choice + .get("finish_reason") + .and_then(Value::as_str) + .and_then(|value| identifier(value, 64)); + if let Some(tools) = choice + .pointer("/message/tool_calls") + .and_then(Value::as_array) + { + for tool in tools.iter().take(MAX_TOOLS + 1) { + add_tool(&mut parsed, &tool["id"], &tool["function"]["name"]); + } + } + } + } + } + Shape::Responses => { + parsed.finish = value + .get("status") + .and_then(Value::as_str) + .and_then(|value| identifier(value, 64)); + if let Some(items) = value.get("output").and_then(Value::as_array) { + parsed.partial |= items.len() > MAX_TOOLS; + for item in items.iter().take(MAX_TOOLS + 1) { + if item["type"] == "function_call" { + add_tool(&mut parsed, &item["call_id"], &item["name"]); + } + } + } + } + Shape::Anthropic => { + parsed.finish = value + .get("stop_reason") + .and_then(Value::as_str) + .and_then(|value| identifier(value, 64)); + if let Some(items) = value.get("content").and_then(Value::as_array) { + parsed.partial |= items.len() > MAX_TOOLS; + for item in items.iter().take(MAX_TOOLS + 1) { + if item["type"] == "tool_use" { + add_tool(&mut parsed, &item["id"], &item["name"]); + } + } + } + } + } + parsed +} + +pub fn request_results(body: &[u8], shape: Shape) -> Vec<(String, Option)> { + if body.len() > MAX_BODY { + return Vec::new(); + } + let Ok(value) = serde_json::from_slice::(body) else { + return Vec::new(); + }; + let mut results = Vec::new(); + let mut add = |id: &Value, ok: Option| { + if results.len() < MAX_TOOLS + && let Some(id) = id.as_str().and_then(|value| identifier(value, 128)) + { + results.push((id, ok)); + } + }; + if matches!(shape, Shape::Responses) { + if let Some(items) = value["input"].as_array() { + for item in items.iter().rev().take(128) { + if item["type"] == "function_call_output" { + add( + &item["call_id"], + item["is_error"].as_bool().map(|error| !error), + ); + } + } + } + } else if let Some(messages) = value["messages"].as_array() { + for message in messages.iter().rev().take(128) { + if message["role"] == "tool" { + add( + &message["tool_call_id"], + message["is_error"].as_bool().map(|error| !error), + ); + } + if let Some(content) = message["content"].as_array() { + for item in content.iter().take(MAX_TOOLS) { + if item["type"] == "tool_result" { + add( + &item["tool_use_id"], + item["is_error"].as_bool().map(|error| !error), + ); + } + } + } + } + } + results +} diff --git a/inference-router/tests/agt_governance_integration.rs b/inference-router/tests/agt_governance_integration.rs index 76c4b15e0..016ec6386 100644 --- a/inference-router/tests/agt_governance_integration.rs +++ b/inference-router/tests/agt_governance_integration.rs @@ -38,6 +38,7 @@ fn test_state(sandbox: &str, admin_token: Option<&str>) -> AppState { let policy_status = Arc::new(kars_inference_router::policy_status::PolicyStatusRegistry::new()); let governance = Arc::new(Governance::new_with_status(sandbox, policy_status.clone())); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/anthropic_buffered_guardrail.rs b/inference-router/tests/anthropic_buffered_guardrail.rs index 750fe6909..2151fddc9 100644 --- a/inference-router/tests/anthropic_buffered_guardrail.rs +++ b/inference-router/tests/anthropic_buffered_guardrail.rs @@ -49,6 +49,7 @@ fn test_state(anthropic_endpoint: String, moderation_endpoint: String) -> AppSta policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/chat_output_guardrail_nonjson.rs b/inference-router/tests/chat_output_guardrail_nonjson.rs index 99db094e3..ad80ae9cf 100644 --- a/inference-router/tests/chat_output_guardrail_nonjson.rs +++ b/inference-router/tests/chat_output_guardrail_nonjson.rs @@ -47,6 +47,7 @@ fn test_state(ollama_endpoint: String, moderation_endpoint: String) -> AppState policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/common/governed_services.rs b/inference-router/tests/common/governed_services.rs new file mode 100644 index 000000000..a59804578 --- /dev/null +++ b/inference-router/tests/common/governed_services.rs @@ -0,0 +1,196 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#![allow(dead_code)] + +use axum::{ + Router, + body::Body, + extract::ConnectInfo, + http::{Request, StatusCode}, +}; +use kars_inference_router::{ + access_request::{Identity, scope::ResourceIdentity}, + auth::WorkloadIdentityAuth, + blocklist::Blocklist, + budget::TokenBudgetTracker, + config::{Config, RegistryMode}, + egress_blocked::BlockedBuffer, + governance::Governance, + governed_services::GovernedServices, + handoff::{DrainState, HandoffSession, HandoffTokenStore, PendingHandoffStore}, + mesh::{MeshInbox, MeshMetrics}, + policy_status::PolicyStatusRegistry, + routes::{self, AppState}, +}; +use serde_json::{Value, json}; +use std::sync::Arc; +use tower::ServiceExt; + +pub const CONTROL: &str = "operator-only-service-control-credential-0123456789abcdef"; +pub const AGENT_ADMIN: &str = "legacy-agent-admin-credential-0123456789abcdef0123456789"; + +pub fn state(workspace: &str, uid: &str) -> AppState { + let registry = Arc::new(PolicyStatusRegistry::new()); + let governance = Arc::new(Governance::new_with_status("test", registry.clone())); + let services = Arc::new(GovernedServices::new( + Identity { + sandbox: ResourceIdentity { + namespace: workspace.into(), + name: "test".into(), + uid: uid.into(), + }, + namespace_uid: format!("namespace-{uid}"), + task: Some(ResourceIdentity { + namespace: workspace.into(), + name: "task".into(), + uid: format!("task-{uid}"), + }), + task_authorization: Some(format!("sha256:{}", "a".repeat(64))), + task_generation: Some(1), + managed: true, + }, + Some(CONTROL.into()), + )); + let blocked = Arc::new(BlockedBuffer::with_defaults()); + blocked.bind_services(&services); + AppState { + services, + auth: Arc::new(WorkloadIdentityAuth::new()), + copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), + client: reqwest::Client::new(), + config: Arc::new(Config { + port: 0, + foundry_endpoint: None, + foundry_project_endpoint: None, + azure_openai_endpoint: None, + default_model: "model".into(), + content_safety_enabled: false, + prompt_shields_enabled: false, + content_safety_endpoint: None, + token_budget_daily: 1000, + token_budget_per_request: 100, + registry_mode: RegistryMode::Local, + registry_url: None, + provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: None, + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: None, + openai_moderation_model: "omni-moderation-latest".into(), + providers: Default::default(), + }), + budget: TokenBudgetTracker::new(1000, 100), + policy_provider: governance.clone(), + audit_sink: governance.clone(), + signing_provider: governance.clone(), + governance, + blocklist: Blocklist::disabled(), + blocked_egress: blocked, + sandbox_name: Arc::new("test".into()), + inbox: Arc::new(MeshInbox::new()), + mesh_metrics: Arc::new(MeshMetrics::new()), + model_override: Default::default(), + admin_token: Some(Arc::new(AGENT_ADMIN.into())), + responses_only_models: Default::default(), + unavailable_models: Default::default(), + handoff_tokens: HandoffTokenStore::new(), + handoff_session: HandoffSession::new(), + drain_state: DrainState::new(), + pending_handoff: PendingHandoffStore::new(), + policy_status: registry, + inference_policy: kars_inference_router::inference_policy_loader::empty_handle(), + memory_binding: kars_inference_router::memory_binding_loader::empty_handle(), + egress_allowlist: kars_inference_router::egress_allowlist_loader::empty_handle(), + deployment_health: Arc::new( + kars_inference_router::deployment_health::DeploymentHealthRegistry::new(), + ), + } +} +pub fn app(state: AppState) -> Router { + routes::governed_service_routes(state.clone()) + .merge(routes::egress_routes()) + .merge(routes::sensitive_agt_routes()) + .with_state(state) +} +pub fn request( + method: &str, + path: &str, + body: Value, + token: Option<&str>, + scope: Option<&str>, + local: bool, +) -> Request { + let mut request = Request::builder() + .method(method) + .uri(path) + .header("content-type", "application/json"); + if let Some(token) = token { + request = request.header("authorization", format!("Bearer {token}")); + } + if let Some(scope) = scope { + request = request.header("x-kars-service-scope", scope); + } + let mut request = request.body(Body::from(body.to_string())).unwrap(); + request + .extensions_mut() + .insert(ConnectInfo::( + if local { + "127.0.0.1:32100" + } else { + "192.0.2.10:32100" + } + .parse() + .unwrap(), + )); + request +} +pub async fn send(app: &Router, request: Request) -> (StatusCode, Value) { + let response = app.clone().oneshot(request).await.unwrap(); + let status = response.status(); + let body = axum::body::to_bytes(response.into_body(), 2 * 1024 * 1024) + .await + .unwrap(); + ( + status, + serde_json::from_slice(&body) + .unwrap_or_else(|_| json!({"text":String::from_utf8_lossy(&body)})), + ) +} +pub async fn queue( + app: &Router, + scope: &str, + kind: &str, + target: &str, + port: Option, +) -> Value { + let (status, body) = send( + app, + request( + "POST", + "/v1/access-request", + json!({"scope_id":scope,"kind":kind,"target":target,"reason":"needed", "port":port}), + None, + None, + true, + ), + ) + .await; + assert_eq!(status, StatusCode::ACCEPTED, "{body}"); + body["request"].clone() +} +pub async fn decision(app: &Router, scope: &str, id: &str, verdict: &str) -> (StatusCode, Value) { + send( + app, + request( + "POST", + "/internal/access-requests/decision", + json!({"scope_id":scope,"request_id":id,"verdict":verdict}), + Some(CONTROL), + None, + true, + ), + ) + .await +} diff --git a/inference-router/tests/egress_blocked_endpoint.rs b/inference-router/tests/egress_blocked_endpoint.rs index 843a17772..6b8c8bfaf 100644 --- a/inference-router/tests/egress_blocked_endpoint.rs +++ b/inference-router/tests/egress_blocked_endpoint.rs @@ -36,6 +36,7 @@ fn test_state() -> AppState { policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/failover_walk.rs b/inference-router/tests/failover_walk.rs index b499862c7..41fe902a5 100644 --- a/inference-router/tests/failover_walk.rs +++ b/inference-router/tests/failover_walk.rs @@ -131,6 +131,7 @@ async fn primary_503_falls_through_to_fallback_200() { let health = Arc::new(DeploymentHealthRegistry::new()); let upstream = UpstreamConfig { + telemetry: None, endpoint: base, deployment: "fallback-up".into(), sandbox_name: "sbx".into(), @@ -198,6 +199,7 @@ async fn unhealthy_primary_is_skipped_in_second_pass() { assert!(!health.is_healthy("Foundry::primary-down")); let upstream = UpstreamConfig { + telemetry: None, endpoint: base, deployment: "fallback-up".into(), sandbox_name: "sbx".into(), @@ -257,6 +259,7 @@ async fn all_unhealthy_still_punches_primary_for_last_resort() { assert!(!health.is_healthy("Foundry::fallback-up")); let upstream = UpstreamConfig { + telemetry: None, endpoint: base, deployment: "primary-down".into(), sandbox_name: "sbx".into(), diff --git a/inference-router/tests/foundry_route_guard.rs b/inference-router/tests/foundry_route_guard.rs index 76d5e3f06..b5913fecb 100644 --- a/inference-router/tests/foundry_route_guard.rs +++ b/inference-router/tests/foundry_route_guard.rs @@ -54,6 +54,7 @@ fn test_state() -> AppState { policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/governed_access_services.rs b/inference-router/tests/governed_access_services.rs new file mode 100644 index 000000000..5c9bec61b --- /dev/null +++ b/inference-router/tests/governed_access_services.rs @@ -0,0 +1,533 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#[path = "common/governed_services.rs"] +mod support; +use axum::http::StatusCode; +use kars_inference_router::access_request::{ + AccessRequestBuffer, Error, Identity, Request, Status, +}; +use serde_json::json; +use std::{sync::Arc, time::Duration}; +use support::*; +use tower::ServiceExt; + +#[tokio::test] +async fn privileged_http_endpoints_reject_agent_admin_even_on_loopback() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state); + for token in [None, Some(AGENT_ADMIN)] { + for path in [ + "/internal/access-requests/reset", + "/internal/access-requests/decision", + ] { + let (status, _) = send( + &app, + request("POST", path, json!({"scope_id":scope}), token, None, true), + ) + .await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + } + assert_eq!( + send( + &app, + request( + "GET", + "/internal/access-requests", + json!({}), + token, + None, + true + ) + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + } + assert_eq!( + send( + &app, + request("GET", "/v1/access-requests", json!({}), None, None, false) + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + send( + &app, + request( + "GET", + "/internal/access-requests", + json!({}), + Some(CONTROL), + None, + false + ) + ) + .await + .0, + StatusCode::OK + ); +} + +#[tokio::test] +async fn reset_fences_old_requests_and_preserves_uid_qualified_identity() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + state.budget.record_usage("test", 17).await; + let app = app(state.clone()); + let entry = queue(&app, &scope, "egress", "example.test", Some(443)).await; + let (status, reset) = send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope,"assignment_id":"assignment-2"}), + Some(CONTROL), + None, + true, + ), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!(reset["scope"]["identity"]["sandbox"]["uid"], "uid-a"); + assert_eq!(reset["scope"]["identity"]["task"]["uid"], "task-uid-a"); + assert_eq!(reset["enforcement_changed"], false); + assert_ne!(reset["scope"]["id"], scope); + assert_eq!( + decision( + &app, + &scope, + entry["request_id"].as_str().unwrap(), + "approved" + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + send( + &app, + request( + "POST", + "/v1/access-request", + json!({"scope_id":scope,"kind":"egress","target":"example.test"}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope}), + Some(CONTROL), + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!(state.budget.get_usage("test").await.0, 17); +} + +#[tokio::test] +async fn request_identity_port_tier_and_terminal_decisions_cannot_be_rebound() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state); + let first = queue(&app, &scope, "egress", "example.test", Some(443)).await; + let id = first["request_id"].as_str().unwrap(); + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + assert_eq!( + decision(&app, &scope, id, "denied").await.0, + StatusCode::CONFLICT + ); + let other = queue(&app, &scope, "egress", "example.test", Some(80)).await; + assert_ne!(first["request_id"], other["request_id"]); + assert_eq!(other["status"], "pending"); + let (status,body)=send(&app,request("POST","/v1/access-request", + json!({"scope_id":scope,"kind":"egress","target":"example.test","reason":"changed","port":443}),None,None,true)).await; + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!(body["request"]["reason"], "needed"); + assert_eq!(body["request"]["status"], "approved"); + for (invalid, expected) in [ + ( + json!({"scope_id":scope,"kind":"tier","tier":6}), + StatusCode::BAD_REQUEST, + ), + ( + json!({"scope_id":scope,"kind":"egress","target":"https://example.test/token"}), + StatusCode::BAD_REQUEST, + ), + ( + json!({"scope_id":scope,"kind":"tool","target":"test","sandbox_uid":"other"}), + StatusCode::UNPROCESSABLE_ENTITY, + ), + ] { + assert_eq!( + send( + &app, + request("POST", "/v1/access-request", invalid, None, None, true) + ) + .await + .0, + expected + ); + } +} + +#[tokio::test] +async fn same_named_sandboxes_in_other_workspaces_cannot_share_request_scope() { + let a = state("workspace-a", "uid-a"); + let b = state("workspace-b", "uid-b"); + let scope = a.services.requests.scope().unwrap().id; + let app_b = app(b); + assert_eq!( + send( + &app_b, + request( + "POST", + "/v1/access-request", + json!({"scope_id":scope,"kind":"tool","target":"fetch"}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); +} + +#[tokio::test] +async fn decision_wait_wakes_and_agent_cancel_does_not_rewrite_an_operator_decision() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state.clone()); + let entry = queue(&app, &scope, "tool", "fetch", None).await; + let id = entry["request_id"].as_str().unwrap(); + let waiting = app.clone().oneshot(request( + "GET", + &format!("/v1/access-requests/{id}/wait?scope_id={scope}&timeout_ms=1000"), + json!({}), + None, + None, + true, + )); + let wait = tokio::spawn(waiting); + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + assert_eq!(wait.await.unwrap().unwrap().status(), StatusCode::OK); + assert_eq!( + send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::OK + ); + let cancelled = state.services.requests.entry(&scope, id).unwrap(); + assert_eq!(cancelled.status, Status::Cancelled); + assert_eq!(cancelled.decision, Some(Status::Approved)); +} + +#[tokio::test] +async fn storage_rate_and_expiry_limits_fail_closed() { + let buffer = AccessRequestBuffer::with_limits( + Identity::standalone("test"), + 1, + 128, + Duration::from_millis(10), + ); + let scope = buffer.scope().unwrap().id; + let request = |target: &str| Request { + scope_id: scope.clone(), + kind: "tool".into(), + target: target.into(), + reason: "reason".into(), + tier: None, + port: None, + }; + let (first, _) = buffer.record(request("first")).unwrap(); + assert!(matches!(buffer.record(request("second")), Err(Error::Full))); + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(matches!( + buffer.transition(&scope, &first.request_id, Status::Approved), + Err(Error::Expired) + )); + let rate = AccessRequestBuffer::with_limits( + Identity::standalone("test"), + 64, + 1, + Duration::from_secs(30), + ); + let scope = rate.scope().unwrap().id; + let request = Request { + scope_id: scope, + kind: "tool".into(), + target: "first".into(), + reason: String::new(), + tier: None, + port: None, + }; + rate.record(request.clone()).unwrap(); + assert!(matches!(rate.record(request), Err(Error::RateLimited))); +} + +#[tokio::test] +async fn dropping_waiter_and_reset_release_bounded_wait_capacity() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state.clone()); + let entry = queue(&app, &scope, "tool", "fetch", None).await; + let id = entry["request_id"].as_str().unwrap().to_string(); + let services = state.services.clone(); + let scope_copy = scope.clone(); + let task = tokio::spawn(async move { + services + .wait_for_decision(&scope_copy, &id, Duration::from_secs(30)) + .await + }); + for _ in 0..30 { + if state.services.wait_slots() == 15 { + break; + } + tokio::task::yield_now().await; + } + assert_eq!(state.services.wait_slots(), 15); + task.abort(); + let _ = task.await; + assert_eq!(state.services.wait_slots(), 16); + let scope_copy = scope.clone(); + let id = entry["request_id"].as_str().unwrap().to_string(); + let services = state.services.clone(); + let task = tokio::spawn(async move { + services + .wait_for_decision(&scope_copy, &id, Duration::from_secs(30)) + .await + }); + state.services.reset(&scope, Some("next".into())).unwrap(); + assert!(matches!(task.await.unwrap(), Err(Error::StaleScope))); + assert_eq!(state.services.wait_slots(), 16); +} + +#[tokio::test] +async fn missing_control_configuration_never_falls_back_to_shared_agent_admin() { + let mut state = state("workspace-a", "uid-a"); + state.services = Arc::new( + kars_inference_router::governed_services::GovernedServices::new( + Identity::standalone("test"), + None, + ), + ); + let app = app(state); + assert_eq!( + send( + &app, + request( + "GET", + "/internal/access-requests", + json!({}), + Some(AGENT_ADMIN), + None, + true + ) + ) + .await + .0, + StatusCode::SERVICE_UNAVAILABLE + ); +} + +#[tokio::test] +async fn http_payload_rate_and_operator_origin_limits_are_enforced() { + let mut state = state("workspace-a", "uid-a"); + let identity = state.services.requests.scope().unwrap().identity; + let mut services = kars_inference_router::governed_services::GovernedServices::new( + identity, + Some(CONTROL.into()), + ); + services.allow_ips = Some(vec!["192.0.2.11".parse().unwrap()]); + state.services = Arc::new(services); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state.clone()); + assert_eq!( + send( + &app, + request( + "GET", + "/internal/access-requests", + json!({}), + Some(CONTROL), + None, + false + ) + ) + .await + .0, + StatusCode::FORBIDDEN + ); + assert_eq!( + send( + &app, + request( + "GET", + "/telemetry/cursor", + json!({}), + Some(CONTROL), + None, + false + ) + ) + .await + .0, + StatusCode::FORBIDDEN + ); + assert_eq!( + send( + &app, + request( + "POST", + "/v1/access-request", + json!({ + "scope_id":scope,"kind":"tool","target":"fetch","reason":"x".repeat(9000) + }), + None, + None, + true + ) + ) + .await + .0, + StatusCode::PAYLOAD_TOO_LARGE + ); + for _ in 0..32 { + assert_eq!( + send( + &app, + request( + "POST", + "/v1/access-request", + json!({ + "scope_id":scope,"kind":"tool","target":"fetch" + }), + None, + None, + true + ) + ) + .await + .0, + StatusCode::ACCEPTED + ); + } + assert_eq!( + send( + &app, + request( + "POST", + "/v1/access-request", + json!({ + "scope_id":scope,"kind":"tool","target":"fetch" + }), + None, + None, + true + ) + ) + .await + .0, + StatusCode::TOO_MANY_REQUESTS + ); +} + +#[tokio::test] +async fn unknown_decision_cannot_synthesize_an_approved_request() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state.clone()); + assert_eq!( + decision(&app, &scope, "unknown-id", "approved").await.0, + StatusCode::NOT_FOUND + ); + assert!(state.services.requests.snapshot(&scope).unwrap().is_empty()); +} + +#[tokio::test] +async fn live_http_control_auth_and_reset_use_the_same_production_router() { + let state = state("workspace-a", "uid-a"); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let shutdown = tokio_util::sync::CancellationToken::new(); + let stop = shutdown.clone(); + let server = tokio::spawn(async move { + axum::serve( + listener, + app(state).into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(stop.cancelled_owned()) + .await + .unwrap(); + }); + let client = reqwest::Client::builder().no_proxy().build().unwrap(); + let base = format!("http://{address}"); + let info = client + .get(format!("{base}/v1/access-requests")) + .send() + .await + .unwrap(); + assert_eq!(info.status(), StatusCode::OK); + let scope = info.json::().await.unwrap()["scope"]["id"] + .as_str() + .unwrap() + .to_string(); + let response = client + .post(format!("{base}/internal/access-requests/reset")) + .bearer_auth(AGENT_ADMIN) + .json(&json!({"scope_id":scope})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + let response = client + .post(format!("{base}/internal/access-requests/reset")) + .bearer_auth(CONTROL) + .json(&json!({"scope_id":scope,"assignment_id":"live-test"})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_ne!( + response.json::().await.unwrap()["scope"]["id"], + scope + ); + shutdown.cancel(); + server.await.unwrap(); +} diff --git a/inference-router/tests/governed_egress_wait.rs b/inference-router/tests/governed_egress_wait.rs new file mode 100644 index 000000000..5d9fda648 --- /dev/null +++ b/inference-router/tests/governed_egress_wait.rs @@ -0,0 +1,409 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#[path = "common/governed_services.rs"] +mod support; +use axum::http::StatusCode; +use kars_inference_router::{blocklist::Blocklist, routes::AppState}; +use serde_json::{Value, json}; +use std::time::Duration; +use support::*; +use tower::ServiceExt; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + +async fn setup() -> (AppState, axum::Router, MockServer, String, String) { + let upstream = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(200).set_body_string("response-secret-marker")) + .mount(&upstream) + .await; + let address = reqwest::Url::parse(&upstream.uri()) + .unwrap() + .socket_addrs(|| None) + .unwrap()[0]; + let mut state = state("workspace", "uid"); + state.blocklist = Blocklist::new(None).await; + state.client = reqwest::Client::builder() + .no_proxy() + .redirect(reqwest::redirect::Policy::none()) + .resolve("approved.example", address) + .build() + .unwrap(); + let scope = state.services.requests.scope().unwrap().id; + let url = format!( + "http://approved.example:{}/data?key=request-secret-marker", + address.port() + ); + (state.clone(), app(state), upstream, scope, url) +} + +async fn pending(app: &axum::Router) -> Value { + for _ in 0..100 { + let (_, body) = send( + app, + request( + "GET", + "/internal/access-requests", + json!({}), + Some(CONTROL), + None, + true, + ), + ) + .await; + if let Some(entry) = body["entries"] + .as_array() + .and_then(|entries| entries.first()) + { + return entry.clone(); + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + panic!("request was not queued"); +} + +#[tokio::test] +async fn approved_decision_alone_cannot_widen_egress_but_signed_policy_wakes_the_request() { + let (state, app, upstream, scope, url) = setup().await; + let waiting = tokio::spawn(app.clone().oneshot(request( + "POST", + "/egress/fetch", + json!({"url":url,"scope_id":scope,"wait_for_approval_ms":2000, + "headers":{"Authorization":"Bearer header-secret-marker"}}), + None, + None, + true, + ))); + let entry = pending(&app).await; + assert_eq!( + decision( + &app, + &scope, + entry["request_id"].as_str().unwrap(), + "approved" + ) + .await + .0, + StatusCode::OK + ); + tokio::time::sleep(Duration::from_millis(30)).await; + assert!(!waiting.is_finished()); + assert!(upstream.received_requests().await.unwrap().is_empty()); + // This is the primitive used by the existing signed allowlist loader; the + // access-request service has no API capable of writing it. + state + .blocklist + .replace_allowlist(vec!["approved.example".into()]) + .await; + let response = waiting.await.unwrap().unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(upstream.received_requests().await.unwrap().len(), 1); + let trace = state + .services + .telemetry + .snapshot(&scope, 0) + .unwrap() + .to_string(); + for secret in [ + "request-secret-marker", + "response-secret-marker", + "header-secret-marker", + "approved.example", + ] { + assert!(!trace.contains(secret), "{trace}"); + } + assert!(!state.blocklist.is_learn_mode()); +} + +#[tokio::test] +async fn denial_cancel_reset_and_timeout_do_not_send_unapproved_traffic() { + for action in ["deny", "cancel", "reset", "timeout"] { + let (state, app, upstream, scope, url) = setup().await; + let waiting=tokio::spawn(app.clone().oneshot(request("POST","/egress/fetch", + json!({"url":url,"scope_id":scope,"wait_for_approval_ms":if action=="timeout"{30}else{1000}}),None,None,true))); + let entry = pending(&app).await; + let id = entry["request_id"].as_str().unwrap(); + match action { + "deny" => { + assert_eq!(decision(&app, &scope, id, "denied").await.0, StatusCode::OK); + } + "cancel" => { + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + assert_eq!( + send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::OK + ); + } + "reset" => { + assert_eq!( + send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope,"assignment_id":"next"}), + Some(CONTROL), + None, + true + ) + ) + .await + .0, + StatusCode::OK + ); + } + _ => {} + } + let response = waiting.await.unwrap().unwrap(); + assert_eq!( + response.status(), + if action == "timeout" { + StatusCode::REQUEST_TIMEOUT + } else { + StatusCode::CONFLICT + } + ); + assert!(upstream.received_requests().await.unwrap().is_empty()); + assert_eq!(state.services.wait_slots(), 16); + } +} + +#[tokio::test] +async fn ordinary_denied_fetch_is_still_immediate_and_private_targets_never_enter_approval_wait() { + let (_state, app, upstream, scope, url) = setup().await; + let response = tokio::time::timeout( + Duration::from_millis(500), + send( + &app, + request( + "POST", + "/egress/fetch", + json!({"url":url}), + None, + None, + true, + ), + ), + ) + .await + .unwrap(); + assert_eq!(response.0, StatusCode::FORBIDDEN); + let response = send( + &app, + request( + "POST", + "/egress/fetch", + json!({"url":upstream.uri(),"scope_id":scope,"wait_for_approval_ms":1000}), + None, + None, + true, + ), + ) + .await; + assert_eq!(response.0, StatusCode::FORBIDDEN); + assert!(upstream.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn cancel_after_readiness_but_before_dispatch_claim_prevents_a_post() { + let (state, app, upstream, scope, url) = setup().await; + let port = reqwest::Url::parse(&url) + .unwrap() + .port_or_known_default() + .unwrap(); + let entry = queue(&app, &scope, "egress", "approved.example", Some(port)).await; + let id = entry["request_id"].as_str().unwrap(); + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + state + .blocklist + .replace_allowlist(vec!["approved.example".into()]) + .await; + state + .services + .wait_for_egress( + &state.blocklist, + &scope, + id, + &url, + "test", + Duration::from_secs(1), + ) + .await + .unwrap(); + let cancelled = send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true, + ), + ) + .await; + assert_eq!(cancelled.0, StatusCode::OK); + assert_eq!(cancelled.1["status"], "cancelled"); + assert!(matches!( + state.services.claim_egress_dispatch(&scope, Some(id)), + Err(kars_inference_router::access_request::Error::Terminal) + )); + assert!(upstream.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn cancel_and_reset_do_not_acknowledge_prevention_after_dispatch_has_started() { + let (state, app, upstream, scope, url) = setup().await; + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(200) + .set_body_string("accepted") + .set_delay(Duration::from_millis(300)), + ) + .expect(1) + .mount(&upstream) + .await; + let waiting=tokio::spawn(app.clone().oneshot(request("POST","/egress/fetch", + json!({"url":url,"method":"POST","body":"side-effect","scope_id":scope,"wait_for_approval_ms":2000}), + None,None,true))); + let entry = pending(&app).await; + let id = entry["request_id"].as_str().unwrap(); + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + state + .blocklist + .replace_allowlist(vec!["approved.example".into()]) + .await; + for _ in 0..100 { + if !upstream.received_requests().await.unwrap().is_empty() { + break; + } + tokio::time::sleep(Duration::from_millis(2)).await; + } + assert_eq!(upstream.received_requests().await.unwrap().len(), 1); + assert_eq!( + send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope,"assignment_id":"next"}), + Some(CONTROL), + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!(waiting.await.unwrap().unwrap().status(), StatusCode::OK); + let completed = state.services.requests.entry(&scope, id).unwrap(); + assert_eq!( + completed.status, + kars_inference_router::access_request::Status::DispatchClaimed + ); + assert!(!completed.dispatch_active); + assert_eq!( + send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope,"assignment_id":"next"}), + Some(CONTROL), + None, + true + ) + ) + .await + .0, + StatusCode::OK + ); +} + +#[tokio::test] +async fn wait_rejects_wrong_port_request_and_shutdown_cancels_waiters() { + let (state, app, _upstream, scope, url) = setup().await; + let entry = queue(&app, &scope, "egress", "approved.example", Some(443)).await; + let id = entry["request_id"].as_str().unwrap(); + assert!(matches!( + state + .services + .wait_for_egress( + &state.blocklist, + &scope, + id, + &url, + "test", + Duration::from_millis(20) + ) + .await, + Err(kars_inference_router::access_request::Error::Invalid) + )); + let services = state.services.clone(); + let id = id.to_string(); + let old = scope.clone(); + let waiter = tokio::spawn(async move { + services + .wait_for_decision(&old, &id, Duration::from_secs(30)) + .await + }); + state.services.shutdown.cancel(); + assert!(matches!( + waiter.await.unwrap(), + Err(kars_inference_router::access_request::Error::Unavailable) + )); +} diff --git a/inference-router/tests/governed_telemetry.rs b/inference-router/tests/governed_telemetry.rs new file mode 100644 index 000000000..6fb55ec7a --- /dev/null +++ b/inference-router/tests/governed_telemetry.rs @@ -0,0 +1,601 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#[path = "common/governed_services.rs"] +mod support; +use axum::{ + body::Body, + http::{HeaderMap, Method, Request, StatusCode}, +}; +use bytes::Bytes; +use futures::{StreamExt, TryStreamExt}; +use kars_inference_router::{ + auth::WorkloadIdentityAuth, + config::ProviderEndpoint, + proxy::{self, AuthenticationProvenance, UpstreamConfig}, + routes::{self, McpRouteState}, + task_telemetry::{TaskTelemetry, observe::wrap_stream}, +}; +use serde_json::{Value, json}; +use std::sync::Arc; +use support::*; +use tower::ServiceExt; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + +fn telemetry() -> Arc { + Arc::new(TaskTelemetry::new("scope-a".into())) +} +fn upstream(endpoint: String, telemetry: Arc) -> UpstreamConfig { + let mut target = UpstreamConfig::azure(endpoint, "model".into(), "test".into()); + target.authentication = AuthenticationProvenance::Named { + provider_id: "qualified-provider".into(), + }; + target.telemetry = Some(telemetry); + target +} +fn events(telemetry: &TaskTelemetry) -> Vec { + telemetry.snapshot("scope-a", 0).unwrap()["events"] + .as_array() + .unwrap() + .clone() +} + +#[tokio::test] +async fn buffered_http_producer_observes_usage_but_never_api_bodies_or_tool_arguments() { + let server = MockServer::start().await; + Mock::given(method("POST")).respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "choices":[{"finish_reason":"tool_calls","message":{"content":"output-secret-marker","tool_calls":[ + {"id":"call-1","function":{"name":"fetch","arguments":"argument-secret-marker"}} + ]}}],"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18} + }))).mount(&server).await; + let telemetry = telemetry(); + let target = upstream(server.uri(), telemetry.clone()); + let client = reqwest::Client::builder().no_proxy().build().unwrap(); + let (_, _, body) = proxy::forward( + &WorkloadIdentityAuth::new(), + None, + &client, + &target, + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from(r#"{"messages":[{"role":"user","content":"prompt-secret-marker"}]}"#), + ) + .await + .unwrap(); + assert!(String::from_utf8_lossy(&body).contains("output-secret-marker")); + let trace = events(&telemetry); + assert_eq!(trace[0]["usage"]["total_tokens"], 18); + assert_eq!(trace[0]["provider"], "qualified-provider"); + assert_eq!(trace[1]["kind"], "tool_proposed"); + assert!(trace[1]["ok"].is_null()); + proxy::forward(&WorkloadIdentityAuth::new(),None,&client,&target,Method::POST, + "chat/completions",&HeaderMap::new(),Bytes::from(r#"{"messages":[{"role":"tool","tool_call_id":"call-1","content":"result-secret-marker","is_error":true}]}"#)).await.unwrap(); + let trace = events(&telemetry); + assert!(trace.iter().any(|event| event["kind"] == "tool_result" + && event["ok"] == false + && event["source"] == "harness-reported")); + let serialized = serde_json::to_string(&trace).unwrap(); + for secret in [ + "prompt-secret-marker", + "output-secret-marker", + "argument-secret-marker", + "result-secret-marker", + ] { + assert!(!serialized.contains(secret)); + } +} + +#[tokio::test] +async fn buffered_http_200_responses_distinguish_failed_and_incomplete_model_outcomes() { + for (status, outcome) in [("failed", "upstream_error"), ("incomplete", "incomplete")] { + let server = MockServer::start().await; + let body = json!({ + "status":status,"error":{"message":"credential-secret-marker"}, + "incomplete_details":{"reason":"private-details-marker"}, + "usage":if status=="failed"{json!({"input_tokens":7})}else{Value::Null}, + "output":[], + }); + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(body.clone())) + .expect(1) + .mount(&server) + .await; + let telemetry = telemetry(); + let target = upstream(server.uri(), telemetry.clone()); + let (http, _, bytes) = proxy::forward( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::builder().no_proxy().build().unwrap(), + &target, + Method::POST, + "responses", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap(); + assert_eq!(http, StatusCode::OK); + assert_eq!(serde_json::from_slice::(&bytes).unwrap(), body); + let trace = events(&telemetry); + assert_eq!(trace.len(), 1); + assert_eq!(trace[0]["accepted"], true); + assert_eq!(trace[0]["outcome"], outcome); + assert_eq!(trace[0]["finish_reason"], status); + assert_eq!( + trace[0]["usage_state"], + if status == "failed" { + "partial" + } else { + "missing" + } + ); + assert!(trace[0]["usage"]["completion_tokens"].is_null()); + assert!(trace[0]["usage"]["total_tokens"].is_null()); + assert!(!trace[0].to_string().contains("secret-marker")); + assert!(!trace[0].to_string().contains("private-details-marker")); + } +} + +#[tokio::test] +async fn fragmented_openai_error_frames_cannot_be_overwritten_by_done() { + for error in [ + "data: {\"error\":{\"message\":\"credential-secret-marker🚫\"}}\n\n", + "event: error\ndata: {\"message\":\"credential-secret-marker🚫\"}\n\n", + ] { + let body = format!( + "data: {{\"choices\":[],\"usage\":{{\"prompt_tokens\":3}}}}\n\n{error}data: [DONE]\n\n" + ); + let telemetry = telemetry(); + let mut observation = telemetry + .begin("chat/completions", "provider", "model", b"{}") + .unwrap(); + observation.headers(200); + let chunks = body + .as_bytes() + .iter() + .map(|byte| Ok::<_, reqwest::Error>(Bytes::from(vec![*byte]))) + .collect::>(); + let returned = wrap_stream( + futures::stream::iter(chunks).boxed(), + Some(observation), + true, + ) + .try_collect::>() + .await + .unwrap(); + assert_eq!(returned.concat(), body.as_bytes()); + let trace = events(&telemetry); + assert_eq!(trace.len(), 1); + assert_eq!(trace[0]["accepted"], true); + assert_eq!(trace[0]["outcome"], "upstream_error"); + assert_eq!(trace[0]["usage"]["prompt_tokens"], 3); + assert!(trace[0]["usage"]["completion_tokens"].is_null()); + assert!(trace[0]["usage"]["total_tokens"].is_null()); + assert!(!trace[0].to_string().contains("credential-secret-marker")); + } +} + +#[tokio::test] +async fn responses_incomplete_stream_is_not_reported_as_a_complete_model_response() { + let telemetry = telemetry(); + let mut observation = telemetry + .begin("responses", "provider", "model", b"{}") + .unwrap(); + observation.headers(200); + let body = "data: {\"type\":\"response.incomplete\",\"response\":{\"status\":\"incomplete\",\"usage\":{\"input_tokens\":4,\"output_tokens\":2},\"output\":[]}}\n\n"; + let returned = wrap_stream( + futures::stream::once(async move { + Ok::<_, reqwest::Error>(Bytes::from_static(body.as_bytes())) + }) + .boxed(), + Some(observation), + true, + ) + .try_collect::>() + .await + .unwrap(); + assert_eq!(returned.concat(), body.as_bytes()); + let trace = events(&telemetry); + assert_eq!(trace[0]["accepted"], true); + assert_eq!(trace[0]["outcome"], "incomplete"); + assert_eq!(trace[0]["usage_state"], "present"); + assert!(trace[0]["usage"]["total_tokens"].is_null()); +} + +#[tokio::test] +async fn fragmented_streams_preserve_bytes_and_report_each_supported_usage_shape_once() { + let cases = [ + ( + "chat/completions", + concat!( + "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"🚀stream-secret-marker\"},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":4,\"completion_tokens\":5,\"total_tokens\":9}}\n\n", + "data: [DONE]\n\n" + ), + 4, + 5, + Some(9), + ), + ( + "v1/messages", + concat!( + "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":12,\"cache_read_input_tokens\":3}}}\n\n", + "data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":8},\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n", + "data: {\"type\":\"message_stop\"}\n\n" + ), + 12, + 8, + None, + ), + ( + "responses", + "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"total_tokens\":5},\"output\":[]}}\n\n", + 2, + 3, + Some(5), + ), + ]; + for (path, body, input, output, total) in cases { + let telemetry = telemetry(); + let mut observation = telemetry.begin(path, "provider", "model", b"{}").unwrap(); + observation.headers(200); + let chunks = body + .as_bytes() + .iter() + .map(|byte| Ok::<_, reqwest::Error>(Bytes::from(vec![*byte]))) + .collect::>(); + let bytes = wrap_stream( + futures::stream::iter(chunks).boxed(), + Some(observation), + true, + ) + .try_collect::>() + .await + .unwrap(); + assert_eq!(bytes.concat(), body.as_bytes()); + let events = events(&telemetry); + assert_eq!(events.len(), 1); + assert_eq!(events[0]["usage"]["prompt_tokens"], input); + assert_eq!(events[0]["usage"]["completion_tokens"], output); + assert_eq!(events[0]["usage"]["total_tokens"], json!(total)); + assert_eq!(events[0]["outcome"], "complete"); + assert!(!events[0].to_string().contains("stream-secret-marker")); + } +} + +#[tokio::test] +async fn streaming_http_producer_reports_missing_usage_and_no_replay_after_incomplete_stream() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_raw( + "data: {\"choices\":[{\"delta\":{\"content\":\"private\"}}]}\n\n", + "text/event-stream", + )) + .expect(1) + .mount(&server) + .await; + let telemetry = telemetry(); + let (_, _, stream) = proxy::forward_stream( + Arc::new(WorkloadIdentityAuth::new()), + None, + reqwest::Client::builder().no_proxy().build().unwrap(), + upstream(server.uri(), telemetry.clone()), + "chat/completions", + HeaderMap::new(), + Bytes::from(r#"{"stream":true}"#), + ) + .await + .unwrap(); + let _: Vec<_> = stream.try_collect().await.unwrap(); + let events = events(&telemetry); + assert_eq!(events.len(), 1); + assert_eq!(events[0]["outcome"], "incomplete"); + assert_eq!(events[0]["accepted"], true); + assert_eq!(events[0]["usage_state"], "missing"); + assert!(events[0]["usage"]["total_tokens"].is_null()); +} + +#[tokio::test] +async fn reset_discards_old_inflight_observations_and_ring_overflow_is_explicit() { + let telemetry = telemetry(); + let mut old = telemetry + .begin("responses", "provider", "model", b"{}") + .unwrap(); + telemetry.reset("scope-b".into()); + old.buffered( + 200, + br#"{"usage":{"input_tokens":10,"output_tokens":4,"total_tokens":14}}"#, + ); + assert!(telemetry.snapshot("scope-a", 0).is_none()); + assert_eq!( + telemetry.snapshot("scope-b", 0).unwrap()["events"], + json!([]) + ); + for _ in 0..1100 { + telemetry.record_policy("scope-b", "tool:fetch", false); + } + let snapshot = telemetry.snapshot("scope-b", 0).unwrap(); + assert_eq!(snapshot["dropped_events"], 76); + assert_eq!(snapshot["events"].as_array().unwrap().len(), 256); + assert_eq!(snapshot["has_more"], true); +} + +#[tokio::test] +async fn physical_failover_attempts_keep_actual_provider_identity_without_counting_false_successes() +{ + use kars_inference_router::{ + failover::forward_with_failover, + inference_policy_loader::{InferencePolicySnapshot, ModelPreference, ModelRef}, + }; + let first = MockServer::start().await; + let second = MockServer::start().await; + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(503) + .set_body_json(json!({"error":{"message":"secret-error-body"}})), + ) + .expect(1) + .mount(&first) + .await; + Mock::given(method("POST")).respond_with(ResponseTemplate::new(200).set_body_json(json!({"choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}))).expect(1).mount(&second).await; + let state = state("workspace", "uid"); + let mut config = Arc::try_unwrap(state.config) + .ok() + .expect("fixture config is uniquely owned"); + for (id, endpoint) in [("first", first.uri()), ("second", second.uri())] { + config.providers.insert( + id.into(), + ProviderEndpoint { + tag: id.into(), + endpoint, + api_key: None, + }, + ); + } + let telemetry = telemetry(); + let base = upstream("http://127.0.0.1:1".into(), telemetry.clone()); + let policy = InferencePolicySnapshot { + model_preference: Some(ModelPreference { + primary: ModelRef { + provider: "first".into(), + deployment: "same-model".into(), + }, + fallback: vec![ModelRef { + provider: "second".into(), + deployment: "same-model".into(), + }], + }), + ..Default::default() + }; + let result = forward_with_failover( + &state.auth, + None, + &reqwest::Client::builder().no_proxy().build().unwrap(), + &state.deployment_health, + &base, + &config, + &policy, + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap(); + assert_eq!(result.0, StatusCode::OK); + let events = events(&telemetry); + assert_eq!(events.len(), 2); + assert_eq!(events[0]["provider"], "first"); + assert_eq!(events[0]["outcome"], "http_error"); + assert_eq!(events[1]["provider"], "second"); + assert_eq!(events[1]["usage"]["total_tokens"], 3); + assert!( + !serde_json::to_string(&events) + .unwrap() + .contains("secret-error-body") + ); +} + +#[tokio::test] +async fn telemetry_http_requires_scope_and_only_correlated_native_outcomes_are_accepted() { + let state = state("workspace", "uid"); + let scope = state.services.telemetry.cursor().0; + let app = app(state.clone()); + assert_eq!( + send( + &app, + request("GET", "/telemetry/trace", json!({}), None, None, true) + ) + .await + .0, + StatusCode::CONFLICT + ); + let body = json!({"scope_id":scope,"call_id":"call-a","ok":true,"latency_ms":5}); + assert_eq!( + send( + &app, + request("POST", "/telemetry/tool", body.clone(), None, None, true) + ) + .await + .0, + StatusCode::CONFLICT + ); + let (status,evaluation)=send(&app,request("POST","/agt/evaluate", + json!({"agent_id":"test","action":"tool:fetch","context":{ + "scope_id":scope,"tool_call_id":"call-a","tool_name":"spoofed-name","args_preview":"native-secret-marker" + }}),None,None,true)).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(evaluation["allowed"], true); + assert_eq!( + send( + &app, + request("POST", "/telemetry/tool", body.clone(), None, None, true) + ) + .await + .0, + StatusCode::OK + ); + assert_eq!( + send( + &app, + request("POST", "/telemetry/tool", body, None, None, true) + ) + .await + .0, + StatusCode::CONFLICT + ); + let (status, trace) = send( + &app, + request( + "GET", + "/telemetry/trace", + json!({}), + None, + Some(&scope), + true, + ), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert!( + trace["events"] + .as_array() + .unwrap() + .iter() + .any(|event| event["source"] == "harness-reported" && event["name"] == "fetch") + ); + assert!(!trace.to_string().contains("native-secret-marker")); + assert!(!trace.to_string().contains("spoofed-name")); + assert_eq!(trace["durable"], false); +} + +#[tokio::test] +async fn actual_mcp_http_errors_are_not_mistaken_for_success_and_bodies_are_not_retained() { + let telemetry = telemetry(); + let mut state = McpRouteState::standard(); + state.task_telemetry = Some(telemetry.clone()); + let app = routes::mcp_route().with_state(state); + for (id, name) in [(1, "echo"), (2, "missing-tool")] { + let request = Request::post("/mcp") + .header("content-type", "application/json") + .header("accept", "application/json, text/event-stream") + .body(Body::from( + json!({"jsonrpc":"2.0","id":id,"method":"tools/call", + "params":{"name":name,"arguments":{"text":"mcp-secret-body"}}}) + .to_string(), + )) + .unwrap(); + assert_eq!( + app.clone().oneshot(request).await.unwrap().status(), + StatusCode::OK + ); + } + + let events = events(&telemetry); + assert_eq!(events.len(), 2); + assert_eq!(events[0]["ok"], true); + assert_eq!(events[1]["ok"], false); + assert!( + !serde_json::to_string(&events) + .unwrap() + .contains("mcp-secret-body") + ); +} + +#[tokio::test] +async fn mcp_is_error_and_dispatch_transport_failure_are_both_failure_observations() { + use kars_inference_router::mcp::tools::{ + DispatchError, SyncToAsync, ToolCallOutput, ToolCatalog, ToolDefinition, ToolDispatcher, + }; + struct Outcomes(ToolCatalog); + impl ToolDispatcher for Outcomes { + fn catalog(&self) -> &ToolCatalog { + &self.0 + } + fn invoke(&self, name: &str, _args: &Value) -> Result { + if name == "tool_error" { + Ok(ToolCallOutput { + content: vec![], + is_error: true, + }) + } else { + Err(DispatchError::ExecutionFailed { + tool: name.into(), + reason: "transport-secret-marker".into(), + }) + } + } + } + let catalog = ToolCatalog::new( + ["tool_error", "transport_error"] + .map(|name| ToolDefinition { + name: name.into(), + description: "test".into(), + input_schema: json!({"type":"object"}), + }) + .to_vec(), + ) + .unwrap(); + let telemetry = telemetry(); + let mut state = + McpRouteState::standard().with_tools(Arc::new(SyncToAsync::new(Outcomes(catalog)))); + state.task_telemetry = Some(telemetry.clone()); + let app = routes::mcp_route().with_state(state); + for (id, name) in [(1, "tool_error"), (2, "transport_error")] { + let request=Request::post("/mcp").header("content-type","application/json").header("accept","application/json, text/event-stream") + .body(Body::from(json!({"jsonrpc":"2.0","id":id,"method":"tools/call","params":{"name":name,"arguments":{}}}).to_string())).unwrap(); + assert_eq!( + app.clone().oneshot(request).await.unwrap().status(), + StatusCode::OK + ); + } + let events = events(&telemetry); + assert_eq!(events.len(), 2); + assert!( + events + .iter() + .all(|event| event["ok"] == false && event["http_status"] == 200) + ); + assert!( + !serde_json::to_string(&events) + .unwrap() + .contains("transport-secret-marker") + ); +} + +#[tokio::test] +async fn transport_failures_and_cancelled_partial_streams_never_become_success() { + let telemetry = telemetry(); + let result = proxy::forward( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::builder().no_proxy().build().unwrap(), + &upstream("http://127.0.0.1:1".into(), telemetry.clone()), + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await; + assert!(result.is_err()); + assert_eq!(events(&telemetry)[0]["outcome"], "transport_error"); + let mut observation = telemetry + .begin("v1/messages", "provider", "model", b"{}") + .unwrap(); + observation.headers(200); + let chunks = futures::stream::once(async { + Ok::<_, reqwest::Error>(Bytes::from_static( + b"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":12}}}\n\n", + )) + }) + .chain(futures::stream::pending()) + .boxed(); + let mut wrapped = wrap_stream(chunks, Some(observation), true); + wrapped.next().await.unwrap().unwrap(); + drop(wrapped); + let trace = events(&telemetry); + assert_eq!(trace[1]["outcome"], "cancelled"); + assert_eq!(trace[1]["usage"]["prompt_tokens"], 12); + assert!(trace[1]["usage"]["completion_tokens"].is_null()); +} diff --git a/inference-router/tests/multi_provider_guardrails.rs b/inference-router/tests/multi_provider_guardrails.rs index 61601ecd2..6f18201d8 100644 --- a/inference-router/tests/multi_provider_guardrails.rs +++ b/inference-router/tests/multi_provider_guardrails.rs @@ -61,6 +61,7 @@ async fn ollama_provider_forwards_openai_compat_without_auth() { .await; let upstream = UpstreamConfig { + telemetry: None, endpoint: server.uri(), deployment: "llama3.1".into(), sandbox_name: "test-sandbox".into(), @@ -118,6 +119,7 @@ async fn anthropic_provider_forwards_messages_with_router_held_key() { .await; let upstream = UpstreamConfig { + telemetry: None, endpoint: server.uri(), deployment: "claude-sonnet-4-5".into(), sandbox_name: "test-sandbox".into(), diff --git a/inference-router/tests/policy_status_endpoint.rs b/inference-router/tests/policy_status_endpoint.rs index b40f633b3..8c1699b68 100644 --- a/inference-router/tests/policy_status_endpoint.rs +++ b/inference-router/tests/policy_status_endpoint.rs @@ -47,6 +47,7 @@ fn test_state() -> (AppState, Arc) { policy_status.clone(), )); let state = AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/proxy_fake_upstream.rs b/inference-router/tests/proxy_fake_upstream.rs index 16c878a38..4475036d1 100644 --- a/inference-router/tests/proxy_fake_upstream.rs +++ b/inference-router/tests/proxy_fake_upstream.rs @@ -82,6 +82,7 @@ async fn api_key_mode_proxies_chat_completion_with_filter_results() { let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { + telemetry: None, endpoint, deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox".to_string(), @@ -163,6 +164,7 @@ async fn wi_mode_falls_back_to_imds_and_proxies_embeddings() { let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { + telemetry: None, endpoint, deployment: "text-embedding-3-small".to_string(), sandbox_name: "test-sandbox-wi".to_string(), @@ -236,6 +238,7 @@ async fn upstream_error_status_is_propagated() { let auth = WorkloadIdentityAuth::new(); let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { + telemetry: None, endpoint, deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox-429".to_string(), diff --git a/tests/e2e/governed-services.sh b/tests/e2e/governed-services.sh new file mode 100644 index 000000000..2371013ef --- /dev/null +++ b/tests/e2e/governed-services.sh @@ -0,0 +1,134 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +# The subshell keeps private fixture credentials and its cleanup trap local. +test_governed_services() ( + set +x + local k=(kubectl --context kind-kars-e2e) + local scratch forward_pid="" port="" token agent_token scope request_id new_scope code + local sandbox_uid namespace_uid + scratch=$(mktemp -d) || return 1 + cleanup_governed_smoke() { + if [ -n "$forward_pid" ]; then + kill "$forward_pid" 2>/dev/null || true + wait "$forward_pid" 2>/dev/null || true + fi + rm -f "$scratch/forward.log" "$scratch/response.json" "$scratch/request.json" + rmdir "$scratch" + } + trap cleanup_governed_smoke EXIT + + # Values travel only through the test process and curl's stdin, not argv/logs. + token=$("${k[@]}" get secret router-services-admin -n kars-e2e-test \ + --request-timeout=20s -o go-template='{{index .data "control-token" | base64decode}}') || return 1 + agent_token=$("${k[@]}" get secret router-admin-token -n kars-e2e-test \ + --request-timeout=20s -o go-template='{{index .data "token" | base64decode}}') || return 1 + [ -n "$token" ] && [ -n "$agent_token" ] && [ "$token" != "$agent_token" ] || return 1 + sandbox_uid=$("${k[@]}" get karssandbox e2e-test -n kars-system \ + --request-timeout=20s -o jsonpath='{.metadata.uid}') || return 1 + namespace_uid=$("${k[@]}" get namespace kars-e2e-test \ + --request-timeout=20s -o jsonpath='{.metadata.uid}') || return 1 + [ -n "$sandbox_uid" ] && [ -n "$namespace_uid" ] || return 1 + "${k[@]}" get deployment e2e-test -n kars-e2e-test --request-timeout=20s -o json \ + >"$scratch/response.json" || return 1 + python3 - "$scratch/response.json" <<'PY' || return 1 +import json, sys +pod = json.load(open(sys.argv[1]))["spec"]["template"]["spec"] +volume = next(v for v in pod["volumes"] if v.get("secret", {}).get("secretName") == "router-services-admin") +for container in pod["containers"]: + mounts = [m for m in container.get("volumeMounts", []) if m["name"] == volume["name"]] + if container["name"] == "inference-router": + assert len(mounts) == 1 and mounts[0]["mountPath"] == "/etc/kars/services" + assert mounts[0]["readOnly"] is True + else: + assert not mounts +PY + + "${k[@]}" port-forward --address 127.0.0.1 service/e2e-test -n kars-e2e-test :8443 \ + >"$scratch/forward.log" 2>&1 & + forward_pid=$! + local deadline=$(($(date +%s) + 30)) + while [ "$(date +%s)" -lt "$deadline" ]; do + kill -0 "$forward_pid" 2>/dev/null || return 1 + port=$(sed -n 's/^Forwarding from 127\.0\.0\.1:\([0-9]*\) ->.*/\1/p' "$scratch/forward.log" | head -1) + [ -z "$port" ] || break + sleep 1 + done + [ -n "$port" ] || return 1 + + service_request() { + local expected="$1" method="$2" path="$3" bearer="${4:-}" + local args=(--disable --silent --show-error --noproxy 127.0.0.1 --connect-timeout 5 --max-time 15 + --config - --request "$method" --url "http://127.0.0.1:$port$path" + --output "$scratch/response.json" --write-out '%{http_code}') + if [ "$method" = POST ]; then + args+=(--header 'Content-Type: application/json' --data-binary "@$scratch/request.json") + fi + code=$( + { [ -z "$bearer" ] || printf 'header = "Authorization: Bearer %s"\n' "$bearer"; } \ + | curl "${args[@]}" + ) || return 1 + if [ "$code" != "$expected" ]; then + printf 'Governed service %s %s returned %s, expected %s\n' "$method" "$path" "$code" "$expected" >&2 + return 1 + fi + } + service_body() { + python3 - "$scratch/request.json" "$@" <<'PY' +import json, sys +args = sys.argv[2:] +assert len(args) % 2 == 0 +with open(sys.argv[1], "w") as output: + json.dump(dict(zip(args[::2], args[1::2])), output) +PY + } + response_field() { + python3 - "$scratch/response.json" "$1" <<'PY' +import json, sys +value = json.load(open(sys.argv[1])) +for field in sys.argv[2].split("."): + value = value[field] +assert isinstance(value, str) and value +print(value) +PY + } + + service_request 401 GET /internal/access-requests || return 1 + service_request 401 GET /internal/access-requests "$agent_token" || return 1 + service_request 200 GET /internal/access-requests "$token" || return 1 + python3 - "$scratch/response.json" "$sandbox_uid" "$namespace_uid" <<'PY' || return 1 +import json, sys +response = json.load(open(sys.argv[1])) +identity = response["scope"]["identity"] +assert identity["sandbox"] == {"namespace": "kars-system", "name": "e2e-test", "uid": sys.argv[2]} +assert identity["namespace_uid"] == sys.argv[3] +assert identity.get("task") is None +assert response["enforcement_changed"] is False +PY + scope=$(response_field scope.id) || return 1 + service_body scope_id "$scope" kind egress target example.invalid reason fixture || return 1 + service_request 202 POST /v1/access-request || return 1 + request_id=$(response_field request.request_id) || return 1 + service_body scope_id "$scope" request_id "$request_id" verdict approved || return 1 + service_request 401 POST /internal/access-requests/decision "$agent_token" || return 1 + service_request 200 POST /internal/access-requests/decision "$token" || return 1 + python3 - "$scratch/response.json" <<'PY' || return 1 +import json, sys +response = json.load(open(sys.argv[1])) +assert response["request"]["status"] == "approved" +assert response["enforcement_changed"] is False +PY + service_body scope_id "$scope" assignment_id fixture-assignment || return 1 + service_request 401 POST /internal/access-requests/reset "$agent_token" || return 1 + service_request 200 POST /internal/access-requests/reset "$token" || return 1 + new_scope=$(response_field scope.id) || return 1 + [ "$new_scope" != "$scope" ] || return 1 + [ "$(response_field scope.identity.sandbox.uid)" = "$sandbox_uid" ] || return 1 + service_body scope_id "$scope" request_id "$request_id" verdict approved || return 1 + service_request 409 POST /internal/access-requests/decision "$token" || return 1 + service_body scope_id "$scope" kind egress target example.invalid || return 1 + service_request 409 POST /v1/access-request || return 1 + service_request 200 GET /telemetry/cursor || return 1 + [ "$(response_field scope_id)" = "$new_scope" ] || return 1 +) diff --git a/tests/e2e/run.sh b/tests/e2e/run.sh index 00b97055b..028203ab0 100755 --- a/tests/e2e/run.sh +++ b/tests/e2e/run.sh @@ -3041,6 +3041,7 @@ EOF source "$SCRIPT_DIR/sre-authority.sh" source "$SCRIPT_DIR/namespace-ownership.sh" source "$SCRIPT_DIR/credential-sources.sh" +source "$SCRIPT_DIR/governed-services.sh" main() { umask 077 @@ -3079,6 +3080,11 @@ main() { test_sandbox_namespace_labels || true test_sandbox_deployment_exists || true test_sandbox_pod_starts || true + if test_governed_services; then + pass "Service API smoke: router-only token mount, credential checks, and scope reset" + else + fail "Governed service authentication and scope lifecycle gate failed" + fi test_sandbox_networkpolicy_denies_ingress || true test_sandbox_suspended_lifecycle || true test_secondary_resource_watch || true