fix: authorize restore coordination with worker input capability

crates/workspace-server/src/hosts.rs

@@ -3951,6 +3951,7 @@ fn workspace_runtime_operation(method: &str, path_and_query: &str) -> &'static s
     }
     if path.ends_with("/input")
         || path.ends_with("/restore")
+        || path.ends_with("/restore/coordinate")
         || (path.contains("/attachments") && method != "GET")
         || path.contains("/workdir-attachments")
     {
@@ -6513,6 +6514,8 @@ pub fn placeholder_spawn_response(host_id: impl Into<String>) -> WorkerSpawnResu
 
 #[cfg(test)]
 mod tests {
+    mod restore_auth_tests;
+
     use super::*;
     use serde_json::json;
     use std::collections::HashMap;

crates/workspace-server/src/hosts/tests/restore_auth_tests.rs

@@ -0,0 +1,175 @@
+use super::*;
+use axum::{
+    body::{Body, to_bytes},
+    http::Request,
+};
+use tower::ServiceExt;
+use worker_runtime::{
+    Runtime,
+    auth::RuntimeIdentityMaterial,
+    http_server::{WorkspaceRuntimeHttpAuth, runtime_http_router_with_workspace_auth},
+    workspace_issuer::{
+        InMemoryWorkspaceClaimReplayProtection, InMemoryWorkspaceRuntimeVerificationAuthority,
+        RuntimeVerificationSigner, WorkspaceCapabilityVerifier, WorkspaceIssuerTrustRecord,
+        WorkspaceIssuerTrustState, WorkspaceRuntimeVerificationAuthority,
+        WorkspaceRuntimeVerificationRecord, issue_workspace_capability_token,
+    },
+};
+
+#[tokio::test]
+async fn restore_capabilities_authorize_runtime_routes_without_relaxing_request_binding() {
+    let workspace_identity = RuntimeIdentityMaterial::generate("workspace-key").unwrap();
+    let runtime_identity = RuntimeIdentityMaterial::generate("runtime-test").unwrap();
+    let signer = RuntimeVerificationSigner::from_identity(&runtime_identity).unwrap();
+    let now = Utc::now().timestamp();
+    let verifications = Arc::new(InMemoryWorkspaceRuntimeVerificationAuthority::default());
+    verifications
+        .record(WorkspaceRuntimeVerificationRecord {
+            workspace_id: "workspace-test".into(),
+            runtime_id: "runtime-test".into(),
+            binding_revision: 1,
+            workspace_key_id: "workspace-key".into(),
+            workspace_identity_revision: 1,
+            workspace_trust_generation: 1,
+            runtime_public_key_fingerprint: signer.public_key_fingerprint().to_owned(),
+            runtime_identity_revision: 1,
+            verified_at: now,
+        })
+        .unwrap();
+    let public_key =
+        worker_runtime::auth::decode_public_key(&workspace_identity.public_key).unwrap();
+    let verifier = WorkspaceCapabilityVerifier::new(
+        vec![WorkspaceIssuerTrustRecord {
+            workspace_id: "workspace-test".into(),
+            backend_url: "https://backend.test".into(),
+            key_id: "workspace-key".into(),
+            algorithm: "ed25519".into(),
+            public_key: workspace_identity.public_key.clone(),
+            public_key_fingerprint: format!(
+                "sha256:{}",
+                Sha256::digest(public_key)
+                    .iter()
+                    .map(|byte| format!("{byte:02x}"))
+                    .collect::<String>()
+            ),
+            identity_revision: 1,
+            trust_generation: 1,
+            state: WorkspaceIssuerTrustState::Active,
+            registered_at_unix: now,
+            updated_at_unix: now,
+        }],
+        Arc::new(InMemoryWorkspaceClaimReplayProtection::default()),
+    )
+    .unwrap();
+    // No Worker execution is needed: a typed worker_not_found proves the request passed real auth.
+    let app = runtime_http_router_with_workspace_auth(
+        Runtime::new_memory(),
+        None,
+        WorkspaceRuntimeHttpAuth {
+            verifier,
+            signer,
+            verifications,
+        },
+    );
+    let worker_id = EmbeddedWorkerId::from_legacy_u64(1).to_string();
+    let body = serde_json::to_vec(&runtime_api::WorkerRestoreCoordinationRequest {
+        expected_observation_token: "observed-generation".into(),
+        request_id: "restore-request".into(),
+        preparation: None,
+    })
+    .unwrap();
+
+    for suffix in [
+        "restore",
+        "restore/coordinate",
+        "restore/coordinate?probe=1",
+    ] {
+        let path = format!("/v1/workers/{worker_id}/{suffix}");
+        for case in [
+            "authorized",
+            "wrong-operation",
+            "wrong-worker",
+            "wrong-body",
+        ] {
+            let claims = WorkspaceCapabilityClaims {
+                issuer: "https://backend.test".into(),
+                issuer_workspace_id: "workspace-test".into(),
+                issuer_key_id: "workspace-key".into(),
+                issuer_identity_revision: 1,
+                trust_generation: 1,
+                binding_revision: 1,
+                runtime_id: "runtime-test".into(),
+                worker_id: if case == "wrong-worker" {
+                    Some(EmbeddedWorkerId::from_legacy_u64(2).to_string())
+                } else {
+                    worker_id_from_remote_path(&path)
+                },
+                operation: if case == "wrong-operation" {
+                    "runtime:read"
+                } else {
+                    workspace_runtime_operation("POST", &path)
+                }
+                .into(),
+                method: "POST".into(),
+                path_and_query: path.clone(),
+                body_digest: workspace_request_body_digest(&body),
+                iat: now,
+                exp: now + 60,
+                jti: uuid::Uuid::now_v7().to_string(),
+            };
+            let token = issue_workspace_capability_token(
+                &workspace_identity.signing_key().unwrap(),
+                &claims,
+            )
+            .unwrap();
+            let mut sent_body = body.clone();
+            if case == "wrong-body" {
+                sent_body.push(b' ');
+            }
+            let response = app
+                .clone()
+                .oneshot(
+                    Request::builder()
+                        .method("POST")
+                        .uri(&path)
+                        .header("authorization", format!("Bearer {token}"))
+                        .header("content-type", "application/json")
+                        .body(Body::from(sent_body))
+                        .unwrap(),
+                )
+                .await
+                .unwrap();
+            let status = response.status();
+            let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
+            let error: RuntimeHttpErrorResponse = serde_json::from_slice(&bytes).unwrap();
+            if case == "authorized" {
+                assert_eq!(
+                    status,
+                    StatusCode::NOT_FOUND,
+                    "{path}: {}",
+                    error.error.message
+                );
+                assert_eq!(error.error.code, "worker_not_found", "{path}");
+            } else {
+                assert_eq!(
+                    status,
+                    StatusCode::UNAUTHORIZED,
+                    "{path} {case}: {}",
+                    error.error.message
+                );
+                assert_eq!(error.error.code, "unauthorized");
+                let expected = match case {
+                    "wrong-operation" => "does not authorize this operation",
+                    "wrong-worker" => "targets another Worker",
+                    "wrong-body" => "does not bind this request body",
+                    _ => unreachable!(),
+                };
+                assert!(
+                    error.error.message.contains(expected),
+                    "{path} {case}: {}",
+                    error.error.message
+                );
+            }
+        }
+    }
+}