fix(desktop): attest proxy before sending admin token

desktop/src-tauri/Cargo.toml

@@ -15,9 +15,12 @@ crate-type = ["staticlib", "cdylib", "rlib"]
 tauri-build = { version = "=2.6.3", features = [] }
 
 [dependencies]
+base64 = "=0.22.1"
+hmac = "=0.12.1"
 reqwest = { version = "=0.12.24", default-features = false, features = ["json", "rustls-tls"] }
 serde = { version = "=1.0.219", features = ["derive"] }
 serde_json = "=1.0.140"
+sha2 = "=0.10.9"
 uuid = { version = "=1.18.1", features = ["v4"] }
 tauri = { version = "=2.11.6", features = ["tray-icon", "image-png"] }
 tauri-plugin-autostart = "=2.5.0"

desktop/src-tauri/src/auth.rs

@@ -1,5 +1,15 @@
 use std::path::PathBuf;
 
+use serde::Deserialize;
+
+#[derive(Clone, Debug, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "camelCase")]
+pub struct RuntimeIdentity {
+    pub pid: u32,
+    pub port: u16,
+    pub attestation_secret: String,
+}
+
 #[derive(Clone, Debug)]
 pub struct Auth {
     home: PathBuf,
@@ -25,6 +35,15 @@ impl Auth {
         })
     }
 
+    pub fn runtime_identity(&self) -> Option<RuntimeIdentity> {
+        let value = std::fs::read(self.home.join("runtime-port.json")).ok()?;
+        let identity: RuntimeIdentity = serde_json::from_slice(&value).ok()?;
+        if identity.pid == 0 || identity.attestation_secret.len() != 43 {
+            return None;
+        }
+        Some(identity)
+    }
+
     pub fn user_agent() -> &'static str {
         concat!("OpenCodexDesktop/", env!("CARGO_PKG_VERSION"))
     }

desktop/src-tauri/src/proxy.rs

@@ -1,6 +1,12 @@
-use crate::{auth::Auth, discovery::ProxyEndpoint};
+use crate::{
+    auth::{Auth, RuntimeIdentity},
+    discovery::ProxyEndpoint,
+};
+use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
+use hmac::{Hmac, Mac};
 use reqwest::{Client, Method, StatusCode};
 use serde_json::Value;
+use sha2::Sha256;
 use std::time::Duration;
 
 #[derive(Clone)]
@@ -73,13 +79,50 @@ impl ProxyClient {
     async fn request(&self, method: Method, path: &str) -> Result<Value, ProxyError> {
         let response = self.send(&method, path, None).await?;
         if response.status() == StatusCode::UNAUTHORIZED {
+            self.authenticate_target().await?;
             let token = self.auth.token().ok_or(ProxyError::Unauthorized)?;
             let response = self.send(&method, path, Some(token)).await?;
             return decode(response).await;
         }
         decode(response).await
     }
 
+    async fn authenticate_target(&self) -> Result<(), ProxyError> {
+        let identity = self
+            .auth
+            .runtime_identity()
+            .ok_or(ProxyError::Unauthorized)?;
+        if identity.port != self.endpoint.port {
+            return Err(ProxyError::Unauthorized);
+        }
+        let mut challenge_bytes = [0_u8; 32];
+        challenge_bytes[..16].copy_from_slice(uuid::Uuid::new_v4().as_bytes());
+        challenge_bytes[16..].copy_from_slice(uuid::Uuid::new_v4().as_bytes());
+        let challenge = URL_SAFE_NO_PAD.encode(challenge_bytes);
+        let response = self
+            .client
+            .get(self.endpoint.url("/healthz"))
+            .header("x-opencodex-attestation-challenge", &challenge)
+            .send()
+            .await
+            .map_err(map_request_error)?;
+        let proof = response
+            .headers()
+            .get("x-opencodex-attestation-proof")
+            .and_then(|value| value.to_str().ok())
+            .map(str::to_owned);
+        let health: Value = decode(response).await?;
+        if health.get("service").and_then(Value::as_str) != Some("opencodex")
+            || health.get("pid").and_then(Value::as_u64) != Some(identity.pid.into())
+            || health.get("port").and_then(Value::as_u64) != Some(identity.port.into())
+            || !valid_attestation_proof(&identity, &challenge, proof.as_deref())
+            || self.auth.runtime_identity().as_ref() != Some(&identity)
+        {
+            return Err(ProxyError::Unauthorized);
+        }
+        Ok(())
+    }
+
     async fn send(
         &self,
         method: &Method,
@@ -90,13 +133,35 @@ impl ProxyClient {
         if let Some(value) = token {
             request = request.header("X-OpenCodex-API-Key", value);
         }
-        request.send().await.map_err(|error| {
-            if error.is_connect() {
-                ProxyError::Unreachable
-            } else {
-                ProxyError::Decode(error)
-            }
-        })
+        request.send().await.map_err(map_request_error)
+    }
+}
+
+fn valid_attestation_proof(
+    identity: &RuntimeIdentity,
+    challenge: &str,
+    proof: Option<&str>,
+) -> bool {
+    let Ok(mut mac) = Hmac::<Sha256>::new_from_slice(identity.attestation_secret.as_bytes()) else {
+        return false;
+    };
+    mac.update(
+        format!(
+            "opencodex-local-management-v1\n{challenge}\n{}\n{}",
+            identity.pid, identity.port
+        )
+        .as_bytes(),
+    );
+    proof
+        .and_then(|value| URL_SAFE_NO_PAD.decode(value).ok())
+        .is_some_and(|value| mac.verify_slice(&value).is_ok())
+}
+
+fn map_request_error(error: reqwest::Error) -> ProxyError {
+    if error.is_connect() {
+        ProxyError::Unreachable
+    } else {
+        ProxyError::Decode(error)
     }
 }
 
@@ -109,3 +174,32 @@ async fn decode(response: reqwest::Response) -> Result<Value, ProxyError> {
     }
     response.json().await.map_err(ProxyError::Decode)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn identity() -> RuntimeIdentity {
+        RuntimeIdentity {
+            pid: 4242,
+            port: 10100,
+            attestation_secret: "BwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwcHBwc".into(),
+        }
+    }
+
+    #[test]
+    fn accepts_only_a_proof_bound_to_the_runtime_identity() {
+        let challenge = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";
+        let proof = "Yr9EKHjeAFfsFMsF8Xsd7J6LxBYnObweKZlLyTMk0Lo";
+        assert!(valid_attestation_proof(&identity(), challenge, Some(proof)));
+
+        let mut replacement = identity();
+        replacement.pid += 1;
+        assert!(!valid_attestation_proof(
+            &replacement,
+            challenge,
+            Some(proof)
+        ));
+        assert!(!valid_attestation_proof(&identity(), challenge, None));
+    }
+}