fix: fence attachment grants and lifecycle cleanup
This commit is contained in:
@@ -393,6 +393,7 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
|
||||
_file_name: &str,
|
||||
_media_type: &str,
|
||||
_content: &[u8],
|
||||
_context: Option<&session_store::UploadedFileUploadContext>,
|
||||
) -> Result<UploadedFileRef, WorkerExecutionResult> {
|
||||
Err(WorkerExecutionResult::unsupported(
|
||||
WorkerExecutionOperation::UploadFile,
|
||||
@@ -546,9 +547,10 @@ impl WorkerExecutionBackendRef {
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
context: Option<&session_store::UploadedFileUploadContext>,
|
||||
) -> Result<UploadedFileRef, WorkerExecutionResult> {
|
||||
self.backend
|
||||
.upload_file(handle, file_name, media_type, content)
|
||||
.upload_file(handle, file_name, media_type, content, context)
|
||||
}
|
||||
|
||||
pub(crate) fn delete_uploaded_file(
|
||||
|
||||
@@ -390,6 +390,16 @@ pub struct RuntimeHttpWorkerInputResponse {
|
||||
pub struct RuntimeHttpUploadFileQuery {
|
||||
pub file_name: String,
|
||||
pub media_type: String,
|
||||
#[serde(default)]
|
||||
pub upload_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub principal_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub workspace_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub runtime_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub owner_worker_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
@@ -1455,14 +1465,74 @@ async fn upload_worker_file(
|
||||
body: Bytes,
|
||||
) -> RestResult<RuntimeHttpUploadedFileResponse> {
|
||||
let worker_ref = worker_ref_for(&state.runtime, worker_id)?;
|
||||
let context = match (
|
||||
query.upload_id,
|
||||
query.principal_id,
|
||||
query.workspace_id,
|
||||
query.runtime_id,
|
||||
query.owner_worker_id,
|
||||
) {
|
||||
(None, None, None, None, None) => None,
|
||||
(
|
||||
Some(upload_id),
|
||||
Some(principal_id),
|
||||
Some(workspace_id),
|
||||
Some(runtime_id),
|
||||
Some(owner_worker_id),
|
||||
) => {
|
||||
if owner_worker_id != worker_ref.worker_id.to_string() {
|
||||
return Err(RuntimeHttpRestError::new(
|
||||
StatusCode::FORBIDDEN,
|
||||
"uploaded_file_owner_mismatch",
|
||||
"uploaded file context does not match the target Worker",
|
||||
));
|
||||
}
|
||||
Some(session_store::UploadedFileUploadContext {
|
||||
upload_id,
|
||||
principal_id,
|
||||
workspace_id,
|
||||
runtime_id,
|
||||
worker_id: owner_worker_id,
|
||||
})
|
||||
}
|
||||
_ => {
|
||||
return Err(RuntimeHttpRestError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"uploaded_file_context_incomplete",
|
||||
"uploaded file context fields must be provided together",
|
||||
));
|
||||
}
|
||||
};
|
||||
let file = match auth_workspace_scope(&state, auth.as_ref())? {
|
||||
Some(scope) => state.runtime.upload_worker_file_scoped(
|
||||
&scope,
|
||||
&worker_ref,
|
||||
&query.file_name,
|
||||
&query.media_type,
|
||||
&body,
|
||||
),
|
||||
Some(scope) => {
|
||||
if context
|
||||
.as_ref()
|
||||
.is_some_and(|context| context.workspace_id != scope.workspace_id)
|
||||
{
|
||||
return Err(RuntimeHttpRestError::new(
|
||||
StatusCode::FORBIDDEN,
|
||||
"uploaded_file_workspace_mismatch",
|
||||
"uploaded file context does not match the authenticated Workspace",
|
||||
));
|
||||
}
|
||||
match context.as_ref() {
|
||||
Some(context) => state.runtime.upload_worker_file_with_context_scoped(
|
||||
&scope,
|
||||
&worker_ref,
|
||||
&query.file_name,
|
||||
&query.media_type,
|
||||
&body,
|
||||
context,
|
||||
),
|
||||
None => state.runtime.upload_worker_file_scoped(
|
||||
&scope,
|
||||
&worker_ref,
|
||||
&query.file_name,
|
||||
&query.media_type,
|
||||
&body,
|
||||
),
|
||||
}
|
||||
}
|
||||
None => state.runtime.upload_worker_file(
|
||||
&worker_ref,
|
||||
&query.file_name,
|
||||
|
||||
@@ -33,3 +33,4 @@ pub mod working_directory;
|
||||
pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions};
|
||||
pub use management::RuntimeOptions;
|
||||
pub use runtime::{Runtime, RuntimeWorkspaceScope};
|
||||
pub use session_store::UploadedFileUploadContext;
|
||||
|
||||
@@ -1224,6 +1224,41 @@ impl Runtime {
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
) -> Result<protocol::UploadedFileRef, RuntimeError> {
|
||||
self.upload_worker_file_inner(worker_ref, file_name, media_type, content, None)
|
||||
}
|
||||
|
||||
pub fn upload_worker_file_with_context_scoped(
|
||||
&self,
|
||||
scope: &RuntimeWorkspaceScope,
|
||||
worker_ref: &WorkerRef,
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
context: &session_store::UploadedFileUploadContext,
|
||||
) -> Result<protocol::UploadedFileRef, RuntimeError> {
|
||||
self.ensure_worker_in_workspace(scope, worker_ref)?;
|
||||
self.upload_worker_file_inner(worker_ref, file_name, media_type, content, Some(context))
|
||||
}
|
||||
|
||||
pub fn upload_worker_file_with_context(
|
||||
&self,
|
||||
worker_ref: &WorkerRef,
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
context: &session_store::UploadedFileUploadContext,
|
||||
) -> Result<protocol::UploadedFileRef, RuntimeError> {
|
||||
self.upload_worker_file_inner(worker_ref, file_name, media_type, content, Some(context))
|
||||
}
|
||||
|
||||
fn upload_worker_file_inner(
|
||||
&self,
|
||||
worker_ref: &WorkerRef,
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
context: Option<&session_store::UploadedFileUploadContext>,
|
||||
) -> Result<protocol::UploadedFileRef, RuntimeError> {
|
||||
let (backend, handle) = {
|
||||
let state = self.lock()?;
|
||||
@@ -1244,7 +1279,7 @@ impl Runtime {
|
||||
}
|
||||
};
|
||||
backend
|
||||
.upload_file(&handle, file_name, media_type, content)
|
||||
.upload_file(&handle, file_name, media_type, content, context)
|
||||
.map_err(|result| RuntimeError::WorkerExecutionRejected {
|
||||
worker_id: worker_ref.worker_id.clone(),
|
||||
operation: result.operation,
|
||||
|
||||
@@ -2031,19 +2031,24 @@ where
|
||||
file_name: &str,
|
||||
media_type: &str,
|
||||
content: &[u8],
|
||||
context: Option<&session_store::UploadedFileUploadContext>,
|
||||
) -> Result<protocol::UploadedFileRef, WorkerExecutionResult> {
|
||||
let (worker, _, _) = self.get_execution(handle).map_err(|mut result| {
|
||||
result.operation = WorkerExecutionOperation::UploadFile;
|
||||
result
|
||||
})?;
|
||||
worker
|
||||
.upload_file(file_name, media_type, content)
|
||||
.map_err(|error| {
|
||||
WorkerExecutionResult::rejected(
|
||||
WorkerExecutionOperation::UploadFile,
|
||||
format!("uploaded_file_rejected: {error}"),
|
||||
)
|
||||
})
|
||||
let uploaded = match context {
|
||||
Some(context) => {
|
||||
worker.upload_file_with_context(file_name, media_type, content, context)
|
||||
}
|
||||
None => worker.upload_file(file_name, media_type, content),
|
||||
};
|
||||
uploaded.map_err(|error| {
|
||||
WorkerExecutionResult::rejected(
|
||||
WorkerExecutionOperation::UploadFile,
|
||||
format!("uploaded_file_rejected: {error}"),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn delete_uploaded_file(
|
||||
@@ -2145,25 +2150,6 @@ where
|
||||
),
|
||||
);
|
||||
}
|
||||
let cleanup_handle = match self.workers.lock() {
|
||||
Ok(workers) => workers
|
||||
.get(handle.worker_ref())
|
||||
.map(|execution| execution.handle.clone()),
|
||||
Err(_) => {
|
||||
return WorkerExecutionResult::errored(
|
||||
WorkerExecutionOperation::Stop,
|
||||
"worker adapter registry lock is poisoned",
|
||||
);
|
||||
}
|
||||
};
|
||||
if let Some(worker) = cleanup_handle
|
||||
&& let Err(error) = worker.delete_uncommitted_uploaded_files()
|
||||
{
|
||||
return WorkerExecutionResult::errored(
|
||||
WorkerExecutionOperation::Stop,
|
||||
format!("uploaded_file_cleanup_failed: {error}"),
|
||||
);
|
||||
}
|
||||
let execution = match self.workers.lock() {
|
||||
Ok(mut workers) => workers.remove(handle.worker_ref()),
|
||||
Err(_) => {
|
||||
@@ -2179,6 +2165,7 @@ where
|
||||
"execution handle does not reference a live Worker",
|
||||
);
|
||||
};
|
||||
let artifact_cleanup = execution.handle.clone();
|
||||
let shutdown = execution.shutdown.clone();
|
||||
let result = self.send_method(
|
||||
WorkerExecutionOperation::Stop,
|
||||
@@ -2198,7 +2185,13 @@ where
|
||||
}
|
||||
Ok(())
|
||||
}) {
|
||||
Ok(()) => result,
|
||||
Ok(()) => match artifact_cleanup.delete_uncommitted_uploaded_files() {
|
||||
Ok(_) => result,
|
||||
Err(error) => WorkerExecutionResult::errored(
|
||||
WorkerExecutionOperation::Stop,
|
||||
format!("uploaded_file_cleanup_failed: {error}"),
|
||||
),
|
||||
},
|
||||
Err(message) => WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message),
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user