feat: integrate Workspace API authorization

crates/workspace-server/src/auth.rs

@@ -176,8 +176,28 @@ fn actor_for_user<S: ControlPlaneStore + ?Sized>(
     }))
 }
 
-pub fn session_set_cookie(cookie_name: &str, token: &str, max_age_seconds: i64) -> String {
-    format!("{cookie_name}={token}; Max-Age={max_age_seconds}; Path=/; HttpOnly; SameSite=Lax")
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct SessionCookiePolicy<'a> {
+    pub cookie_name: &'a str,
+    pub path: &'a str,
+    pub domain: Option<&'a str>,
+    pub secure: bool,
+}
+
+pub fn session_set_cookie(
+    policy: SessionCookiePolicy<'_>,
+    token: &str,
+    max_age_seconds: i64,
+) -> String {
+    let domain = policy
+        .domain
+        .map(|domain| format!("; Domain={domain}"))
+        .unwrap_or_default();
+    let secure = if policy.secure { "; Secure" } else { "" };
+    format!(
+        "{}={token}; Max-Age={max_age_seconds}; Path={}; HttpOnly; SameSite=Lax{domain}{secure}",
+        policy.cookie_name, policy.path
+    )
 }
 
 pub fn auth_error(code: &str, message: &str) -> Error {

crates/workspace-server/src/store.rs

@@ -761,6 +761,7 @@ pub trait ControlPlaneStore: Send + Sync {
         record: &WorkspaceBootstrapRecord,
     ) -> Result<WorkspaceBootstrapResult>;
     async fn get_trusted_runtime(&self, runtime_id: &str) -> Result<Option<TrustedRuntimeRecord>>;
+    async fn upsert_trusted_runtime_record(&self, record: &TrustedRuntimeRecord) -> Result<()>;
     async fn consume_worker_mutation_source_jti(
         &self,
         runtime_id: &str,
@@ -1975,6 +1976,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
         })
     }
 
+    async fn upsert_trusted_runtime_record(&self, record: &TrustedRuntimeRecord) -> Result<()> {
+        SqliteWorkspaceStore::upsert_trusted_runtime(self, record)
+    }
+
     async fn consume_worker_mutation_source_jti(
         &self,
         runtime_id: &str,

crates/workspace-server/src/worker_source.rs

@@ -3,14 +3,118 @@ use std::time::{SystemTime, UNIX_EPOCH};
 
 use axum::http::HeaderMap;
 use worker_runtime::auth::{
-    WorkerMutationActorKind, WorkerMutationOperation, WorkerMutationSourceClaims,
-    WorkerMutationSourceExpectation, decode_worker_mutation_source_claims,
-    verify_worker_mutation_source_proof,
+    RuntimeRequestSourceExpectation, WorkerMutationActorKind, WorkerMutationOperation,
+    WorkerMutationSourceClaims, WorkerMutationSourceExpectation,
+    decode_runtime_request_source_claims, decode_worker_mutation_source_claims,
+    verify_runtime_request_source, verify_worker_mutation_source_proof,
 };
 use worker_runtime::worker_source::InProcessWorkerMutationProof;
 
 use crate::hosts::RemoteRuntimeConfig;
-use crate::server::WorkspaceApi;
+use crate::server::{ServerConfig, WorkspaceApi};
+use crate::store::ControlPlaneStore;
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct VerifiedRuntimeRequestSource {
+    pub runtime_id: String,
+    pub worker_id: Option<String>,
+}
+
+pub async fn verify_runtime_request_source_proof(
+    api: &WorkspaceApi,
+    proof: &str,
+    workspace_id: &str,
+    permission: &str,
+    method: &str,
+    path: &str,
+    body_digest: &str,
+) -> Result<VerifiedRuntimeRequestSource, WorkerMutationSourceProofError> {
+    verify_runtime_request_source_proof_with_store(
+        api.store.as_ref(),
+        &api.config,
+        proof,
+        workspace_id,
+        permission,
+        method,
+        path,
+        body_digest,
+    )
+    .await
+}
+
+#[allow(clippy::too_many_arguments)]
+pub async fn verify_runtime_request_source_proof_with_store(
+    store: &dyn ControlPlaneStore,
+    config: &ServerConfig,
+    proof: &str,
+    workspace_id: &str,
+    permission: &str,
+    method: &str,
+    path: &str,
+    body_digest: &str,
+) -> Result<VerifiedRuntimeRequestSource, WorkerMutationSourceProofError> {
+    let unverified = decode_runtime_request_source_claims(proof)
+        .map_err(|_| WorkerMutationSourceProofError::Invalid)?;
+    let audience = remote_audience(config, &unverified.iss, workspace_id)?;
+    let trusted = store
+        .get_trusted_runtime(&unverified.iss)
+        .await
+        .map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?
+        .filter(|record| record.revoked_at.is_none())
+        .ok_or(WorkerMutationSourceProofError::RevokedRuntimeTrust)?;
+    let trusted_for_workspace = trusted.workspace_id.as_deref() == Some(workspace_id)
+        || (unverified.iss == crate::hosts::EMBEDDED_RUNTIME_ID && trusted.workspace_id.is_none());
+    if !trusted_for_workspace {
+        return Err(WorkerMutationSourceProofError::WrongWorkspace);
+    }
+    let expected = RuntimeRequestSourceExpectation {
+        identity_id: &unverified.iss,
+        audience: audience.as_ref(),
+        workspace_id,
+        worker_id: unverified.worker_id.as_deref(),
+        permission,
+        method,
+        path,
+        body_digest,
+        now_unix: i64::try_from(unix_now_seconds()).unwrap_or(i64::MAX),
+    };
+    let claims = verify_runtime_request_source(proof, &trusted.public_key, &expected)
+        .map_err(map_auth_error)?;
+    let now_seconds = u64::try_from(expected.now_unix).unwrap_or(u64::MAX);
+    let expires_at = u64::try_from(claims.exp).unwrap_or(0);
+    let consumed_at = chrono::DateTime::from_timestamp(expected.now_unix, 0)
+        .ok_or(WorkerMutationSourceProofError::Expired)?
+        .to_rfc3339();
+    if !store
+        .consume_worker_mutation_source_jti(
+            &claims.iss,
+            &claims.jti,
+            expires_at,
+            now_seconds,
+            &consumed_at,
+        )
+        .await
+        .map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?
+    {
+        return Err(WorkerMutationSourceProofError::Replay);
+    }
+    if let Some(worker_id) = claims.worker_id.as_deref() {
+        let worker = worker_runtime::identity::RuntimeWorkerRef {
+            runtime_id: claims.iss.clone(),
+            worker_id: worker_id.to_owned(),
+        };
+        let member = store
+            .get_worker_registry(workspace_id, &worker)
+            .map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?;
+        if member.is_none() {
+            return Err(WorkerMutationSourceProofError::WorkerCatalogMembership);
+        }
+    }
+    Ok(VerifiedRuntimeRequestSource {
+        runtime_id: claims.iss,
+        worker_id: claims.worker_id,
+    })
+}
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum PresentedWorkerMutationSourceProof<'a> {
@@ -97,16 +201,19 @@ async fn verify_worker_remove_source_with(
         PresentedWorkerMutationSourceProof::Remote(token) => {
             let unverified = decode_worker_mutation_source_claims(token)
                 .map_err(|_| WorkerMutationSourceProofError::Invalid)?;
-            let audience = remote_audience(config, &unverified.iss)?;
+            let audience = remote_audience(config, &unverified.iss, &config.workspace_id)?;
             let trusted = store
                 .get_trusted_runtime(&unverified.iss)
                 .await
                 .map_err(|error| WorkerMutationSourceProofError::Authority(error.to_string()))?
                 .filter(|record| record.revoked_at.is_none())
                 .ok_or(WorkerMutationSourceProofError::RevokedRuntimeTrust)?;
+            if trusted.workspace_id.as_deref() != Some(config.workspace_id.as_str()) {
+                return Err(WorkerMutationSourceProofError::WrongWorkspace);
+            }
             let expected = WorkerMutationSourceExpectation {
                 runtime_id: &unverified.iss,
-                audience,
+                audience: audience.as_ref(),
                 workspace_id: &config.workspace_id,
                 worker_id: None,
                 actor_kind: WorkerMutationActorKind::Worker,
@@ -247,13 +354,17 @@ impl worker_runtime::worker_source::EmbeddedWorkerMutationDispatcher
 fn remote_audience<'a>(
     config: &'a crate::server::ServerConfig,
     runtime_id: &str,
-) -> Result<&'a str, WorkerMutationSourceProofError> {
+    workspace_id: &str,
+) -> Result<std::borrow::Cow<'a, str>, WorkerMutationSourceProofError> {
+    if runtime_id == crate::hosts::EMBEDDED_RUNTIME_ID {
+        return Ok(std::borrow::Cow::Owned(format!("embedded:{workspace_id}")));
+    }
     config
         .remote_runtime_sources
         .iter()
         .find(|runtime| runtime.runtime_id == runtime_id)
         .and_then(|runtime: &RemoteRuntimeConfig| runtime.auth.as_ref())
-        .map(|auth| auth.server_id.as_str())
+        .map(|auth| std::borrow::Cow::Borrowed(auth.server_id.as_str()))
         .ok_or(WorkerMutationSourceProofError::RevokedRuntimeTrust)
 }