diff --git a/crates/worker-runtime/src/lib.rs b/crates/worker-runtime/src/lib.rs index bd4ed723..a2b49dcb 100644 --- a/crates/worker-runtime/src/lib.rs +++ b/crates/worker-runtime/src/lib.rs @@ -30,6 +30,8 @@ pub mod worker_backend; pub mod worker_source; pub mod working_directory; pub mod workspace_issuer; +#[cfg(feature = "http-server")] +pub mod workspace_request; #[cfg(feature = "fs-store")] pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions}; diff --git a/crates/worker-runtime/src/main.rs b/crates/worker-runtime/src/main.rs index 7019d3c6..7c05ecb7 100644 --- a/crates/worker-runtime/src/main.rs +++ b/crates/worker-runtime/src/main.rs @@ -33,6 +33,7 @@ use worker_runtime::workspace_issuer::{ WorkspaceIssuerTrustState, add_workspace_issuer_trust, replace_workspace_issuer_trust, revoke_workspace_issuer_trust, validate_workspace_issuer_trust_records, }; +use worker_runtime::workspace_request::RuntimeWorkspaceRequestClient; use worker_runtime::{Runtime, RuntimeOptions}; fn main() -> ExitCode { @@ -211,6 +212,22 @@ fn build_runtime(config: &ProcessConfig) -> Result { if let Some(identity) = runtime_auth.identity.clone() { factory = factory.with_remote_worker_mutation_identity(identity); } + if let Some(identity) = runtime_auth.identity.as_ref() { + for issuer in runtime_auth + .workspace_issuers + .iter() + .filter(|issuer| issuer.state == WorkspaceIssuerTrustState::Active) + { + factory = factory.with_workspace_request_client( + RuntimeWorkspaceRequestClient::new( + issuer.workspace_id.clone(), + issuer.backend_url.clone(), + identity.identity_id.clone(), + ) + .with_runtime_request_source(identity, issuer.backend_url.clone()), + ); + } + } let mut backend_resource_client: Option< Arc, > = None; @@ -231,7 +248,14 @@ fn build_runtime(config: &ProcessConfig) -> Result { endpoint, config.backend_resource_token.clone(), ) - .with_runtime_request_source(identity, workspace_issuer.backend_url.clone()), + .with_workspace_request_client( + RuntimeWorkspaceRequestClient::new( + workspace_issuer.workspace_id.clone(), + workspace_issuer.backend_url.clone(), + identity.identity_id.clone(), + ) + .with_runtime_request_source(identity, workspace_issuer.backend_url.clone()), + ), ); factory = factory.with_resource_client(client.clone()); backend_resource_client = Some(client); @@ -257,7 +281,14 @@ fn build_runtime(config: &ProcessConfig) -> Result { endpoint, config.backend_resource_token.clone(), ) - .with_runtime_request_source(identity, workspace_issuer.backend_url.clone()), + .with_workspace_request_client( + RuntimeWorkspaceRequestClient::new( + workspace_issuer.workspace_id.clone(), + workspace_issuer.backend_url.clone(), + identity.identity_id.clone(), + ) + .with_runtime_request_source(identity, workspace_issuer.backend_url.clone()), + ), ); workspace_backend_resource_clients .push((workspace_issuer.workspace_id.clone(), client)); diff --git a/crates/worker-runtime/src/resource.rs b/crates/worker-runtime/src/resource.rs index 4f1c9487..9f76ce34 100644 --- a/crates/worker-runtime/src/resource.rs +++ b/crates/worker-runtime/src/resource.rs @@ -1,9 +1,7 @@ -use crate::auth::{ - BACKEND_RESOURCE_FETCH_PERMISSION, RUNTIME_REQUEST_SOURCE_PROOF_HEADER, - RuntimeIdentityMaterial, RuntimeRequestSourceSigner, unix_now_seconds, -}; +use crate::auth::BACKEND_RESOURCE_FETCH_PERMISSION; use crate::identity::WorkerId; use crate::profile_archive::{ProfileSourceArchive, ProfileSourceArchiveRef, sha256_hex}; +use crate::workspace_request::{RuntimeWorkspaceRequest, RuntimeWorkspaceRequestClient}; use async_trait::async_trait; use serde::{Deserialize, Serialize}; @@ -188,10 +186,8 @@ pub trait BackendResourceClient: Send + Sync + 'static { pub struct HttpBackendResourceClient { endpoint: String, bearer_token: Option, - request_source_signer: Option, - request_source_audience: Option, + workspace_request_client: Option, request_timeout: std::time::Duration, - client: reqwest::Client, } #[cfg(feature = "http-server")] @@ -200,10 +196,8 @@ impl HttpBackendResourceClient { Self { endpoint: endpoint.into(), bearer_token, - request_source_signer: None, - request_source_audience: None, + workspace_request_client: None, request_timeout: DEFAULT_BACKEND_RESOURCE_FETCH_TIMEOUT, - client: reqwest::Client::new(), } } @@ -212,13 +206,8 @@ impl HttpBackendResourceClient { self } - pub fn with_runtime_request_source( - mut self, - identity: &RuntimeIdentityMaterial, - audience: impl Into, - ) -> Self { - self.request_source_signer = Some(RuntimeRequestSourceSigner::from_identity(identity)); - self.request_source_audience = Some(audience.into()); + pub fn with_workspace_request_client(mut self, client: RuntimeWorkspaceRequestClient) -> Self { + self.workspace_request_client = Some(client); self } } @@ -240,59 +229,73 @@ impl BackendResourceClient for HttpBackendResourceClient { message: error.to_string(), } })?; - let mut builder = self - .client - .post(endpoint.clone()) - .timeout(self.request_timeout) - .header(reqwest::header::CONTENT_TYPE, "application/json") - .body(body.clone()); - if let Some(signer) = self.request_source_signer.as_ref() { - let audience = self.request_source_audience.as_deref().ok_or_else(|| { - BackendResourceError::Unauthorized { - message: "Runtime request proof audience is unavailable".to_owned(), - } - })?; - let proof = signer - .issue( - audience, - &request.handle.workspace_id, - None, - BACKEND_RESOURCE_FETCH_PERMISSION, - "POST", - endpoint.path(), - &body, - i64::try_from(unix_now_seconds()).unwrap_or(i64::MAX), - 30, - ) - .map_err(|error| BackendResourceError::Unauthorized { - message: error.to_string(), - })?; - builder = builder.header(RUNTIME_REQUEST_SOURCE_PROOF_HEADER, proof); - } - let builder = if let Some(token) = self.bearer_token.as_deref() { - builder.bearer_auth(token) - } else { - builder - }; - let response = builder.send().await.map_err(|error| { - if error.is_timeout() { - BackendResourceError::Timeout - } else { - BackendResourceError::Transport { - message: error.to_string(), - } + let client = self.workspace_request_client.as_ref().ok_or_else(|| { + BackendResourceError::Unauthorized { + message: "Workspace request client is unavailable".to_string(), } })?; - if response.status().is_success() { - response - .json::() - .await - .map_err(|err| BackendResourceError::InvalidResponse { + if client.workspace_id() != request.handle.workspace_id { + return Err(BackendResourceError::Unauthorized { + message: "Workspace request client does not match the resource workspace" + .to_string(), + }); + } + let base_url = client.base_url().trim_end_matches('/'); + let endpoint_text = endpoint.as_str(); + let endpoint_suffix = endpoint_text.strip_prefix(base_url).ok_or_else(|| { + BackendResourceError::Unauthorized { + message: "Workspace resource endpoint does not match its request client" + .to_string(), + } + })?; + if !endpoint_suffix.starts_with('/') { + return Err(BackendResourceError::Unauthorized { + message: "Workspace resource endpoint does not match its request client" + .to_string(), + }); + } + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert( + reqwest::header::CONTENT_TYPE, + reqwest::header::HeaderValue::from_static("application/json"), + ); + if let Some(token) = self.bearer_token.as_deref() { + let value = reqwest::header::HeaderValue::from_str(&format!("Bearer {token}")) + .map_err(|error| BackendResourceError::Transport { + message: error.to_string(), + })?; + headers.insert(reqwest::header::AUTHORIZATION, value); + } + let response = client + .execute(RuntimeWorkspaceRequest { + method: reqwest::Method::POST, + path_and_query: endpoint_suffix.to_string(), + body, + headers, + permission: BACKEND_RESOURCE_FETCH_PERMISSION.to_string(), + worker_id: None, + timeout: Some(self.request_timeout), + max_response_bytes: 8 * 1024 * 1024, + }) + .await + .map_err(|error| { + if error.is_timeout() { + BackendResourceError::Timeout + } else { + BackendResourceError::Transport { + message: error.to_string(), + } + } + })?; + if response.status.is_success() { + serde_json::from_slice::(&response.body).map_err(|err| { + BackendResourceError::InvalidResponse { message: err.to_string(), - }) + } + }) } else { - let status = response.status(); - match response.json::().await { + let status = response.status; + match serde_json::from_slice::(&response.body) { Ok(error) => Err(error), Err(err) => Err(BackendResourceError::Transport { message: format!("backend resource fetch failed with HTTP {status}: {err}"), @@ -383,6 +386,7 @@ pub fn validate_resource_handle_text(label: &str, value: &str) -> Result<(), Str #[cfg(test)] mod tests { use super::*; + use crate::auth::RuntimeIdentityMaterial; use crate::profile_archive::ProfileSourceGraphSummary; use std::collections::BTreeMap; @@ -435,7 +439,14 @@ mod tests { let handle = handle_for(b"archive-bytes"); let client = HttpBackendResourceClient::new(format!("{base_url}/fetch"), None) .with_request_timeout(std::time::Duration::from_millis(25)) - .with_runtime_request_source(&identity, base_url); + .with_workspace_request_client( + RuntimeWorkspaceRequestClient::new( + "workspace-test", + base_url.clone(), + "runtime-test", + ) + .with_runtime_request_source(&identity, base_url), + ); let error = client .fetch_resource(BackendResourceFetchRequest { diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 8f15bd67..758f4188 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -14,10 +14,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock, mpsc}; use std::time::Duration; -use crate::auth::{ - BACKEND_RESOURCE_FETCH_PERMISSION, RUNTIME_REQUEST_SOURCE_PROOF_HEADER, - RuntimeIdentityMaterial, RuntimeRequestSourceSigner, unix_now_seconds, -}; +use crate::auth::{BACKEND_RESOURCE_FETCH_PERMISSION, RuntimeIdentityMaterial}; use crate::catalog::{ CreateWorkerRequest, ProfileSourceArchiveSource, RepositoryRefObservation, RepositoryRefObservationRequest, WorkingDirectoryRepositoryAccessRequest, @@ -38,9 +35,8 @@ use crate::worker_source::{ use crate::working_directory::{ WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer, }; +use crate::workspace_request::{RuntimeWorkspaceRequest, RuntimeWorkspaceRequestClient}; use async_trait::async_trait; -#[cfg(feature = "http-server")] -use futures::StreamExt; #[cfg(test)] use protocol::WorkerStatus; use protocol::{Event, Method, Segment, WorkerCommandEnvelope}; @@ -323,7 +319,7 @@ pub struct ProfileRuntimeWorkerFactory { prompt_projection_cache: Arc, runtime_id: Option, worker_mutation_identity: Option, - runtime_request_audience: Option, + workspace_request_clients: Arc>, embedded_worker_mutation_dispatcher: Option>, controller_transport: WorkerControllerTransport, } @@ -339,7 +335,7 @@ impl ProfileRuntimeWorkerFactory { prompt_projection_cache: Arc::new(WorkspacePromptProjectionCache::default()), runtime_id: None, worker_mutation_identity: None, - runtime_request_audience: None, + workspace_request_clients: Arc::new(HashMap::new()), embedded_worker_mutation_dispatcher: None, controller_transport: WorkerControllerTransport::UnixSocket, } @@ -360,14 +356,10 @@ impl ProfileRuntimeWorkerFactory { self } - pub fn with_runtime_request_identity( - mut self, - identity: RuntimeIdentityMaterial, - audience: impl Into, - ) -> Self { - self.runtime_id = Some(identity.identity_id.clone()); - self.worker_mutation_identity = Some(identity); - self.runtime_request_audience = Some(audience.into()); + pub fn with_workspace_request_client(mut self, client: RuntimeWorkspaceRequestClient) -> Self { + self.runtime_id = Some(client.runtime_id().to_string()); + Arc::make_mut(&mut self.workspace_request_clients) + .insert(client.workspace_id().to_string(), client); self } @@ -550,7 +542,7 @@ impl RuntimeWorkspaceBackendRef { worker_ref: &WorkerRef, workspace_scope: Option<&crate::runtime::RuntimeWorkspaceScope>, mutation_identity: Option<&RuntimeIdentityMaterial>, - runtime_request_audience: Option<&str>, + workspace_request_client: Option<&RuntimeWorkspaceRequestClient>, embedded_dispatcher: Option<&Arc>, prompt_projection_cache: Option>, ) -> WorkerWorkspaceContext { @@ -561,28 +553,33 @@ impl RuntimeWorkspaceBackendRef { base_url, runtime_id, } => { - let mut client = RuntimeOwnedWorkspaceClient::new( - workspace_id.clone(), - base_url.clone(), - runtime_id.clone(), - worker_ref.worker_id.to_string(), - ); + let mut client = workspace_request_client + .cloned() + .map(|request_client| { + RuntimeOwnedWorkspaceClient::from_request_client( + request_client, + worker_ref.worker_id.to_string(), + ) + }) + .unwrap_or_else(|| { + RuntimeOwnedWorkspaceClient::new( + workspace_id.clone(), + base_url.clone(), + runtime_id.clone(), + worker_ref.worker_id.to_string(), + ) + }); if let Some(cache) = prompt_projection_cache { client = client.with_prompt_projection_cache(cache); } - if let Some(identity) = mutation_identity { - let audience = runtime_request_audience - .or_else(|| workspace_scope.map(|scope| scope.server_id.as_str())); - if let Some(audience) = audience { - client = client.with_runtime_request_source(identity, audience.to_owned()); - } - } - if let (Some(scope), Some(identity)) = (workspace_scope, mutation_identity) { + if let (Some(scope), Some(identity), Some(request_client)) = + (workspace_scope, mutation_identity, workspace_request_client) + { client = client.with_worker_remove(RuntimeWorkerMutationForwarder::remote( identity, scope.clone(), worker_ref.worker_id.to_string(), - base_url.clone(), + request_client.clone(), )); } else if let (Some(scope), Some(dispatcher)) = (workspace_scope, embedded_dispatcher) @@ -606,10 +603,18 @@ impl RuntimeWorkspaceBackendRef { #[cfg(feature = "http-server")] async fn fetch_workspace_config_http( request: &WorkspaceConfigFetchRequest, - identity: Option<&RuntimeIdentityMaterial>, - audience: Option<&str>, + client: &RuntimeWorkspaceRequestClient, ) -> Result { - let mut url = reqwest::Url::parse(&request.workspace_api.base_url) + if !client.matches_workspace( + &request.workspace_api.workspace_id, + &request.workspace_api.base_url, + ) { + return Err(format!( + "Workspace request route does not match Workspace Config source: workspace={} base_url={}", + request.workspace_api.workspace_id, request.workspace_api.base_url + )); + } + let mut url = reqwest::Url::parse(client.base_url()) .map_err(|error| format!("Workspace API base URL is invalid: {error}"))?; url.set_path(&format!( "/api/w/{}/runtime-config", @@ -620,79 +625,48 @@ async fn fetch_workspace_config_http( | crate::catalog::ProfileSelector::Named(value) => value.clone(), }; url.query_pairs_mut().append_pair("profile", &profile); - - let path = url.path().to_owned(); - let request_target = match url.query() { - Some(query) => format!("{path}?{query}"), - None => path.clone(), - }; - let client = reqwest::Client::builder() - .connect_timeout(Duration::from_secs(3)) - .timeout(WORKSPACE_CONFIG_HTTP_TIMEOUT) - .build() - .map_err(|error| format!("failed to build Workspace Config HTTP client: {error}"))?; - let mut http_request = client.get(url); - if let Some(identity) = identity { - let audience = audience - .ok_or_else(|| "Workspace Config request proof audience is unavailable".to_owned())?; - let proof = RuntimeRequestSourceSigner::from_identity(identity) - .issue( - audience, - &request.workspace_api.workspace_id, - None, - BACKEND_RESOURCE_FETCH_PERMISSION, - "GET", - &request_target, - b"", - i64::try_from(unix_now_seconds()).unwrap_or(i64::MAX), - 30, - ) - .map_err(|error| error.to_string())?; - http_request = http_request.header(RUNTIME_REQUEST_SOURCE_PROOF_HEADER, proof); - } + let mut headers = reqwest::header::HeaderMap::new(); if let Some(cached) = request.cached.as_ref() { - http_request = http_request.header( + headers.insert( reqwest::header::IF_NONE_MATCH, - workspace_config_etag(&cached.digest), + reqwest::header::HeaderValue::from_str(&workspace_config_etag(&cached.digest)) + .map_err(|error| format!("Workspace Config ETag is invalid: {error}"))?, ); } - - let response = http_request - .send() + let mut path_and_query = url.path().to_string(); + if let Some(query) = url.query() { + path_and_query.push('?'); + path_and_query.push_str(query); + } + let response = client + .execute(RuntimeWorkspaceRequest { + method: reqwest::Method::GET, + path_and_query, + body: Vec::new(), + headers, + permission: BACKEND_RESOURCE_FETCH_PERMISSION.to_string(), + worker_id: None, + timeout: Some(WORKSPACE_CONFIG_HTTP_TIMEOUT), + max_response_bytes: MAX_WORKSPACE_CONFIG_RESPONSE_BYTES, + }) .await .map_err(|error| format!("failed to fetch latest Workspace Config: {error}"))?; - if response.status() == reqwest::StatusCode::NOT_MODIFIED { + if response.status == reqwest::StatusCode::NOT_MODIFIED { return Ok(WorkspaceConfigFetchResult::NotModified); } - if !response.status().is_success() { + if !response.status.is_success() { return Err(format!( "latest Workspace Config fetch failed with HTTP {}", - response.status() + response.status )); } - if response - .content_length() - .is_some_and(|size| size > MAX_WORKSPACE_CONFIG_RESPONSE_BYTES as u64) - { - return Err("latest Workspace Config response exceeds the size limit".to_string()); - } let response_etag = response - .headers() + .headers .get(reqwest::header::ETAG) .and_then(|value| value.to_str().ok()) .map(str::to_string) .ok_or_else(|| "latest Workspace Config response is missing its ETag".to_string())?; - let mut body = Vec::new(); - let mut stream = response.bytes_stream(); - while let Some(chunk) = stream.next().await { - let chunk = - chunk.map_err(|error| format!("failed to read latest Workspace Config: {error}"))?; - if body.len().saturating_add(chunk.len()) > MAX_WORKSPACE_CONFIG_RESPONSE_BYTES { - return Err("latest Workspace Config response exceeds the size limit".to_string()); - } - body.extend_from_slice(&chunk); - } - let bundle = serde_json::from_slice::(&body) + let bundle = serde_json::from_slice::(&response.body) .map_err(|error| format!("failed to decode latest Workspace Config: {error}"))?; let expected_etag = workspace_config_etag(&bundle.metadata.digest); if response_etag != expected_etag { @@ -704,10 +678,9 @@ async fn fetch_workspace_config_http( } #[cfg(not(feature = "http-server"))] -async fn fetch_workspace_config_http( +async fn fetch_workspace_config_http( _request: &WorkspaceConfigFetchRequest, - _identity: Option<&RuntimeIdentityMaterial>, - _audience: Option<&str>, + _client: &T, ) -> Result { Err("Workspace Config fetch requires the worker-runtime http-server feature".to_string()) } @@ -805,12 +778,16 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { &self, request: WorkspaceConfigFetchRequest, ) -> Result { - fetch_workspace_config_http( - &request, - self.worker_mutation_identity.as_ref(), - self.runtime_request_audience.as_deref(), - ) - .await + let client = self + .workspace_request_clients + .get(&request.workspace_api.workspace_id) + .ok_or_else(|| { + format!( + "Workspace request client is unavailable for workspace {}", + request.workspace_api.workspace_id + ) + })?; + fetch_workspace_config_http(&request, client).await } fn observe_workspace_prompt_projection( @@ -854,11 +831,16 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { .map(|api| api.workspace_id.clone()); let observation_grants = request.request.worker_observation_grants.clone(); let observation_enabled = request.request.worker_observation_enabled; + let workspace_request_client = request + .request + .workspace_api + .as_ref() + .and_then(|api| self.workspace_request_clients.get(&api.workspace_id)); let workspace_context = workspace_backend_ref.worker_context( &request.worker_ref, request.workspace_scope.as_ref(), self.worker_mutation_identity.as_ref(), - self.runtime_request_audience.as_deref(), + workspace_request_client, self.embedded_worker_mutation_dispatcher.as_ref(), Some(self.prompt_projection_cache.clone()), ); @@ -1039,11 +1021,16 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { .map(|api| api.workspace_id.clone()); let observation_grants = request.request.worker_observation_grants.clone(); let observation_enabled = request.request.worker_observation_enabled; + let workspace_request_client = request + .request + .workspace_api + .as_ref() + .and_then(|api| self.workspace_request_clients.get(&api.workspace_id)); let workspace_context = workspace_backend_ref.worker_context( &request.worker_ref, request.workspace_scope.as_ref(), self.worker_mutation_identity.as_ref(), - self.runtime_request_audience.as_deref(), + workspace_request_client, self.embedded_worker_mutation_dispatcher.as_ref(), Some(self.prompt_projection_cache.clone()), ); @@ -2241,7 +2228,7 @@ mod tests { use crate::catalog::{ ConfigBundleRef, CreateWorkerRequest, MaterializerKind, ProfileSelector, RepositorySelector, WorkingDirectoryClaim, WorkingDirectoryRepository, - WorkingDirectoryRequest, + WorkingDirectoryRequest, WorkspaceApiRef, }; use crate::execution::WorkerExecutionContext; use crate::identity::WorkerId; @@ -2257,6 +2244,115 @@ mod tests { use manifest::{Scope, WorkerManifest}; use session_store::{LogEntry, WorkerMetadataStore}; + #[test] + fn profile_factory_routes_workspace_requests_by_workspace_id() { + let profiles = tempfile::tempdir().unwrap(); + let identity = RuntimeIdentityMaterial::generate("runtime-a").unwrap(); + let factory = ProfileRuntimeWorkerFactory::new(profiles.path()) + .with_workspace_request_client( + RuntimeWorkspaceRequestClient::new( + "workspace-a", + "https://workspace-a.example.test", + "runtime-a", + ) + .with_runtime_request_source(&identity, "server-a"), + ) + .with_workspace_request_client( + RuntimeWorkspaceRequestClient::new( + "workspace-b", + "https://workspace-b.example.test", + "runtime-a", + ) + .with_runtime_request_source(&identity, "server-b"), + ); + + assert_eq!( + factory + .workspace_request_clients + .get("workspace-a") + .and_then(RuntimeWorkspaceRequestClient::audience), + Some("server-a") + ); + assert_eq!( + factory + .workspace_request_clients + .get("workspace-b") + .and_then(RuntimeWorkspaceRequestClient::audience), + Some("server-b") + ); + } + + #[tokio::test] + async fn workspace_config_refresh_uses_workspace_scoped_request_client() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut chunk = [0_u8; 1024]; + loop { + let read = stream.read(&mut chunk).await.unwrap(); + if read == 0 { + break; + } + request.extend_from_slice(&chunk[..read]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + stream + .write_all(b"HTTP/1.1 304 Not Modified\r\nConnection: close\r\n\r\n") + .await + .unwrap(); + String::from_utf8(request).unwrap() + }); + let identity = RuntimeIdentityMaterial::generate("runtime-a").unwrap(); + let base_url = format!("http://{address}"); + let client = + RuntimeWorkspaceRequestClient::new("workspace-b", base_url.clone(), "runtime-a") + .with_runtime_request_source(&identity, "server-b"); + let bundle = test_bundle(); + let bundle_ref = ConfigBundleRef { + id: bundle.metadata.id.clone(), + digest: bundle.metadata.digest.clone(), + }; + let request = WorkspaceConfigFetchRequest { + workspace_api: WorkspaceApiRef { + workspace_id: "workspace-b".to_string(), + base_url, + }, + profile: crate::catalog::ProfileSelector::Named("coder".to_string()), + expected: bundle_ref.clone(), + cached: Some(bundle_ref), + }; + + let result = fetch_workspace_config_http(&request, &client) + .await + .unwrap(); + assert!(matches!(result, WorkspaceConfigFetchResult::NotModified)); + let raw_request = server.await.unwrap(); + let proof = raw_request + .lines() + .find_map(|line| { + line.split_once(':').and_then(|(name, value)| { + name.eq_ignore_ascii_case(crate::auth::RUNTIME_REQUEST_SOURCE_PROOF_HEADER) + .then(|| value.trim().to_string()) + }) + }) + .unwrap(); + let claims = crate::auth::decode_runtime_request_source_claims(&proof).unwrap(); + assert_eq!(claims.aud, "server-b"); + assert_eq!(claims.workspace_id, "workspace-b"); + assert_eq!(claims.worker_id, None); + assert_eq!(claims.method, "GET"); + assert_eq!( + claims.path, + "/api/w/workspace-b/runtime-config?profile=coder" + ); + } + fn test_command() -> WorkerCommandEnvelope { WorkerCommandEnvelope { command_id: 1, diff --git a/crates/worker-runtime/src/worker_source.rs b/crates/worker-runtime/src/worker_source.rs index 1905e233..f57e3a9e 100644 --- a/crates/worker-runtime/src/worker_source.rs +++ b/crates/worker-runtime/src/worker_source.rs @@ -1,16 +1,16 @@ use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use crate::workspace_request::{RuntimeWorkspaceRequest, RuntimeWorkspaceRequestClient}; use worker::{ WorkspaceClient, WorkspaceClientError, WorkspacePromptCatalogResolution, WorkspacePromptProjection, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, }; use crate::auth::{ - RUNTIME_REQUEST_SOURCE_PROOF_HEADER, RuntimeAuthError, RuntimeIdentityMaterial, - RuntimeRequestSourceSigner, RuntimeWorkerMutationSourceSigner, WORKER_REMOVE_PERMISSION, - WORKSPACE_REQUEST_PERMISSION, WORKSPACE_WORKER_DISCOVERY_PERMISSION, WorkerMutationActorKind, - WorkerMutationOperation, WorkerMutationSourceClaims, new_token_id, + RuntimeAuthError, RuntimeIdentityMaterial, RuntimeWorkerMutationSourceSigner, + WORKER_REMOVE_PERMISSION, WORKSPACE_REQUEST_PERMISSION, WORKSPACE_WORKER_DISCOVERY_PERMISSION, + WorkerMutationActorKind, WorkerMutationOperation, WorkerMutationSourceClaims, new_token_id, }; use crate::runtime::RuntimeWorkspaceScope; use crate::worker_backend::WorkspacePromptProjectionCache; @@ -133,9 +133,7 @@ pub trait EmbeddedWorkerMutationDispatcher: Send + Sync { #[derive(Clone)] enum RuntimeWorkerMutationTransport { Remote { - base_url: String, - request_source_signer: RuntimeRequestSourceSigner, - request_source_audience: String, + request_client: RuntimeWorkspaceRequestClient, }, Embedded { dispatcher: Arc, @@ -155,17 +153,13 @@ impl RuntimeWorkerMutationForwarder { identity: &RuntimeIdentityMaterial, scope: RuntimeWorkspaceScope, source_worker_id: impl Into, - base_url: impl Into, + request_client: RuntimeWorkspaceRequestClient, ) -> Self { Self { authority: RuntimeWorkerMutationSourceAuthority::remote(identity), scope: scope.clone(), source_worker_id: source_worker_id.into(), - transport: RuntimeWorkerMutationTransport::Remote { - base_url: base_url.into().trim_end_matches('/').to_string(), - request_source_signer: RuntimeRequestSourceSigner::from_identity(identity), - request_source_audience: scope.server_id, - }, + transport: RuntimeWorkerMutationTransport::Remote { request_client }, } } @@ -201,18 +195,11 @@ impl RuntimeWorkerMutationForwarder { )?; match (&self.transport, proof) { ( - RuntimeWorkerMutationTransport::Remote { - base_url, - request_source_signer, - request_source_audience, - }, + RuntimeWorkerMutationTransport::Remote { request_client }, RuntimeOwnedWorkerMutationProof::Remote(token), ) => execute_remote_worker_remove_http(RemoteWorkerRemoveHttpRequest { - base_url: base_url.clone(), - workspace_id: self.scope.workspace_id.clone(), + request_client: request_client.clone(), source_worker_id: self.source_worker_id.clone(), - request_source_signer: request_source_signer.clone(), - request_source_audience: request_source_audience.clone(), token, target_runtime_id: target_runtime_id.to_string(), target_worker_id: target_worker_id.to_string(), @@ -233,11 +220,8 @@ impl RuntimeWorkerMutationForwarder { } struct RemoteWorkerRemoveHttpRequest { - base_url: String, - workspace_id: String, + request_client: RuntimeWorkspaceRequestClient, source_worker_id: String, - request_source_signer: RuntimeRequestSourceSigner, - request_source_audience: String, token: String, target_runtime_id: String, target_worker_id: String, @@ -270,54 +254,54 @@ fn execute_remote_worker_remove_http( fn execute_remote_worker_remove_http_blocking( request: RemoteWorkerRemoveHttpRequest, ) -> Result { - let path = format!("/api/w/{}/workers/remove", request.workspace_id); - let url = format!("{}{}", request.base_url, path); - let body = serde_json::to_string(&serde_json::json!({ + let path = format!( + "/api/w/{}/workers/remove", + request.request_client.workspace_id() + ); + let body = serde_json::to_vec(&serde_json::json!({ "target_runtime_id": request.target_runtime_id, "target_worker_id": request.target_worker_id, "reason": request.reason, })) .map_err(|error| RuntimeWorkerMutationForwardError::Transport(error.to_string()))?; - let request_source_proof = request.request_source_signer.issue( - &request.request_source_audience, - &request.workspace_id, - Some(&request.source_worker_id), - WORKSPACE_REQUEST_PERMISSION, - "POST", - &path, - body.as_bytes(), - i64::try_from(unix_now_seconds()).unwrap_or(i64::MAX), - 30, - )?; - let client = reqwest::blocking::Client::new(); - let response = client - .post(url) - .header(RUNTIME_REQUEST_SOURCE_PROOF_HEADER, request_source_proof) - .header( - crate::auth::WORKER_MUTATION_SOURCE_PROOF_HEADER, - request.token, - ) - .header(reqwest::header::CONTENT_TYPE, "application/json") - .body(body) - .send() + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert( + crate::auth::WORKER_MUTATION_SOURCE_PROOF_HEADER, + reqwest::header::HeaderValue::from_str(&request.token) + .map_err(|error| RuntimeWorkerMutationForwardError::Transport(error.to_string()))?, + ); + headers.insert( + reqwest::header::CONTENT_TYPE, + reqwest::header::HeaderValue::from_static("application/json"), + ); + let response = request + .request_client + .execute_blocking(RuntimeWorkspaceRequest { + method: reqwest::Method::POST, + path_and_query: path, + body, + headers, + permission: WORKSPACE_REQUEST_PERMISSION.to_string(), + worker_id: Some(request.source_worker_id), + timeout: Some(Duration::from_secs(5)), + max_response_bytes: 8 * 1024 * 1024, + }) .map_err(|error| RuntimeWorkerMutationForwardError::Transport(error.to_string()))?; - let status = response.status().as_u16(); - let body = response - .text() + let body = String::from_utf8(response.body) .map_err(|error| RuntimeWorkerMutationForwardError::Transport(error.to_string()))?; - Ok(WorkspaceResponse { status, body }) + Ok(WorkspaceResponse { + status: response.status.as_u16(), + body, + }) } #[derive(Clone)] pub struct RuntimeOwnedWorkspaceClient { workspace_id: String, - base_url: String, - runtime_id: String, + request_client: RuntimeWorkspaceRequestClient, worker_id: String, request_timeout: Option, worker_remove: Option, - request_source_signer: Option, - request_source_audience: Option, prompt_projection_cache: Option>, } @@ -328,15 +312,32 @@ impl RuntimeOwnedWorkspaceClient { runtime_id: impl Into, worker_id: impl Into, ) -> Self { + let workspace_id = workspace_id.into(); Self { - workspace_id: workspace_id.into(), - base_url: base_url.into().trim_end_matches('/').to_string(), - runtime_id: runtime_id.into(), + request_client: RuntimeWorkspaceRequestClient::new( + workspace_id.clone(), + base_url, + runtime_id, + ), + workspace_id, + worker_id: worker_id.into(), + request_timeout: None, + worker_remove: None, + prompt_projection_cache: None, + } + } + + pub(crate) fn from_request_client( + request_client: RuntimeWorkspaceRequestClient, + worker_id: impl Into, + ) -> Self { + let workspace_id = request_client.workspace_id().to_string(); + Self { + workspace_id, + request_client, worker_id: worker_id.into(), request_timeout: None, worker_remove: None, - request_source_signer: None, - request_source_audience: None, prompt_projection_cache: None, } } @@ -351,8 +352,9 @@ impl RuntimeOwnedWorkspaceClient { identity: &RuntimeIdentityMaterial, audience: impl Into, ) -> Self { - self.request_source_signer = Some(RuntimeRequestSourceSigner::from_identity(identity)); - self.request_source_audience = Some(audience.into()); + self.request_client = self + .request_client + .with_runtime_request_source(identity, audience); self } @@ -375,44 +377,43 @@ impl RuntimeOwnedWorkspaceClient { request: WorkspaceRequest, permission: &'static str, ) -> Result { - let base_url = self.base_url.clone(); - let workspace_id = self.workspace_id.clone(); - let runtime_id = self.runtime_id.clone(); - let worker_id = self.worker_id.clone(); - let request_source_signer = self.request_source_signer.clone(); - let request_source_audience = self.request_source_audience.clone(); - let request_timeout = self.request_timeout; - if tokio::runtime::Handle::try_current().is_ok() { - std::thread::spawn(move || { - execute_runtime_owned_workspace_http( - &base_url, - &workspace_id, - &runtime_id, - &worker_id, - request_source_signer.as_ref(), - request_source_audience.as_deref(), - request_timeout, - permission, - request, - ) - }) - .join() - .map_err(|_| { - WorkspaceClientError::Request("workspace request thread panicked".to_string()) - })? - } else { - execute_runtime_owned_workspace_http( - &self.base_url, - &self.workspace_id, - &self.runtime_id, - &self.worker_id, - self.request_source_signer.as_ref(), - self.request_source_audience.as_deref(), - self.request_timeout, - permission, - request, - ) + let method = match request.method { + WorkspaceRequestMethod::Get => reqwest::Method::GET, + WorkspaceRequestMethod::Post => reqwest::Method::POST, + WorkspaceRequestMethod::Put => reqwest::Method::PUT, + WorkspaceRequestMethod::Patch => reqwest::Method::PATCH, + WorkspaceRequestMethod::Delete => reqwest::Method::DELETE, + }; + let body = request.body.unwrap_or_default().into_bytes(); + let mut headers = reqwest::header::HeaderMap::new(); + if !body.is_empty() { + headers.insert( + reqwest::header::CONTENT_TYPE, + reqwest::header::HeaderValue::from_static("application/json"), + ); } + let request_label = format!("{method} {}", request.path); + let response = self + .request_client + .execute_blocking(RuntimeWorkspaceRequest { + method, + path_and_query: request.path, + body, + headers, + permission: permission.to_string(), + worker_id: Some(self.worker_id.clone()), + timeout: self.request_timeout, + max_response_bytes: 8 * 1024 * 1024, + }) + .map_err(|error| { + WorkspaceClientError::Request(format!("{request_label} failed: {error}")) + })?; + let body = String::from_utf8(response.body) + .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; + Ok(WorkspaceResponse { + status: response.status.as_u16(), + body, + }) } } @@ -420,8 +421,8 @@ impl std::fmt::Debug for RuntimeOwnedWorkspaceClient { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter .debug_struct("RuntimeOwnedWorkspaceClient") - .field("workspace_id", &self.workspace_id) - .field("base_url", &self.base_url) + .field("workspace_id", &self.request_client.workspace_id()) + .field("base_url", &self.request_client.base_url()) .field("source", &"Runtime-owned") .field( "worker_remove", @@ -433,7 +434,7 @@ impl std::fmt::Debug for RuntimeOwnedWorkspaceClient { impl WorkspaceClient for RuntimeOwnedWorkspaceClient { fn workspace_id(&self) -> Option<&str> { - Some(&self.workspace_id) + Some(self.request_client.workspace_id()) } fn kind(&self) -> &str { @@ -588,117 +589,6 @@ fn percent_encode_query(value: &str) -> String { encoded } -fn execute_runtime_owned_workspace_http( - base_url: &str, - workspace_id: &str, - runtime_id: &str, - worker_id: &str, - request_source_signer: Option<&RuntimeRequestSourceSigner>, - request_source_audience: Option<&str>, - request_timeout: Option, - permission: &'static str, - request: WorkspaceRequest, -) -> Result { - if !request.path.starts_with('/') || request.path.starts_with("//") { - return Err(WorkspaceClientError::InvalidPath(request.path)); - } - let url = format!("{base_url}{}", request.path); - let method = match request.method { - WorkspaceRequestMethod::Get => reqwest::Method::GET, - WorkspaceRequestMethod::Post => reqwest::Method::POST, - WorkspaceRequestMethod::Put => reqwest::Method::PUT, - WorkspaceRequestMethod::Patch => reqwest::Method::PATCH, - WorkspaceRequestMethod::Delete => reqwest::Method::DELETE, - }; - let client = reqwest::blocking::Client::builder() - .timeout(request_timeout) - .build() - .map_err(|error| { - WorkspaceClientError::Unavailable(format!( - "failed to build Workspace API HTTP client: {}", - reqwest_error_chain(&error) - )) - })?; - let request_label = format!("{method} {}", request.path); - let body = request.body.unwrap_or_default(); - let mut request_builder = client - .request(method.clone(), url) - .header("x-yoi-runtime-id", runtime_id) - .header("x-yoi-worker-id", worker_id); - if let Some(signer) = request_source_signer { - let audience = request_source_audience.ok_or_else(|| { - WorkspaceClientError::Request( - "runtime request proof audience is unavailable".to_owned(), - ) - })?; - let proof = signer - .issue( - audience, - workspace_id, - Some(worker_id), - permission, - method.as_str(), - &request.path, - body.as_bytes(), - i64::try_from(unix_now_seconds()).unwrap_or(i64::MAX), - 30, - ) - .map_err(|error| WorkspaceClientError::Request(error.to_string()))?; - request_builder = request_builder.header(RUNTIME_REQUEST_SOURCE_PROOF_HEADER, proof); - } - if !body.is_empty() { - request_builder = request_builder - .header(reqwest::header::CONTENT_TYPE, "application/json") - .body(body); - } - let response = request_builder - .send() - .map_err(|error| workspace_http_error(&request_label, "waiting for response", error))?; - let status = response.status().as_u16(); - let body = response - .text() - .map_err(|error| workspace_http_error(&request_label, "reading response body", error))?; - Ok(WorkspaceResponse { status, body }) -} - -fn workspace_http_error( - request_label: &str, - stage: &str, - error: reqwest::Error, -) -> WorkspaceClientError { - let details = reqwest_error_chain(&error); - if error.is_timeout() { - WorkspaceClientError::Request(format!( - "Workspace API {request_label} timed out while {stage}: {details}" - )) - } else if error.is_connect() { - WorkspaceClientError::Unavailable(format!( - "Workspace API {request_label} could not connect while {stage}: {details}" - )) - } else { - WorkspaceClientError::Request(format!( - "Workspace API {request_label} transport failed while {stage}: {details}" - )) - } -} - -fn reqwest_error_chain(error: &reqwest::Error) -> String { - let mut details = error.to_string(); - let mut source = std::error::Error::source(error); - for _ in 0..4 { - let Some(current) = source else { - break; - }; - let current_text = current.to_string(); - if !current_text.is_empty() && !details.ends_with(¤t_text) { - details.push_str(": "); - details.push_str(¤t_text); - } - source = std::error::Error::source(current); - } - details -} - #[derive(Debug, thiserror::Error)] pub enum RuntimeWorkerMutationForwardError { #[error(transparent)] @@ -722,9 +612,9 @@ fn unix_now_seconds() -> u64 { mod tests { use super::*; use crate::auth::{ - WorkerMutationSourceExpectation, decode_runtime_request_source_claims, - decode_worker_mutation_source_claims, request_body_digest, - verify_worker_mutation_source_proof, + RUNTIME_REQUEST_SOURCE_PROOF_HEADER, WorkerMutationSourceExpectation, + decode_runtime_request_source_claims, decode_worker_mutation_source_claims, + request_body_digest, verify_worker_mutation_source_proof, }; #[test] @@ -1132,7 +1022,12 @@ mod tests { &identity, scope, "worker-source", - format!("http://{address}"), + RuntimeWorkspaceRequestClient::new( + "workspace-a", + format!("http://{address}"), + "runtime-a", + ) + .with_runtime_request_source(&identity, "server-a"), ); let response = forwarder .execute_worker_remove("runtime-target", "worker-target", "retire obsolete Worker") diff --git a/crates/worker-runtime/src/workspace_request.rs b/crates/worker-runtime/src/workspace_request.rs new file mode 100644 index 00000000..507e69c7 --- /dev/null +++ b/crates/worker-runtime/src/workspace_request.rs @@ -0,0 +1,360 @@ +use std::error::Error as _; +use std::io::Read; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use futures::StreamExt; +use reqwest::header::HeaderMap; +use thiserror::Error; + +use crate::auth::{ + RUNTIME_REQUEST_SOURCE_PROOF_HEADER, RuntimeAuthError, RuntimeIdentityMaterial, + RuntimeRequestSourceSigner, +}; + +const DEFAULT_REQUEST_PROOF_TTL_SECONDS: u64 = 60; +const RUNTIME_ID_HEADER: &str = "x-yoi-runtime-id"; +const WORKER_ID_HEADER: &str = "x-yoi-worker-id"; + +#[derive(Clone, Debug)] +pub struct RuntimeWorkspaceRequestClient { + workspace_id: String, + base_url: String, + runtime_id: String, + request_source: Option<(RuntimeRequestSourceSigner, String)>, +} + +#[derive(Clone, Debug)] +pub(crate) struct RuntimeWorkspaceRequest { + pub method: reqwest::Method, + pub path_and_query: String, + pub body: Vec, + pub headers: HeaderMap, + pub permission: String, + pub worker_id: Option, + pub timeout: Option, + pub max_response_bytes: usize, +} + +#[derive(Debug)] +pub(crate) struct RuntimeWorkspaceResponse { + pub status: reqwest::StatusCode, + pub headers: HeaderMap, + pub body: Vec, +} + +#[derive(Debug, Error)] +pub(crate) enum RuntimeWorkspaceRequestError { + #[error("invalid Workspace request: {0}")] + InvalidRequest(String), + #[error("failed to sign Workspace request: {0}")] + Sign(#[from] RuntimeAuthError), + #[error("Workspace request failed: {message}")] + Transport { message: String, timeout: bool }, + #[error("Workspace response exceeded {max_response_bytes} bytes")] + ResponseTooLarge { max_response_bytes: usize }, +} + +impl RuntimeWorkspaceRequestError { + fn transport(error: reqwest::Error) -> Self { + let timeout = error.is_timeout(); + Self::Transport { + message: reqwest_error_chain(&error), + timeout, + } + } + + pub(crate) fn is_timeout(&self) -> bool { + matches!(self, Self::Transport { timeout: true, .. }) + } +} + +impl RuntimeWorkspaceRequestClient { + pub fn new( + workspace_id: impl Into, + base_url: impl Into, + runtime_id: impl Into, + ) -> Self { + Self { + workspace_id: workspace_id.into(), + base_url: base_url.into().trim_end_matches('/').to_string(), + runtime_id: runtime_id.into(), + request_source: None, + } + } + + pub fn with_runtime_request_source( + mut self, + identity: &RuntimeIdentityMaterial, + audience: impl Into, + ) -> Self { + self.request_source = Some(( + RuntimeRequestSourceSigner::from_identity(identity), + audience.into(), + )); + self + } + + pub fn workspace_id(&self) -> &str { + &self.workspace_id + } + + pub fn base_url(&self) -> &str { + &self.base_url + } + + pub fn runtime_id(&self) -> &str { + &self.runtime_id + } + + pub fn audience(&self) -> Option<&str> { + self.request_source + .as_ref() + .map(|(_, audience)| audience.as_str()) + } + + pub fn matches_workspace(&self, workspace_id: &str, base_url: &str) -> bool { + self.workspace_id == workspace_id + && self.base_url.trim_end_matches('/') == base_url.trim_end_matches('/') + } + + pub(crate) async fn execute( + &self, + request: RuntimeWorkspaceRequest, + ) -> Result { + let prepared = self.prepare(&request)?; + let mut client_builder = reqwest::Client::builder(); + if let Some(timeout) = request.timeout { + client_builder = client_builder.timeout(timeout); + } + let client = client_builder + .build() + .map_err(RuntimeWorkspaceRequestError::transport)?; + let mut builder = client + .request(request.method, prepared.url) + .headers(request.headers) + .header(RUNTIME_ID_HEADER, &self.runtime_id); + if let Some(worker_id) = request.worker_id.as_deref() { + builder = builder.header(WORKER_ID_HEADER, worker_id); + } + if let Some(proof) = prepared.proof { + builder = builder.header(RUNTIME_REQUEST_SOURCE_PROOF_HEADER, proof); + } + if !request.body.is_empty() { + builder = builder.body(request.body); + } + let response = builder + .send() + .await + .map_err(RuntimeWorkspaceRequestError::transport)?; + let status = response.status(); + let headers = response.headers().clone(); + if response + .content_length() + .is_some_and(|size| size > request.max_response_bytes as u64) + { + return Err(RuntimeWorkspaceRequestError::ResponseTooLarge { + max_response_bytes: request.max_response_bytes, + }); + } + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(RuntimeWorkspaceRequestError::transport)?; + if body.len().saturating_add(chunk.len()) > request.max_response_bytes { + return Err(RuntimeWorkspaceRequestError::ResponseTooLarge { + max_response_bytes: request.max_response_bytes, + }); + } + body.extend_from_slice(&chunk); + } + Ok(RuntimeWorkspaceResponse { + status, + headers, + body, + }) + } + + pub(crate) fn execute_blocking( + &self, + request: RuntimeWorkspaceRequest, + ) -> Result { + let client = self.clone(); + std::thread::spawn(move || client.execute_blocking_inner(request)) + .join() + .map_err(|_| RuntimeWorkspaceRequestError::Transport { + message: "Workspace request thread panicked".to_string(), + timeout: false, + })? + } + + fn execute_blocking_inner( + &self, + request: RuntimeWorkspaceRequest, + ) -> Result { + let prepared = self.prepare(&request)?; + let mut client_builder = reqwest::blocking::Client::builder(); + if let Some(timeout) = request.timeout { + client_builder = client_builder.timeout(timeout); + } + let client = client_builder + .build() + .map_err(RuntimeWorkspaceRequestError::transport)?; + let mut builder = client + .request(request.method, prepared.url) + .headers(request.headers) + .header(RUNTIME_ID_HEADER, &self.runtime_id); + if let Some(worker_id) = request.worker_id.as_deref() { + builder = builder.header(WORKER_ID_HEADER, worker_id); + } + if let Some(proof) = prepared.proof { + builder = builder.header(RUNTIME_REQUEST_SOURCE_PROOF_HEADER, proof); + } + if !request.body.is_empty() { + builder = builder.body(request.body); + } + let response = builder + .send() + .map_err(RuntimeWorkspaceRequestError::transport)?; + let status = response.status(); + let headers = response.headers().clone(); + if response + .content_length() + .is_some_and(|size| size > request.max_response_bytes as u64) + { + return Err(RuntimeWorkspaceRequestError::ResponseTooLarge { + max_response_bytes: request.max_response_bytes, + }); + } + let limit = u64::try_from(request.max_response_bytes) + .unwrap_or(u64::MAX) + .saturating_add(1); + let mut body = Vec::new(); + response + .take(limit) + .read_to_end(&mut body) + .map_err(|error| RuntimeWorkspaceRequestError::Transport { + message: error.to_string(), + timeout: false, + })?; + if body.len() > request.max_response_bytes { + return Err(RuntimeWorkspaceRequestError::ResponseTooLarge { + max_response_bytes: request.max_response_bytes, + }); + } + Ok(RuntimeWorkspaceResponse { + status, + headers, + body, + }) + } + + fn prepare( + &self, + request: &RuntimeWorkspaceRequest, + ) -> Result { + if !request.path_and_query.starts_with('/') || request.path_and_query.starts_with("//") { + return Err(RuntimeWorkspaceRequestError::InvalidRequest( + "path must start with '/'".to_string(), + )); + } + let url = reqwest::Url::parse(&format!("{}{}", self.base_url, request.path_and_query)) + .map_err(|error| RuntimeWorkspaceRequestError::InvalidRequest(error.to_string()))?; + let mut request_target = url.path().to_string(); + if let Some(query) = url.query() { + request_target.push('?'); + request_target.push_str(query); + } + let proof = self + .request_source + .as_ref() + .map(|(signer, audience)| { + signer.issue( + audience, + &self.workspace_id, + request.worker_id.as_deref(), + &request.permission, + request.method.as_str(), + &request_target, + &request.body, + unix_now_seconds(), + DEFAULT_REQUEST_PROOF_TTL_SECONDS, + ) + }) + .transpose()?; + Ok(PreparedRuntimeWorkspaceRequest { url, proof }) + } +} + +struct PreparedRuntimeWorkspaceRequest { + url: reqwest::Url, + proof: Option, +} + +fn reqwest_error_chain(error: &reqwest::Error) -> String { + let mut message = error.to_string(); + let mut source = error.source(); + while let Some(error) = source { + message.push_str(": "); + message.push_str(&error.to_string()); + source = error.source(); + } + message +} + +fn unix_now_seconds() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| i64::try_from(duration.as_secs()).unwrap_or(i64::MAX)) + .unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::auth::decode_runtime_request_source_claims; + + #[test] + fn route_issues_workspace_scoped_request_proof() { + let identity = RuntimeIdentityMaterial::generate("runtime-a").unwrap(); + let client = RuntimeWorkspaceRequestClient::new( + "workspace-a", + "https://workspace.example.test/", + "runtime-a", + ) + .with_runtime_request_source(&identity, "workspace-server-a"); + let request = RuntimeWorkspaceRequest { + method: reqwest::Method::GET, + path_and_query: "/api/w/workspace-a/runtime-config?profile=coder".to_string(), + body: Vec::new(), + headers: HeaderMap::new(), + permission: "backend.resource.fetch".to_string(), + worker_id: None, + timeout: Some(Duration::from_secs(5)), + max_response_bytes: 1024, + }; + + let prepared = client.prepare(&request).unwrap(); + let claims = decode_runtime_request_source_claims(&prepared.proof.unwrap()).unwrap(); + assert_eq!(claims.aud, "workspace-server-a"); + assert_eq!(claims.workspace_id, "workspace-a"); + assert_eq!(claims.worker_id, None); + assert_eq!(claims.method, "GET"); + assert_eq!( + claims.path, + "/api/w/workspace-a/runtime-config?profile=coder" + ); + } + + #[test] + fn route_matches_only_its_workspace_and_backend() { + let client = RuntimeWorkspaceRequestClient::new( + "workspace-a", + "https://workspace.example.test/", + "runtime-a", + ); + + assert!(client.matches_workspace("workspace-a", "https://workspace.example.test")); + assert!(!client.matches_workspace("workspace-b", "https://workspace.example.test")); + assert!(!client.matches_workspace("workspace-a", "https://other.example.test")); + } +} diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index b3f26088..78240f19 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -2141,7 +2141,14 @@ impl WorkspaceApi { EMBEDDED_RUNTIME_ID, worker_remove_dispatcher.clone(), ) - .with_runtime_request_identity(embedded_identity, embedded_request_audience) + .with_workspace_request_client( + worker_runtime::workspace_request::RuntimeWorkspaceRequestClient::new( + config.workspace_id.clone(), + embedded_request_audience.clone(), + EMBEDDED_RUNTIME_ID, + ) + .with_runtime_request_source(&embedded_identity, embedded_request_audience), + ) .with_runtime_store_dir(config.embedded_runtime_store_root.clone()) .with_controller_transport(worker::WorkerControllerTransport::InProcess) .with_resource_client(Arc::new(resource_broker.clone())), @@ -27993,7 +28000,14 @@ mod tests { format!("http://{addr}/api/runtime/v1/workspaces/{TEST_WORKSPACE_ID}/resources/fetch"), None, ) - .with_runtime_request_source(&identity, "server-test"); + .with_workspace_request_client( + worker_runtime::workspace_request::RuntimeWorkspaceRequestClient::new( + TEST_WORKSPACE_ID, + format!("http://{addr}"), + runtime_id, + ) + .with_runtime_request_source(&identity, "server-test"), + ); let response = client .fetch_resource(worker_runtime::resource::BackendResourceFetchRequest {