345 lines
13 KiB
Rust
345 lines
13 KiB
Rust
use async_trait::async_trait;
|
|
use chrono::{Duration, Utc};
|
|
use std::collections::HashMap;
|
|
use std::sync::{Arc, Mutex};
|
|
use uuid::Uuid;
|
|
use worker_runtime::identity::WorkerId;
|
|
use worker_runtime::profile_archive::ProfileSourceArchive;
|
|
use worker_runtime::resource::{
|
|
BackendResourceClient, BackendResourceError, BackendResourceFetchRequest,
|
|
BackendResourceFetchResponse, BackendResourceHandle, BackendResourceKind,
|
|
BackendResourceOperation, DEFAULT_PROFILE_SOURCE_ARCHIVE_MAX_BYTES,
|
|
PROFILE_SOURCE_ARCHIVE_CONTENT_TYPE, ResourceRedactionPolicy,
|
|
};
|
|
|
|
#[derive(Clone, Default)]
|
|
pub struct BackendResourceBroker {
|
|
resources: Arc<Mutex<HashMap<String, StoredResource>>>,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct StoredResource {
|
|
runtime_id: Option<String>,
|
|
worker_id: Option<String>,
|
|
handle: BackendResourceHandle,
|
|
archive: ProfileSourceArchive,
|
|
}
|
|
|
|
impl BackendResourceBroker {
|
|
pub fn issue_profile_source_archive_handle(
|
|
&self,
|
|
workspace_id: impl Into<String>,
|
|
runtime_id: Option<&str>,
|
|
worker_id: Option<&WorkerId>,
|
|
archive: ProfileSourceArchive,
|
|
) -> BackendResourceHandle {
|
|
let workspace_id = workspace_id.into();
|
|
let nonce = Uuid::now_v7().to_string();
|
|
let audit_correlation_id = format!("resource-fetch-{nonce}");
|
|
let expires_at = Utc::now() + Duration::minutes(15);
|
|
let handle = BackendResourceHandle {
|
|
kind: BackendResourceKind::ProfileSourceArchive,
|
|
workspace_id: workspace_id.clone(),
|
|
scope_id: Some("workspace-profile-source".to_string()),
|
|
runtime_id: runtime_id.map(|id| id.to_string()),
|
|
worker_id: worker_id.map(|id| id.to_string()),
|
|
resource_id: archive.reference.id.clone(),
|
|
digest: archive.reference.digest.clone(),
|
|
operation: BackendResourceOperation::FetchArchive,
|
|
expires_at_unix_seconds: expires_at.timestamp(),
|
|
nonce: nonce.clone(),
|
|
revision: archive.reference.digest.clone(),
|
|
generation: None,
|
|
max_bytes: DEFAULT_PROFILE_SOURCE_ARCHIVE_MAX_BYTES,
|
|
content_type: PROFILE_SOURCE_ARCHIVE_CONTENT_TYPE.to_string(),
|
|
redaction: ResourceRedactionPolicy::RuntimeInternalOnly,
|
|
audit_correlation_id,
|
|
profile_source_graph: Some(archive.reference.source_graph.clone()),
|
|
};
|
|
let stored = StoredResource {
|
|
runtime_id: runtime_id.map(|id| id.to_string()),
|
|
worker_id: worker_id.map(|id| id.to_string()),
|
|
handle: handle.clone(),
|
|
archive,
|
|
};
|
|
if let Ok(mut resources) = self.resources.lock() {
|
|
resources.insert(nonce, stored);
|
|
}
|
|
handle
|
|
}
|
|
|
|
pub fn profile_source_archive(
|
|
&self,
|
|
digest: &str,
|
|
) -> Option<worker_runtime::profile_archive::ProfileSourceArchive> {
|
|
self.resources
|
|
.lock()
|
|
.ok()?
|
|
.values()
|
|
.find(|resource| resource.handle.digest == digest)
|
|
.map(|resource| resource.archive.clone())
|
|
}
|
|
|
|
pub fn fetch_profile_source_archive(
|
|
&self,
|
|
request: BackendResourceFetchRequest,
|
|
) -> Result<BackendResourceFetchResponse, BackendResourceError> {
|
|
verify_handle_shape(&request.handle)?;
|
|
let stored = self
|
|
.resources
|
|
.lock()
|
|
.map_err(|_| BackendResourceError::Transport {
|
|
message: "resource broker lock poisoned".to_string(),
|
|
})?
|
|
.get(&request.handle.nonce)
|
|
.cloned()
|
|
.ok_or(BackendResourceError::MissingResource)?;
|
|
verify_handle_shape(&stored.handle)?;
|
|
if stored.handle.expires_at_unix_seconds < Utc::now().timestamp() {
|
|
return Err(BackendResourceError::Expired);
|
|
}
|
|
let actual_bytes = stored.archive.content.len() as u64;
|
|
if actual_bytes > stored.handle.max_bytes {
|
|
return Err(BackendResourceError::Oversized {
|
|
max_bytes: stored.handle.max_bytes,
|
|
actual_bytes,
|
|
});
|
|
}
|
|
if request.handle != stored.handle {
|
|
return Err(BackendResourceError::Unauthorized {
|
|
message: "resource handle does not match broker-issued handle".to_string(),
|
|
});
|
|
}
|
|
if let Some(expected_runtime_id) = stored.runtime_id.as_deref() {
|
|
if expected_runtime_id != request.runtime_id {
|
|
return Err(BackendResourceError::Unauthorized {
|
|
message: "runtime id does not match resource handle".to_string(),
|
|
});
|
|
}
|
|
}
|
|
if let Some(expected_worker_id) = stored.worker_id.as_deref() {
|
|
if Some(expected_worker_id) != request.worker_id.as_deref() {
|
|
return Err(BackendResourceError::Unauthorized {
|
|
message: "worker id does not match resource handle".to_string(),
|
|
});
|
|
}
|
|
}
|
|
Ok(BackendResourceFetchResponse {
|
|
kind: BackendResourceKind::ProfileSourceArchive,
|
|
resource_id: stored.archive.reference.id,
|
|
digest: stored.archive.reference.digest,
|
|
content_type: PROFILE_SOURCE_ARCHIVE_CONTENT_TYPE.to_string(),
|
|
bytes: stored.archive.content,
|
|
audit_correlation_id: request.audit_correlation_id,
|
|
})
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl BackendResourceClient for BackendResourceBroker {
|
|
async fn fetch_resource(
|
|
&self,
|
|
request: BackendResourceFetchRequest,
|
|
) -> Result<BackendResourceFetchResponse, BackendResourceError> {
|
|
self.fetch_profile_source_archive(request)
|
|
}
|
|
}
|
|
|
|
fn verify_handle_shape(handle: &BackendResourceHandle) -> Result<(), BackendResourceError> {
|
|
if handle.kind != BackendResourceKind::ProfileSourceArchive {
|
|
return Err(BackendResourceError::UnsupportedKind);
|
|
}
|
|
if handle.operation != BackendResourceOperation::FetchArchive {
|
|
return Err(BackendResourceError::Unauthorized {
|
|
message: "resource handle operation is not fetch_archive".to_string(),
|
|
});
|
|
}
|
|
if handle.content_type != PROFILE_SOURCE_ARCHIVE_CONTENT_TYPE {
|
|
return Err(BackendResourceError::ContentTypeMismatch {
|
|
expected: PROFILE_SOURCE_ARCHIVE_CONTENT_TYPE.to_string(),
|
|
actual: handle.content_type.clone(),
|
|
});
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::collections::BTreeMap;
|
|
use worker_runtime::identity::WorkerId;
|
|
use worker_runtime::profile_archive::{
|
|
ProfileSourceArchive, ProfileSourceArchiveRef, ProfileSourceGraphSummary, sha256_hex,
|
|
};
|
|
|
|
fn archive() -> ProfileSourceArchive {
|
|
let content = b"archive-content".to_vec();
|
|
let mut entrypoints = BTreeMap::new();
|
|
entrypoints.insert("default".to_string(), "profiles/default.dcdl".to_string());
|
|
ProfileSourceArchive {
|
|
reference: ProfileSourceArchiveRef {
|
|
id: "profile-source-archive:test".to_string(),
|
|
digest: sha256_hex(&content),
|
|
size_bytes: content.len() as u64,
|
|
source_graph: ProfileSourceGraphSummary {
|
|
entrypoints,
|
|
source_count: 1,
|
|
import_count: 0,
|
|
total_source_bytes: content.len() as u64,
|
|
},
|
|
},
|
|
content,
|
|
}
|
|
}
|
|
|
|
fn archive_with_len(len: usize) -> ProfileSourceArchive {
|
|
let content = vec![b'x'; len];
|
|
let mut entrypoints = BTreeMap::new();
|
|
entrypoints.insert("default".to_string(), "profiles/default.dcdl".to_string());
|
|
ProfileSourceArchive {
|
|
reference: ProfileSourceArchiveRef {
|
|
id: format!("profile-source-archive:test-{len}"),
|
|
digest: sha256_hex(&content),
|
|
size_bytes: content.len() as u64,
|
|
source_graph: ProfileSourceGraphSummary {
|
|
entrypoints,
|
|
source_count: 1,
|
|
import_count: 0,
|
|
total_source_bytes: content.len() as u64,
|
|
},
|
|
},
|
|
content,
|
|
}
|
|
}
|
|
|
|
fn request(
|
|
handle: BackendResourceHandle,
|
|
runtime_id: &str,
|
|
worker_id: Option<&WorkerId>,
|
|
) -> BackendResourceFetchRequest {
|
|
BackendResourceFetchRequest {
|
|
audit_correlation_id: handle.audit_correlation_id.clone(),
|
|
handle,
|
|
runtime_id: runtime_id.to_string(),
|
|
worker_id: worker_id.map(|id| id.to_string()),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn broker_issues_and_verifies_profile_source_archive_handles() {
|
|
let broker = BackendResourceBroker::default();
|
|
let runtime_id = "runtime-test";
|
|
let handle = broker.issue_profile_source_archive_handle(
|
|
"workspace-test",
|
|
Some(runtime_id),
|
|
None,
|
|
archive(),
|
|
);
|
|
let response = broker
|
|
.fetch_profile_source_archive(BackendResourceFetchRequest {
|
|
handle: handle.clone(),
|
|
runtime_id: runtime_id.to_string(),
|
|
worker_id: None,
|
|
audit_correlation_id: handle.audit_correlation_id.clone(),
|
|
})
|
|
.expect("fetch succeeds");
|
|
assert_eq!(response.digest, handle.digest);
|
|
assert_eq!(response.content_type, PROFILE_SOURCE_ARCHIVE_CONTENT_TYPE);
|
|
}
|
|
|
|
#[test]
|
|
fn broker_rejects_runtime_mismatch() {
|
|
let broker = BackendResourceBroker::default();
|
|
let runtime_a = "runtime-a";
|
|
let handle = broker.issue_profile_source_archive_handle(
|
|
"workspace-test",
|
|
Some(runtime_a),
|
|
None,
|
|
archive(),
|
|
);
|
|
let err = broker
|
|
.fetch_profile_source_archive(request(handle, "runtime-b", None))
|
|
.unwrap_err();
|
|
assert!(matches!(err, BackendResourceError::Unauthorized { .. }));
|
|
}
|
|
|
|
#[test]
|
|
fn broker_rejects_worker_mismatch() {
|
|
let broker = BackendResourceBroker::default();
|
|
let runtime_id = "runtime-test";
|
|
let worker_a = WorkerId::new(1);
|
|
let worker_b = WorkerId::new(2);
|
|
let handle = broker.issue_profile_source_archive_handle(
|
|
"workspace-test",
|
|
Some(runtime_id),
|
|
Some(&worker_a),
|
|
archive(),
|
|
);
|
|
let err = broker
|
|
.fetch_profile_source_archive(request(handle, &runtime_id, Some(&worker_b)))
|
|
.unwrap_err();
|
|
assert!(matches!(err, BackendResourceError::Unauthorized { .. }));
|
|
}
|
|
|
|
#[test]
|
|
fn broker_rejects_expiry_extension_from_request_handle() {
|
|
let broker = BackendResourceBroker::default();
|
|
let runtime_id = "runtime-test";
|
|
let handle = broker.issue_profile_source_archive_handle(
|
|
"workspace-test",
|
|
Some(runtime_id),
|
|
None,
|
|
archive(),
|
|
);
|
|
broker
|
|
.resources
|
|
.lock()
|
|
.unwrap()
|
|
.get_mut(&handle.nonce)
|
|
.unwrap()
|
|
.handle
|
|
.expires_at_unix_seconds = 1;
|
|
let mut extended = handle;
|
|
extended.expires_at_unix_seconds = 4_102_444_800;
|
|
let err = broker
|
|
.fetch_profile_source_archive(request(extended, &runtime_id, None))
|
|
.unwrap_err();
|
|
assert!(matches!(err, BackendResourceError::Expired));
|
|
}
|
|
|
|
#[test]
|
|
fn broker_rejects_policy_tampered_request_handle() {
|
|
let broker = BackendResourceBroker::default();
|
|
let runtime_id = "runtime-test";
|
|
let mut handle = broker.issue_profile_source_archive_handle(
|
|
"workspace-test",
|
|
Some(runtime_id),
|
|
None,
|
|
archive(),
|
|
);
|
|
handle.scope_id = Some("tampered-scope".to_string());
|
|
let err = broker
|
|
.fetch_profile_source_archive(request(handle, &runtime_id, None))
|
|
.unwrap_err();
|
|
assert!(matches!(err, BackendResourceError::Unauthorized { .. }));
|
|
}
|
|
|
|
#[test]
|
|
fn broker_uses_stored_max_bytes_when_request_handle_is_tampered() {
|
|
let broker = BackendResourceBroker::default();
|
|
let runtime_id = "runtime-test";
|
|
let archive = archive_with_len((DEFAULT_PROFILE_SOURCE_ARCHIVE_MAX_BYTES + 1) as usize);
|
|
let mut handle = broker.issue_profile_source_archive_handle(
|
|
"workspace-test",
|
|
Some(runtime_id),
|
|
None,
|
|
archive,
|
|
);
|
|
handle.max_bytes = DEFAULT_PROFILE_SOURCE_ARCHIVE_MAX_BYTES + 1024;
|
|
let err = broker
|
|
.fetch_profile_source_archive(request(handle, &runtime_id, None))
|
|
.unwrap_err();
|
|
assert!(matches!(err, BackendResourceError::Oversized { .. }));
|
|
}
|
|
}
|