Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: harden external workdir lifecycle

Browse Source
This commit is contained in:
Hare
2026-09-22 07:56:20 +09:00
parent ce9534f3a2
commit b5ee6b5804
12 changed files with 1113 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/fs-operation/src/glob.rs
+8
View File
@@ -15,6 +15,10 @@ pub fn run_glob(
if !root.is_absolute() {
return Err(FsError::RelativePath(root.to_path_buf()));
}
access.check_cancelled().map_err(|source| FsError::Io {
path: PathBuf::from(request.path.as_str()),
source,
})?;
let base_resolved = resolve_access_path(base).map_err(|error| FsError::Io {
path: PathBuf::from(request.path.as_str()),
source: error,
@@ -32,6 +36,10 @@ pub fn run_glob(
walker.hidden(false).follow_links(false);
let mut visited = 0_usize;
for entry in walker.build().flatten() {
access.check_cancelled().map_err(|source| FsError::Io {
path: PathBuf::from(request.path.as_str()),
source,
})?;
visited = visited.saturating_add(1);
if visited > crate::MAX_TRAVERSAL_ENTRIES {
return Err(FsError::InvalidArgument(format!(
crates/fs-operation/src/lib.rs
+87
View File
@@ -152,6 +152,12 @@ pub trait FsAccessPolicy: Send + Sync {
fn is_readable(&self, path: &Path) -> bool;
fn is_writable(&self, path: &Path) -> bool;
/// Cooperatively stop provider-owned filesystem work. Implementations may
/// use this for per-operation cancellation and deadlines.
fn check_cancelled(&self) -> std::io::Result<()> {
Ok(())
}
/// Open an already-authorized readable file. Capability providers override
/// this to bind path resolution and open into one root-confined operation.
fn open_read_file(&self, _logical: &Path, resolved: &Path) -> std::io::Result<std::fs::File> {
@@ -338,6 +344,87 @@ mod tests {
}
}
#[test]
fn cancellation_stops_a_read_already_in_progress() {
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex};
struct CancellableAccess {
root: PathBuf,
checks: AtomicUsize,
cancelled: AtomicBool,
state: Mutex<()>,
changed: Condvar,
}
impl FsAccessPolicy for CancellableAccess {
fn is_readable(&self, path: &Path) -> bool {
path.starts_with(&self.root)
}
fn is_writable(&self, _path: &Path) -> bool {
false
}
fn check_cancelled(&self) -> std::io::Result<()> {
if self.checks.fetch_add(1, Ordering::SeqCst) == 1 {
self.changed.notify_all();
let mut state = self.state.lock().unwrap();
while !self.cancelled.load(Ordering::Acquire) {
state = self.changed.wait(state).unwrap();
}
return Err(std::io::Error::new(
std::io::ErrorKind::Interrupted,
"cancelled",
));
}
Ok(())
}
}
let directory = tempfile::tempdir().unwrap();
let root = directory.path().canonicalize().unwrap();
std::fs::write(root.join("large.txt"), vec![b'x'; 128 * 1024]).unwrap();
let access = Arc::new(CancellableAccess {
root: root.clone(),
checks: AtomicUsize::new(0),
cancelled: AtomicBool::new(false),
state: Mutex::new(()),
changed: Condvar::new(),
});
let task_access = access.clone();
let task_root = root.clone();
let task = std::thread::spawn(move || {
run_read_bounded(
&task_root,
ReadRequest {
path: FsPath::new("large.txt").unwrap(),
offset: 0,
limit: usize::MAX,
max_bytes: 1024,
},
task_access.as_ref(),
BoundedReadLimits::new(256 * 1024, 1024).unwrap(),
)
});
let mut state = access.state.lock().unwrap();
while access.checks.load(Ordering::SeqCst) < 2 {
state = access.changed.wait(state).unwrap();
}
access.cancelled.store(true, Ordering::Release);
access.changed.notify_all();
drop(state);
let error = task
.join()
.unwrap()
.expect_err("read must stop on cancellation");
assert!(
matches!(error, FsError::Io { source, .. } if source.kind() == std::io::ErrorKind::Interrupted)
);
}
#[test]
fn logical_paths_reject_absolute_parent_and_backslash_forms() {
assert!(FsPath::new("src/lib.rs").is_ok());
crates/fs-operation/src/local.rs
+21
View File
@@ -20,6 +20,9 @@ pub fn run_stat(
access: &dyn FsAccessPolicy,
) -> Result<StatResult, FsError> {
let logical = request.path;
access
.check_cancelled()
.map_err(|error| map_io(&logical, error))?;
let path = resolve(root, &logical)?;
let resolved = resolve_access_path(&path).map_err(|error| map_io(&logical, error))?;
if !access.is_readable_paths(&path, &resolved) {
@@ -73,6 +76,9 @@ fn run_read_with_limits(
limits: Option<BoundedReadLimits>,
) -> Result<ReadResult, FsError> {
let logical = request.path;
access
.check_cancelled()
.map_err(|error| map_io(&logical, error))?;
let path = resolve(root, &logical)?;
let target = require_access(&path, &logical, access, false, false)?;
let file = access
@@ -111,6 +117,9 @@ fn run_read_with_limits(
let mut reader = BufReader::with_capacity(READ_BUFFER_BYTES, file);
let mut buffer = [0u8; READ_BUFFER_BYTES];
loop {
access
.check_cancelled()
.map_err(|error| map_io(&logical, error))?;
let read_limit = limits
.map(|limits| {
limits
@@ -190,6 +199,9 @@ pub fn run_write(
access: &dyn FsAccessPolicy,
) -> Result<WriteResult, FsError> {
let logical = request.path;
access
.check_cancelled()
.map_err(|error| map_io(&logical, error))?;
let path = resolve(root, &logical)?;
let created = !path.exists();
if path.exists() {
@@ -222,6 +234,9 @@ pub fn run_edit(
access: &dyn FsAccessPolicy,
) -> Result<EditResult, FsError> {
let logical = request.path;
access
.check_cancelled()
.map_err(|error| map_io(&logical, error))?;
let path = resolve(root, &logical)?;
let target = require_access(&path, &logical, access, true, false)?;
let bytes = fs::read(&target).map_err(|error| map_io(&logical, error))?;
@@ -262,6 +277,9 @@ pub fn run_list(
access: &dyn FsAccessPolicy,
) -> Result<ListResult, FsError> {
let logical = request.path;
access
.check_cancelled()
.map_err(|error| map_io(&logical, error))?;
let path = resolve(root, &logical)?;
let logical_base = path.clone();
let path = require_access(&path, &logical, access, false, true)?;
@@ -279,6 +297,9 @@ pub fn run_list(
.open_read_dir(&logical_base, &path)
.map_err(|error| map_io(&logical, error))?;
for (entry_index, entry) in read_dir.enumerate() {
access
.check_cancelled()
.map_err(|error| map_io(&logical, error))?;
if entry_index >= crate::MAX_TRAVERSAL_ENTRIES {
return Err(FsError::InvalidArgument(format!(
"directory traversal exceeds provider limit {}",
crates/fs-operation/src/search.rs
+25
View File
@@ -18,10 +18,12 @@ use crate::{
struct SourceBoundedReader<'a, R> {
inner: R,
remaining: &'a mut u64,
access: &'a dyn FsAccessPolicy,
}
impl<R: Read> Read for SourceBoundedReader<'_, R> {
fn read(&mut self, buffer: &mut [u8]) -> std::io::Result<usize> {
self.access.check_cancelled()?;
if buffer.is_empty() {
return Ok(0);
}
@@ -228,6 +230,9 @@ pub fn run_grep(
request: GrepRequest,
access: &dyn FsAccessPolicy,
) -> Result<GrepResult, FsError> {
access
.check_cancelled()
.map_err(|error| FsError::io(&base, error))?;
let p = GrepParams {
pattern: request.pattern,
path: Some(base.clone()),
@@ -372,6 +377,9 @@ pub fn run_grep(
let mut visited = 0_usize;
let mut source_bytes_remaining = crate::MAX_GREP_SOURCE_BYTES;
for entry in walker.build().flatten() {
access
.check_cancelled()
.map_err(|error| FsError::io(&base, error))?;
visited = visited.saturating_add(1);
if visited > crate::MAX_TRAVERSAL_ENTRIES {
return Err(FsError::InvalidArgument(format!(
@@ -445,6 +453,7 @@ fn scan_path(
let mut reader = SourceBoundedReader {
inner: reader,
remaining: source_bytes_remaining,
access,
};
match mode {
GrepOutputMode::FilesWithMatches => {
@@ -635,8 +644,22 @@ impl Sink for ContentSink<'_> {
#[cfg(test)]
mod bounded_reader_tests {
use std::io::{Cursor, Read as _};
use std::path::Path;
use super::SourceBoundedReader;
use crate::FsAccessPolicy;
struct Permit;
impl FsAccessPolicy for Permit {
fn is_readable(&self, _path: &Path) -> bool {
true
}
fn is_writable(&self, _path: &Path) -> bool {
true
}
}
#[test]
fn source_budget_is_shared_across_candidate_files() {
@@ -644,6 +667,7 @@ mod bounded_reader_tests {
let mut first = SourceBoundedReader {
inner: Cursor::new(b"ab"),
remaining: &mut remaining,
access: &Permit,
};
let mut output = Vec::new();
first.read_to_end(&mut output).unwrap();
@@ -653,6 +677,7 @@ mod bounded_reader_tests {
let mut second = SourceBoundedReader {
inner: Cursor::new(b"cd"),
remaining: &mut remaining,
access: &Permit,
};
let error = second.read_to_end(&mut output).unwrap_err();
assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
crates/workdir/src/external_local_tests.rs
+37 -2
View File
@@ -2,8 +2,8 @@ use crate::http::{
WorkdirSessionOperation, WorkdirSessionOperationResult, dispatch_workdir_session_operation,
};
use crate::{
BoundedReadLimits, CommandRequest, EditRequest, LocalWorkdirSession, ReadRequest, Workdir,
WorkdirError, WorkdirPath, WorkdirSession, WorkdirSessionCapabilities,
BoundedReadLimits, CommandRequest, EditRequest, ExternalWorkdirRoot, LocalWorkdirSession,
ReadRequest, Workdir, WorkdirError, WorkdirPath, WorkdirSession, WorkdirSessionCapabilities,
WorkdirSessionCapability, WriteRequest,
};
@@ -224,6 +224,41 @@ async fn external_local_provider_rejects_symlink_roots_and_traversal() {
assert!(matches!(error, WorkdirError::Denied(_)));
}
#[cfg(target_os = "linux")]
#[tokio::test]
async fn external_local_provider_keeps_pre_grant_root_when_path_changes_before_session_creation() {
let parent = tempfile::tempdir().unwrap();
let approved = parent.path().join("shared");
std::fs::create_dir(&approved).unwrap();
std::fs::write(approved.join("item.txt"), "approved").unwrap();
// This is the CLI's approval/open boundary. Remote grant creation happens
// only after this descriptor has been pinned.
let pinned = ExternalWorkdirRoot::pin(&approved).unwrap();
std::fs::rename(&approved, parent.path().join("approved-original")).unwrap();
std::fs::create_dir(&approved).unwrap();
std::fs::write(approved.join("item.txt"), "replacement").unwrap();
let session = LocalWorkdirSession::external_read_only_pinned(
Workdir::new("external-workdir-pre-grant-pin"),
pinned,
BoundedReadLimits::new(4096, 1024).unwrap(),
)
.unwrap();
let read = WorkdirSession::read(
&session,
ReadRequest {
path: WorkdirPath::new("item.txt").unwrap(),
offset: 0,
limit: 10,
max_bytes: 1024,
},
)
.await
.unwrap();
assert_eq!(read.bytes, b"approved");
}
#[cfg(target_os = "linux")]
#[tokio::test]
async fn external_local_provider_pins_the_approved_root_across_directory_replacement() {
crates/workdir/src/lib.rs
+2 -1
View File
@@ -30,7 +30,8 @@ pub use fs_operation::{
};
pub use http::dispatch_workdir_session_operation;
pub use local::{
LocalWorkdirSession, SymlinkInfo, WorkdirSessionResource, direct_symlink, first_symlink,
ExternalWorkdirRoot, LocalWorkdirSession, SymlinkInfo, WorkdirSessionResource, direct_symlink,
first_symlink,
};
pub use operation::*;
pub use router::{
crates/workdir/src/local.rs
+148 -34
View File
@@ -15,7 +15,7 @@ use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex as StdMutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use manifest::{Permission, Scope, ScopeConfig, ScopeRule, SharedScope, SymlinkPolicy};
@@ -208,6 +208,7 @@ struct ScopeAccess {
root: PathBuf,
pinned_root: Option<Arc<std::fs::File>>,
reject_symlinks: bool,
operation_guard: Option<OperationGuard>,
}
impl ScopeAccess {
@@ -216,12 +217,14 @@ impl ScopeAccess {
root: &Path,
pinned_root: Option<Arc<std::fs::File>>,
reject_symlinks: bool,
operation_guard: Option<OperationGuard>,
) -> Self {
Self {
scope,
root: root.to_path_buf(),
pinned_root,
reject_symlinks,
operation_guard,
}
}
@@ -258,6 +261,26 @@ impl fs_operation::FsAccessPolicy for ScopeAccess {
self.scope.is_writable(path)
}
fn check_cancelled(&self) -> std::io::Result<()> {
let Some(guard) = &self.operation_guard else {
return Ok(());
};
if guard.cancelled.load(Ordering::Acquire) {
return Err(std::io::Error::new(
std::io::ErrorKind::Interrupted,
"External Workdir operation was cancelled",
));
}
if Instant::now() >= guard.deadline {
guard.cancelled.store(true, Ordering::Release);
return Err(std::io::Error::new(
std::io::ErrorKind::TimedOut,
"External Workdir operation exceeded its provider deadline",
));
}
Ok(())
}
fn open_read_file(&self, _logical: &Path, resolved: &Path) -> std::io::Result<std::fs::File> {
if self.reject_symlinks {
self.open_confined(resolved)
@@ -381,6 +404,98 @@ impl Drop for LocalWorkdirSessionInner {
pub trait WorkdirSessionResource: Debug + Send + Sync {}
impl<T> WorkdirSessionResource for T where T: Debug + Send + Sync {}
/// An atomically validated, descriptor-pinned External Workdir root.
///
/// Pinning compares the selected path identity with the opened descriptor, so
/// a rename or replacement between approval and open fails closed. The
/// descriptor, rather than the pathname, remains the filesystem authority for
/// the lifetime of every session created from this root.
#[derive(Debug, Clone)]
pub struct ExternalWorkdirRoot {
canonical_path: PathBuf,
handle: Arc<std::fs::File>,
}
impl ExternalWorkdirRoot {
pub fn pin(directory: impl AsRef<Path>) -> Result<Self, WorkdirError> {
let directory = directory.as_ref();
if !directory.is_absolute() {
return Err(WorkdirError::InvalidArgument(
"External Workdir directory must be absolute".to_string(),
));
}
let selected_metadata = std::fs::symlink_metadata(directory)
.map_err(|error| WorkdirError::io(Path::new("<external-root>"), error))?;
if selected_metadata.file_type().is_symlink() {
return Err(WorkdirError::Denied(
"External Workdir root must not be a symbolic link".to_string(),
));
}
if !selected_metadata.is_dir() {
return Err(WorkdirError::InvalidArgument(
"External Workdir root must be a directory".to_string(),
));
}
let handle = Arc::new(
fs_operation::open_root_no_symlinks(directory)
.map_err(|error| WorkdirError::io(Path::new("<external-root>"), error))?,
);
let opened_metadata = handle
.metadata()
.map_err(|error| WorkdirError::io(Path::new("<external-root>"), error))?;
if !opened_metadata.is_dir() {
return Err(WorkdirError::InvalidArgument(
"External Workdir root must be a directory".to_string(),
));
}
#[cfg(unix)]
{
use std::os::unix::fs::MetadataExt;
if (selected_metadata.dev(), selected_metadata.ino())
!= (opened_metadata.dev(), opened_metadata.ino())
{
return Err(WorkdirError::Denied(
"External Workdir root changed while it was being pinned".to_string(),
));
}
}
#[cfg(target_os = "linux")]
let canonical_path = {
use std::os::fd::AsRawFd;
let path = std::fs::read_link(format!("/proc/self/fd/{}", handle.as_raw_fd()))
.map_err(|error| WorkdirError::io(Path::new("<external-root>"), error))?;
if !path.is_absolute() || path.to_string_lossy().ends_with(" (deleted)") {
return Err(WorkdirError::Denied(
"External Workdir root moved while it was being pinned".to_string(),
));
}
path
};
#[cfg(not(target_os = "linux"))]
let canonical_path = directory.to_path_buf();
Ok(Self {
canonical_path,
handle,
})
}
pub fn canonical_path(&self) -> &Path {
&self.canonical_path
}
}
#[derive(Debug, Clone)]
struct OperationGuard {
cancelled: Arc<AtomicBool>,
deadline: Instant,
}
/// Scope-aware filesystem handle. Clone-cheap (`Arc` inside).
///
/// The wrapped [`SharedScope`] is shared with every clone of this
@@ -391,6 +506,7 @@ impl<T> WorkdirSessionResource for T where T: Debug + Send + Sync {}
#[derive(Debug, Clone)]
pub struct LocalWorkdirSession {
inner: Arc<LocalWorkdirSessionInner>,
operation_guard: Option<OperationGuard>,
}
/// First symlink encountered while resolving a path.
@@ -495,45 +611,29 @@ impl LocalWorkdirSession {
/// Construct a strictly read-only provider session for an
/// authority-assigned Workdir and an operator-selected local directory.
///
/// The selected directory is canonicalized and fixed for the lifetime of
/// the session. A symlink selected as the root is rejected, and all
/// operation paths are confined below the canonical root without following
/// The selected directory is atomically validated and pinned for the
/// session lifetime. A symlink selected as the root is rejected, and all
/// operation paths are confined below the pinned root without following
/// symbolic links. Read source and response sizes are provider-bounded.
pub fn external_read_only(
workdir: Workdir,
directory: impl AsRef<Path>,
read_limits: BoundedReadLimits,
) -> Result<Self, WorkdirError> {
let root = ExternalWorkdirRoot::pin(directory)?;
Self::external_read_only_pinned(workdir, root, read_limits)
}
/// Construct a provider session from a root pinned before remote grant
/// creation, preventing any later pathname replacement from changing the
/// approved filesystem authority.
pub fn external_read_only_pinned(
workdir: Workdir,
pinned: ExternalWorkdirRoot,
read_limits: BoundedReadLimits,
) -> Result<Self, WorkdirError> {
read_limits.validate().map_err(WorkdirError::from)?;
let directory = directory.as_ref();
if !directory.is_absolute() {
return Err(WorkdirError::InvalidArgument(
"External Workdir directory must be absolute".to_string(),
));
}
let selected_metadata = std::fs::symlink_metadata(directory)
.map_err(|error| WorkdirError::io(Path::new("<external-root>"), error))?;
if selected_metadata.file_type().is_symlink() {
return Err(WorkdirError::Denied(
"External Workdir root must not be a symbolic link".to_string(),
));
}
let root = directory
.canonicalize()
.map_err(|error| WorkdirError::io(Path::new("<external-root>"), error))?;
if !root
.metadata()
.map_err(|error| WorkdirError::io(Path::new("<external-root>"), error))?
.is_dir()
{
return Err(WorkdirError::InvalidArgument(
"External Workdir root must be a directory".to_string(),
));
}
let pinned_root = Arc::new(
fs_operation::open_root_no_symlinks(&root)
.map_err(|error| WorkdirError::io(Path::new("<external-root>"), error))?,
);
let root = pinned.canonical_path;
let scope = Scope::from_config(&ScopeConfig {
allow: vec![ScopeRule {
target: root.clone(),
@@ -556,12 +656,24 @@ impl LocalWorkdirSession {
WorkdirSessionCapabilities::READ_ONLY,
BTreeMap::new(),
Vec::new(),
Some(pinned_root),
Some(pinned.handle),
Some(read_limits),
true,
))
}
/// Return an operation-local clone that cooperatively stops filesystem
/// work when cancelled or when its provider-side deadline expires.
pub fn with_operation_guard(&self, cancelled: Arc<AtomicBool>, deadline: Instant) -> Self {
Self {
inner: self.inner.clone(),
operation_guard: Some(OperationGuard {
cancelled,
deadline,
}),
}
}
fn materialized_bound_with_policy(
workdir: Workdir,
root: PathBuf,
@@ -592,6 +704,7 @@ impl LocalWorkdirSession {
command_environment,
resources: StdMutex::new(resources),
}),
operation_guard: None,
}
}
@@ -809,6 +922,7 @@ impl LocalWorkdirSession {
&self.inner.root,
self.inner.pinned_root.clone(),
self.inner.reject_symlinks,
self.operation_guard.clone(),
)
}
}
crates/worker-runtime/src/http_server.rs
+2 -43
View File
@@ -67,7 +67,7 @@ use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use tokio::net::TcpListener;
use workdir::{
CommandOutput, CommandStatus, WorkdirSessionHandle,
CommandOutput, CommandStatus, WorkdirSessionHandle, dispatch_workdir_session_operation,
http::{
OpenWorkdirSessionRequest, OpenWorkdirSessionResponse, WorkdirSessionId,
WorkdirSessionOperation, WorkdirSessionOperationRequest, WorkdirSessionOperationResult,
@@ -1030,41 +1030,6 @@ async fn run_workdir_session_operation(
let operation = request.operation;
let result = match operation {
WorkdirSessionOperation::AuthorizeScope(request) => {
session.authorize_scope_path(request).await?;
WorkdirSessionOperationResult::AuthorizeScope
}
WorkdirSessionOperation::ScopeRulesOverlap(request) => {
let overlaps = session.scope_rules_overlap(request).await?;
WorkdirSessionOperationResult::ScopeRulesOverlap { overlaps }
}
WorkdirSessionOperation::Stat(request) => {
WorkdirSessionOperationResult::Stat(session.stat(request).await?)
}
WorkdirSessionOperation::Read(request) => {
WorkdirSessionOperationResult::Read(session.read(request).await?)
}
WorkdirSessionOperation::Write(request) => {
WorkdirSessionOperationResult::Write(session.write(request).await?)
}
WorkdirSessionOperation::Edit(request) => {
WorkdirSessionOperationResult::Edit(session.edit(request).await?)
}
WorkdirSessionOperation::List(request) => {
WorkdirSessionOperationResult::List(session.list(request).await?)
}
WorkdirSessionOperation::Glob(request) => {
WorkdirSessionOperationResult::Glob(session.glob(request).await?)
}
WorkdirSessionOperation::Grep(request) => {
WorkdirSessionOperationResult::Grep(session.grep(request).await?)
}
WorkdirSessionOperation::CommandStart(request) => {
WorkdirSessionOperationResult::CommandStart(session.start_command(request).await?)
}
WorkdirSessionOperation::CommandStatus(handle) => {
WorkdirSessionOperationResult::CommandStatus(session.command_status(handle).await?)
}
WorkdirSessionOperation::CommandOutput(request) if request.wait => {
let cursor = request.cursor;
let output = match tokio::time::timeout(
@@ -1086,13 +1051,7 @@ async fn run_workdir_session_operation(
};
WorkdirSessionOperationResult::CommandOutput(output)
}
WorkdirSessionOperation::CommandOutput(request) => {
WorkdirSessionOperationResult::CommandOutput(session.command_output(request).await?)
}
WorkdirSessionOperation::CommandCancel(handle) => {
session.cancel_command(handle).await?;
WorkdirSessionOperationResult::CommandCancel
}
operation => dispatch_workdir_session_operation(session, operation).await?,
};
Ok(Json(result))
}
crates/workspace-server/src/latest_schema.sql
+3
View File
@@ -704,6 +704,9 @@ CREATE TABLE external_workdir_grants (
generation INTEGER NOT NULL CHECK (generation > 0),
status TEXT NOT NULL CHECK (status IN ('pending', 'online', 'offline', 'revoked', 'expired')),
updated_at TEXT NOT NULL,
cleanup_state TEXT NOT NULL DEFAULT 'not_required' CHECK (cleanup_state IN ('not_required', 'pending', 'retry_required', 'completed')),
cleanup_error TEXT,
cleanup_updated_at TEXT,
PRIMARY KEY (workspace_id, grant_id),
UNIQUE (workspace_id, workdir_id),
FOREIGN KEY (workspace_id) REFERENCES workspaces(workspace_id) ON DELETE CASCADE,
crates/workspace-server/src/server.rs
+467 -35
View File
@@ -683,7 +683,9 @@ enum ExternalProviderCommand {
Cancel {
operation_id: ExternalWorkdirOperationId,
},
Revoke,
Revoke {
reason: workdir::external::ExternalWorkdirRevokeReason,
},
}
#[derive(Debug)]
@@ -720,6 +722,10 @@ struct ExternalProviderConnection {
sender: tokio::sync::mpsc::Sender<ExternalProviderCommand>,
}
struct ExternalWorkdirExpiryTask {
handle: tokio::task::JoinHandle<()>,
}
#[derive(Debug, Default)]
struct ExternalProviderRegistry {
connections: HashMap<String, Arc<ExternalProviderConnection>>,
@@ -1101,6 +1107,7 @@ pub struct WorkspaceApi {
resource_broker: BackendResourceBroker,
workdir_sessions: Arc<Mutex<WorkdirSessionRegistry>>,
external_workdir_providers: Arc<Mutex<ExternalProviderRegistry>>,
external_workdir_expiry_tasks: Arc<Mutex<HashMap<String, ExternalWorkdirExpiryTask>>>,
workdir_session_locks: Arc<Mutex<HashMap<RuntimeWorkerRef, Arc<tokio::sync::Mutex<()>>>>>,
worker_remove_locks: Arc<Mutex<HashMap<RuntimeWorkerRef, Arc<tokio::sync::Mutex<()>>>>>,
workdir_remove_locks: Arc<Mutex<HashMap<String, Arc<std::sync::Mutex<()>>>>>,
@@ -3039,6 +3046,7 @@ impl WorkspaceApi {
resource_broker,
workdir_sessions: Arc::new(Mutex::new(WorkdirSessionRegistry::default())),
external_workdir_providers: Arc::new(Mutex::new(ExternalProviderRegistry::default())),
external_workdir_expiry_tasks: Arc::new(Mutex::new(HashMap::new())),
workdir_session_locks: Arc::new(Mutex::new(HashMap::new())),
worker_remove_locks: Arc::new(Mutex::new(HashMap::new())),
workdir_remove_locks: Arc::new(Mutex::new(HashMap::new())),
@@ -3056,10 +3064,18 @@ impl WorkspaceApi {
&Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true),
)?;
for workdir_id in expired_external_workdirs {
release_external_workdir_attachments(&api, &workdir_id)
cleanup_external_workdir_attachments(&api, &workdir_id)
.await
.map_err(|error| error.error)?;
}
for grant in api
.store
.list_external_workdir_grants(&api.config.workspace_id)?
.into_iter()
.filter(|grant| !matches!(grant.status.as_str(), "revoked" | "expired"))
{
schedule_external_workdir_expiry(&api, &grant)?;
}
recover_workdir_removals(&api)?;
recover_runtime_removals(&api).await;
Ok(api)
@@ -8996,6 +9012,7 @@ async fn open_current_worker_workdir_session_locked(
"expired",
&updated_at,
)?;
cancel_external_workdir_expiry(api, grant_id);
let connection = {
api.external_workdir_providers
.lock()
@@ -9005,7 +9022,9 @@ async fn open_current_worker_workdir_session_locked(
if let Some(connection) = connection {
let _ = connection
.sender
.send(ExternalProviderCommand::Revoke)
.send(ExternalProviderCommand::Revoke {
reason: workdir::external::ExternalWorkdirRevokeReason::Expired,
})
.await;
}
close_current_worker_command_sessions_for_alias_locked(api, worker, &link.alias)
@@ -11921,6 +11940,132 @@ async fn scoped_cleanup_runtime_working_directory(
.map_err(ApiError::from)
}
fn schedule_external_workdir_expiry(
api: &WorkspaceApi,
grant: &ExternalWorkdirGrantRecord,
) -> Result<()> {
let expires_at = chrono::DateTime::parse_from_rfc3339(&grant.expires_at)
.map_err(|_| Error::Store("External Workdir grant has invalid expiry".to_string()))?
.with_timezone(&Utc);
let grant_id = grant.grant_id.clone();
let task_api = api.clone();
let task_grant_id = grant_id.clone();
let handle = tokio::spawn(async move {
let delay = (expires_at - Utc::now()).to_std().unwrap_or_default();
tokio::time::sleep(delay).await;
loop {
match expire_external_workdir_grant(&task_api, &task_grant_id).await {
Ok(()) => break,
Err(error) => {
tracing::error!(grant_id = %task_grant_id, error = ?error, "Backend External Workdir expiry cleanup requires retry");
let retryable = task_api
.store
.get_external_workdir_grant(&task_api.config.workspace_id, &task_grant_id)
.ok()
.flatten()
.is_some_and(|grant| {
matches!(grant.status.as_str(), "revoked" | "expired")
});
if !retryable {
break;
}
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
}
}
}
});
let mut tasks = api
.external_workdir_expiry_tasks
.lock()
.expect("External Workdir expiry task registry poisoned");
if let Some(existing) = tasks.insert(grant_id, ExternalWorkdirExpiryTask { handle }) {
existing.handle.abort();
}
Ok(())
}
fn cancel_external_workdir_expiry(api: &WorkspaceApi, grant_id: &str) {
if let Some(task) = api
.external_workdir_expiry_tasks
.lock()
.expect("External Workdir expiry task registry poisoned")
.remove(grant_id)
{
task.handle.abort();
}
}
async fn expire_external_workdir_grant(api: &WorkspaceApi, grant_id: &str) -> ApiResult<()> {
loop {
let grant = api
.store
.get_external_workdir_grant(&api.config.workspace_id, grant_id)?
.ok_or_else(|| Error::InvalidInput("Unknown External Workdir grant".to_string()))?;
if matches!(grant.status.as_str(), "revoked" | "expired") {
break;
}
let expires_at = chrono::DateTime::parse_from_rfc3339(&grant.expires_at)
.map(|value| value.with_timezone(&Utc))
.map_err(|_| Error::Store("External Workdir grant has invalid expiry".to_string()))?;
if expires_at > Utc::now() {
return Ok(());
}
let updated_at = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true);
if !api.store.update_external_workdir_grant_state(
&grant.workspace_id,
&grant.grant_id,
&grant.provider_instance_id,
grant.generation,
"expired",
&updated_at,
)? {
// A reconnect may advance generation between the read and fence.
continue;
}
let connection = {
api.external_workdir_providers
.lock()
.expect("External Workdir provider registry poisoned")
.disconnect(&grant.grant_id, grant.generation)
};
if let Some(connection) = connection {
let _ = connection
.sender
.send(ExternalProviderCommand::Revoke {
reason: workdir::external::ExternalWorkdirRevokeReason::Expired,
})
.await;
}
api.store.record_external_workdir_audit(
&grant.workspace_id,
"system",
"workspace-backend",
"external_workdir_grant_expired",
&grant.grant_id,
None,
"expired",
&format!(
"workdir_id={} generation={}",
grant.workdir_id, grant.generation
),
&updated_at,
)?;
break;
}
cleanup_external_workdir_attachments_for_grant(api, grant_id).await
}
async fn cleanup_external_workdir_attachments_for_grant(
api: &WorkspaceApi,
grant_id: &str,
) -> ApiResult<()> {
let grant = api
.store
.get_external_workdir_grant(&api.config.workspace_id, grant_id)?
.ok_or_else(|| Error::InvalidInput("Unknown External Workdir grant".to_string()))?;
cleanup_external_workdir_attachments(api, &grant.workdir_id).await
}
async fn scoped_create_external_workdir_grant(
State(api): State<WorkspaceApi>,
AxumPath(path): AxumPath<ScopedWorkspacePath>,
@@ -11991,6 +12136,7 @@ async fn scoped_create_external_workdir_grant(
updated_at: grant.updated_at.clone(),
};
api.store.create_external_workdir_grant(&grant, &workdir)?;
schedule_external_workdir_expiry(&api, &grant)?;
Ok((
StatusCode::CREATED,
Json(external_workdir_grant_response(&grant)),
@@ -12042,7 +12188,8 @@ async fn scoped_external_workdir_provider_ws(
"expired",
&now.to_rfc3339_opts(SecondsFormat::Nanos, true),
)?;
release_external_workdir_attachments(&api, &grant.workdir_id).await?;
cancel_external_workdir_expiry(&api, &grant_id);
cleanup_external_workdir_attachments(&api, &grant.workdir_id).await?;
return Err(Error::InvalidInput("External Workdir grant has expired".to_string()).into());
}
Ok(ws
@@ -12208,10 +12355,6 @@ async fn serve_external_workdir_provider(
std::result::Result<WorkdirSessionOperationResult, WorkdirTransportError>,
>,
>::new();
let expiry = tokio::time::sleep_until(
tokio::time::Instant::now() + (expires_at - Utc::now()).to_std().unwrap_or_default(),
);
tokio::pin!(expiry);
let mut heartbeat = tokio::time::interval(std::time::Duration::from_secs(10));
heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let mut heartbeat_sequence = 0_u64;
@@ -12231,14 +12374,6 @@ async fn serve_external_workdir_provider(
break "offline";
}
}
_ = &mut expiry => {
let frame = ExternalWorkdirServerFrame::current(ExternalWorkdirServerMessage::Revoke {
generation: connection.generation,
reason: workdir::external::ExternalWorkdirRevokeReason::Expired,
});
let _ = send_external_workdir_server_frame(&mut socket, &frame).await;
break "expired";
}
command = commands.recv() => match command {
Some(ExternalProviderCommand::Operation { operation_id, operation, response }) => {
if pending.len() >= 16 {
@@ -12280,10 +12415,10 @@ async fn serve_external_workdir_provider(
break "offline";
}
}
Some(ExternalProviderCommand::Revoke) => {
Some(ExternalProviderCommand::Revoke { reason }) => {
let frame = ExternalWorkdirServerFrame::current(ExternalWorkdirServerMessage::Revoke {
generation: connection.generation,
reason: workdir::external::ExternalWorkdirRevokeReason::UserRequested,
reason,
});
let _ = send_external_workdir_server_frame(&mut socket, &frame).await;
break "revoked";
@@ -12383,7 +12518,7 @@ async fn serve_external_workdir_provider(
if state_updated
&& matches!(terminal_status, "expired" | "revoked")
&& let Err(error) =
release_external_workdir_attachments(&api, &initial_grant.workdir_id).await
cleanup_external_workdir_attachments(&api, &initial_grant.workdir_id).await
{
tracing::error!(
grant_id = %initial_grant.grant_id,
@@ -12428,11 +12563,12 @@ async fn scoped_get_external_workdir_grant(
grant.status = "expired".to_string();
grant.updated_at = updated_at;
}
cancel_external_workdir_expiry(&api, &grant_id);
}
// Terminal cleanup is retryable: a prior expiry/revoke may have committed
// its fence before Runtime/session projection completed.
if matches!(grant.status.as_str(), "expired" | "revoked") {
release_external_workdir_attachments(&api, &grant.workdir_id).await?;
cleanup_external_workdir_attachments(&api, &grant.workdir_id).await?;
}
Ok(Json(external_workdir_grant_response(&grant)))
}
@@ -12472,11 +12608,14 @@ async fn scoped_revoke_external_workdir_grant(
if let Some(connection) = connection {
let _ = connection
.sender
.send(ExternalProviderCommand::Revoke)
.send(ExternalProviderCommand::Revoke {
reason: workdir::external::ExternalWorkdirRevokeReason::UserRequested,
})
.await;
}
grant.status = "revoked".to_string();
grant.updated_at = updated_at;
cancel_external_workdir_expiry(&api, &grant_id);
api.store.record_external_workdir_audit(
&workspace_id,
"account",
@@ -12494,10 +12633,78 @@ async fn scoped_revoke_external_workdir_grant(
}
// Cleanup is deliberately retried for terminal grants. A prior revoke may
// have committed its fence before a Runtime/session projection failure.
release_external_workdir_attachments(&api, &grant.workdir_id).await?;
cleanup_external_workdir_attachments(&api, &grant.workdir_id).await?;
Ok(Json(external_workdir_grant_response(&grant)))
}
async fn cleanup_external_workdir_attachments(
api: &WorkspaceApi,
workdir_id: &str,
) -> ApiResult<()> {
let workdir = api
.store
.get_workdir_registry(&api.config.workspace_id, workdir_id)?
.ok_or_else(|| Error::InvalidInput("Unknown External Workdir".to_string()))?;
let grant_id = workdir.source.external_grant_id().ok_or_else(|| {
Error::InvalidInput("Workdir is not backed by an External grant".to_string())
})?;
let now = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true);
api.store.update_external_workdir_cleanup_state(
&api.config.workspace_id,
grant_id,
"pending",
None,
&now,
)?;
match release_external_workdir_attachments(api, workdir_id).await {
Ok(()) => {
let completed_at = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true);
api.store.update_external_workdir_cleanup_state(
&api.config.workspace_id,
grant_id,
"completed",
None,
&completed_at,
)?;
api.store.record_external_workdir_audit(
&api.config.workspace_id,
"system",
"workspace-backend",
"external_workdir_cleanup_completed",
grant_id,
None,
"completed",
&format!("workdir_id={workdir_id}"),
&completed_at,
)?;
Ok(())
}
Err(error) => {
let failed_at = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true);
api.store.update_external_workdir_cleanup_state(
&api.config.workspace_id,
grant_id,
"retry_required",
Some("attachment_release_failed"),
&failed_at,
)?;
api.store.record_external_workdir_audit(
&api.config.workspace_id,
"system",
"workspace-backend",
"external_workdir_cleanup_attention_required",
grant_id,
None,
"attention_required",
&format!("workdir_id={workdir_id} category=attachment_release_failed"),
&failed_at,
)?;
Err(error)
}
}
}
async fn release_external_workdir_attachments(
api: &WorkspaceApi,
workdir_id: &str,
@@ -16816,10 +17023,21 @@ fn finalize_spawned_worker_workdir_attachments(
.get_workdir_registry(&api.config.workspace_id, &summary.working_directory_id)
.map_err(ApiError::from),
)?;
preserve_workdir_identity_for_corrupted_summary(
&mut workdir_record,
previous_record.as_ref(),
);
if previous_record.as_ref().is_some_and(|record| {
matches!(record.source, WorkdirRegistrySource::ExternalGrant { .. })
}) {
// External Workdir identity and lifecycle are Backend/provider-owned;
// a Runtime attachment projection must not rewrite it as a
// Runtime repository materialization.
workdir_record = previous_record
.clone()
.expect("External source was checked above");
} else {
preserve_workdir_identity_for_corrupted_summary(
&mut workdir_record,
previous_record.as_ref(),
);
}
finalize_worker_spawn_stage(
api,
worker,
@@ -18978,12 +19196,19 @@ fn available_working_directory_summaries(
for worker in api.runtime.list_workers(limit).items {
let _ = sync_worker_observation(api, &worker);
}
let records = working_directory_summaries(api)?;
let records = api
.store
.list_workdir_registry(&api.config.workspace_id, 200)?;
let mut available = Vec::new();
for summary in records {
if summary.status != WorkingDirectoryStatusKind::Active
|| summary.cleanliness.as_deref() != Some("clean")
{
for record in records {
let summary = projected_workdir_summary_from_record(api, &record)?;
let source_is_available = match record.source {
WorkdirRegistrySource::Repository { .. } => {
summary.cleanliness.as_deref() == Some("clean")
}
WorkdirRegistrySource::ExternalGrant { .. } => true,
};
if summary.status != WorkingDirectoryStatusKind::Active || !source_is_available {
continue;
}
if summary.occupied_by.is_none() {
@@ -20238,10 +20463,7 @@ fn repository_access_request_for_workdir(
repository_id,
} = &record.source
else {
return Err(settings_bad_request(
"working_directory_provider_mismatch",
"External working directories do not require Repository access",
));
return Ok(None);
};
if record_runtime_id != runtime_id {
return Err(settings_bad_request(
@@ -20937,6 +21159,216 @@ mod tests {
));
}
#[tokio::test]
async fn backend_expiry_fences_offline_grant_and_persists_cleanup_completion() {
let dir = tempfile::tempdir().unwrap();
init_clean_git_workspace(dir.path());
let api = test_api(dir.path()).await;
let now = Utc::now();
let created_at = now.to_rfc3339_opts(SecondsFormat::Nanos, true);
let grant = ExternalWorkdirGrantRecord {
grant_id: "offline-expiry-grant".to_string(),
workspace_id: TEST_WORKSPACE_ID.to_string(),
workdir_id: "offline-expiry-workdir".to_string(),
provider_instance_id: "offline-expiry-provider".to_string(),
display_name: "Offline expiry".to_string(),
permissions: "read_only".to_string(),
created_by: test_browser_request_actor().account_id,
created_at: created_at.clone(),
expires_at: (now + Duration::milliseconds(50))
.to_rfc3339_opts(SecondsFormat::Nanos, true),
generation: 1,
status: "offline".to_string(),
updated_at: created_at.clone(),
};
api.store
.create_external_workdir_grant(
&grant,
&WorkdirRegistryRecord {
workspace_id: TEST_WORKSPACE_ID.to_string(),
workdir_id: grant.workdir_id.clone(),
display_name: Some(grant.display_name.clone()),
source: WorkdirRegistrySource::ExternalGrant {
grant_id: grant.grant_id.clone(),
},
creation_selector: None,
creation_ref: None,
creation_tree: None,
current_selector: None,
current_ref: None,
current_tree: None,
observed_at_epoch_seconds: None,
materialization_status: "unknown".to_string(),
cleanliness: "unknown".to_string(),
created_at: created_at.clone(),
updated_at: created_at,
},
)
.unwrap();
schedule_external_workdir_expiry(&api, &grant).unwrap();
tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
let current = api
.store
.get_external_workdir_grant(TEST_WORKSPACE_ID, &grant.grant_id)
.unwrap()
.unwrap();
if current.status == "expired" {
let cleanup_state: String = api
.config_store
.with_conn(|conn| {
conn.query_row(
"SELECT cleanup_state FROM external_workdir_grants WHERE workspace_id = ?1 AND grant_id = ?2",
rusqlite::params![TEST_WORKSPACE_ID, grant.grant_id],
|row| row.get(0),
)
.map_err(Error::from)
})
.unwrap();
if cleanup_state == "completed" {
break;
}
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
})
.await
.expect("Backend-owned expiry must run without a provider socket");
}
#[tokio::test]
async fn external_workdirs_are_launchable_with_real_multiple_attachments() {
let dir = tempfile::tempdir().unwrap();
init_clean_git_workspace(dir.path());
let api = test_api(dir.path()).await;
register_test_runtime(&api, WorkdirlessFixtureRuntime::RUNTIME_ID).await;
let runtime = WorkdirlessFixtureRuntime::default();
api.runtime.register_or_replace(runtime.clone());
let now = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true);
let expires_at = Utc::now() + Duration::minutes(5);
let mut receivers = Vec::new();
let mut workdir_ids = Vec::new();
for index in 1..=2 {
let grant = ExternalWorkdirGrantRecord {
grant_id: format!("launch-grant-{index}"),
workspace_id: TEST_WORKSPACE_ID.to_string(),
workdir_id: format!("launch-external-{index}"),
provider_instance_id: format!("launch-provider-{index}"),
display_name: format!("External launch {index}"),
permissions: "read_only".to_string(),
created_by: test_browser_request_actor().account_id,
created_at: now.clone(),
expires_at: expires_at.to_rfc3339_opts(SecondsFormat::Nanos, true),
generation: 1,
status: "pending".to_string(),
updated_at: now.clone(),
};
api.store
.create_external_workdir_grant(
&grant,
&WorkdirRegistryRecord {
workspace_id: TEST_WORKSPACE_ID.to_string(),
workdir_id: grant.workdir_id.clone(),
display_name: Some(grant.display_name.clone()),
source: WorkdirRegistrySource::ExternalGrant {
grant_id: grant.grant_id.clone(),
},
creation_selector: None,
creation_ref: None,
creation_tree: None,
current_selector: None,
current_ref: None,
current_tree: None,
observed_at_epoch_seconds: None,
materialization_status: "pending".to_string(),
cleanliness: "unknown".to_string(),
created_at: now.clone(),
updated_at: now.clone(),
},
)
.unwrap();
assert!(
api.store
.activate_external_workdir_grant(
TEST_WORKSPACE_ID,
&grant.grant_id,
&grant.provider_instance_id,
1,
&now,
)
.unwrap()
);
let (sender, receiver) = tokio::sync::mpsc::channel(4);
receivers.push(receiver);
let connection = Arc::new(ExternalProviderConnection {
grant_id: grant.grant_id.clone(),
workdir_id: grant.workdir_id.clone(),
provider_instance_id: grant.provider_instance_id.clone(),
generation: 1,
expires_at,
capabilities: workdir::WorkdirSessionCapabilities::READ_ONLY,
admission: Arc::new(tokio::sync::Semaphore::new(16)),
sender,
});
assert!(
api.external_workdir_providers
.lock()
.unwrap()
.connected(connection)
.is_none()
);
workdir_ids.push(grant.workdir_id);
}
let options = worker_launch_options_response(&api).unwrap();
let available = options
.working_directories
.iter()
.map(|summary| summary.working_directory_id.as_str())
.collect::<HashSet<_>>();
assert!(workdir_ids.iter().all(|id| available.contains(id.as_str())));
let Json(worker) = create_workspace_worker(
State(api.clone()),
HeaderMap::new(),
Json(CreateWorkspaceWorkerRequest {
runtime_id: WorkdirlessFixtureRuntime::RUNTIME_ID.to_string(),
display_name: "External multi-attachment Worker".to_string(),
profile: Some("builtin:coder".to_string()),
ticket_assignment: None,
initial_submit: Vec::new(),
workdir_attachments: vec![
server_api::BrowserWorkerWorkingDirectorySelection {
alias: "primary".to_string(),
working_directory_id: workdir_ids[0].clone(),
relative_cwd: None,
},
server_api::BrowserWorkerWorkingDirectorySelection {
alias: "reference".to_string(),
working_directory_id: workdir_ids[1].clone(),
relative_cwd: None,
},
],
control_operation_id: None,
}),
)
.await
.unwrap();
assert_eq!(worker.worker.workdir_attachments.len(), 2);
assert_eq!(
worker
.worker
.workdir_attachments
.iter()
.map(|attachment| attachment.alias.as_str())
.collect::<HashSet<_>>(),
HashSet::from(["primary", "reference"]),
);
drop(receivers);
}
#[tokio::test]
async fn external_provider_websocket_reconnects_with_fresh_generation_and_revokes() {
let dir = tempfile::tempdir().unwrap();
crates/workspace-server/src/store.rs
+157 -10
View File
@@ -21,7 +21,7 @@ use crate::workspace_deletion::WorkspaceDeletionStore;
use crate::{Error, Result};
const OLDEST_SCHEMA_VERSION: i64 = 50;
const LATEST_SCHEMA_VERSION: i64 = 64;
const LATEST_SCHEMA_VERSION: i64 = 65;
const SCHEMA_BASELINE_NAME: &str = "workspace schema baseline";
const WORKSPACE_RUNTIME_BINDINGS_MIGRATION_NAME: &str = "workspace runtime bindings";
const RUNTIME_BINDING_AUDIT_MIGRATION_NAME: &str = "workspace Runtime binding revision and audit";
@@ -44,6 +44,8 @@ const WORKER_REGISTRY_PROJECTION_MIGRATION_NAME: &str =
const MULTI_WORKDIR_ATTACHMENTS_MIGRATION_NAME: &str =
"multiple Worker Workdir attachments and Workdir display names";
const EXTERNAL_WORKDIR_GRANTS_MIGRATION_NAME: &str = "typed client-hosted External Workdir grants";
const EXTERNAL_WORKDIR_CLEANUP_MIGRATION_NAME: &str =
"durable External Workdir expiry cleanup state";
const WORKER_REGISTRY_PROJECTION_SCHEMA: &str = r#"
CREATE TABLE IF NOT EXISTS worker_registry_observations (
workspace_id TEXT NOT NULL,
@@ -159,6 +161,11 @@ const MIGRATIONS: &[Migration] = &[
name: EXTERNAL_WORKDIR_GRANTS_MIGRATION_NAME,
apply: migrate_external_workdir_grants_v63_to_v64,
},
Migration {
version: 65,
name: EXTERNAL_WORKDIR_CLEANUP_MIGRATION_NAME,
apply: migrate_external_workdir_cleanup_v64_to_v65,
},
];
#[derive(Clone, Copy)]
@@ -1644,6 +1651,10 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore {
workspace_id: &str,
grant_id: &str,
) -> Result<Option<ExternalWorkdirGrantRecord>>;
fn list_external_workdir_grants(
&self,
workspace_id: &str,
) -> Result<Vec<ExternalWorkdirGrantRecord>>;
fn activate_external_workdir_grant(
&self,
workspace_id: &str,
@@ -1666,6 +1677,14 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore {
workspace_id: &str,
now: &str,
) -> Result<Vec<String>>;
fn update_external_workdir_cleanup_state(
&self,
workspace_id: &str,
grant_id: &str,
state: &str,
error: Option<&str>,
updated_at: &str,
) -> Result<bool>;
#[allow(clippy::too_many_arguments)]
fn record_external_workdir_audit(
&self,
@@ -7661,6 +7680,38 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
})
}
fn list_external_workdir_grants(
&self,
workspace_id: &str,
) -> Result<Vec<ExternalWorkdirGrantRecord>> {
self.with_conn(|conn| {
let mut statement = conn.prepare(
r#"SELECT grant_id, workspace_id, workdir_id, provider_instance_id, display_name,
permissions, created_by, created_at, expires_at, generation, status, updated_at
FROM external_workdir_grants WHERE workspace_id = ?1 ORDER BY grant_id"#,
)?;
statement
.query_map(params![workspace_id], |row| {
Ok(ExternalWorkdirGrantRecord {
grant_id: row.get(0)?,
workspace_id: row.get(1)?,
workdir_id: row.get(2)?,
provider_instance_id: row.get(3)?,
display_name: row.get(4)?,
permissions: row.get(5)?,
created_by: row.get(6)?,
created_at: row.get(7)?,
expires_at: row.get(8)?,
generation: row.get::<_, i64>(9)? as u64,
status: row.get(10)?,
updated_at: row.get(11)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()
.map_err(Error::from)
})
}
fn activate_external_workdir_grant(
&self,
workspace_id: &str,
@@ -7706,7 +7757,11 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
) -> Result<bool> {
self.with_conn(|conn| {
let changed = conn.execute(
r#"UPDATE external_workdir_grants SET status = ?5, updated_at = ?6
r#"UPDATE external_workdir_grants
SET status = ?5, updated_at = ?6,
cleanup_state = CASE WHEN ?5 IN ('revoked', 'expired') THEN 'pending' ELSE cleanup_state END,
cleanup_error = CASE WHEN ?5 IN ('revoked', 'expired') THEN NULL ELSE cleanup_error END,
cleanup_updated_at = CASE WHEN ?5 IN ('revoked', 'expired') THEN ?6 ELSE cleanup_updated_at END
WHERE workspace_id = ?1 AND grant_id = ?2 AND provider_instance_id = ?3
AND generation = ?4 AND status NOT IN ('revoked', 'expired')"#,
params![
@@ -7741,7 +7796,9 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
self.with_conn_mut(|conn| {
let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
tx.execute(
"UPDATE external_workdir_grants SET status = 'expired', updated_at = ?2
"UPDATE external_workdir_grants
SET status = 'expired', updated_at = ?2, cleanup_state = 'pending',
cleanup_error = NULL, cleanup_updated_at = ?2
WHERE workspace_id = ?1 AND status IN ('pending', 'online', 'offline')
AND expires_at <= ?2",
params![workspace_id, now],
@@ -7762,14 +7819,11 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
)?;
let cleanup = {
let mut statement = tx.prepare(
"SELECT DISTINCT grant.workdir_id
"SELECT grant.workdir_id
FROM external_workdir_grants grant
JOIN worker_workdir_links link
ON link.workspace_id = grant.workspace_id
AND link.workdir_id = grant.workdir_id
AND link.unlinked_at IS NULL
WHERE grant.workspace_id = ?1
AND grant.status IN ('revoked', 'expired')
AND grant.cleanup_state IN ('pending', 'retry_required')
ORDER BY grant.workdir_id",
)?;
statement
@@ -7781,6 +7835,32 @@ impl ControlPlaneStore for SqliteWorkspaceStore {
})
}
fn update_external_workdir_cleanup_state(
&self,
workspace_id: &str,
grant_id: &str,
state: &str,
error: Option<&str>,
updated_at: &str,
) -> Result<bool> {
if !matches!(state, "pending" | "retry_required" | "completed") {
return Err(Error::InvalidInput(
"invalid External Workdir cleanup state".to_string(),
));
}
self.with_conn(|conn| {
conn.execute(
"UPDATE external_workdir_grants
SET cleanup_state = ?3, cleanup_error = ?4, cleanup_updated_at = ?5
WHERE workspace_id = ?1 AND grant_id = ?2
AND status IN ('revoked', 'expired')",
params![workspace_id, grant_id, state, error, updated_at],
)
.map(|changed| changed > 0)
.map_err(Error::from)
})
}
#[allow(clippy::too_many_arguments)]
fn record_external_workdir_audit(
&self,
@@ -11226,6 +11306,29 @@ fn migrate_external_workdir_grants_v63_to_v64(conn: &Connection) -> Result<()> {
Ok(())
}
fn migrate_external_workdir_cleanup_v64_to_v65(conn: &Connection) -> Result<()> {
let tx = rusqlite::Transaction::new_unchecked(conn, TransactionBehavior::Exclusive)?;
if !column_exists(&tx, "external_workdir_grants", "cleanup_state")? {
tx.execute_batch(
"ALTER TABLE external_workdir_grants ADD COLUMN cleanup_state TEXT NOT NULL DEFAULT 'not_required' CHECK (cleanup_state IN ('not_required', 'pending', 'retry_required', 'completed'));\
ALTER TABLE external_workdir_grants ADD COLUMN cleanup_error TEXT;\
ALTER TABLE external_workdir_grants ADD COLUMN cleanup_updated_at TEXT;",
)?;
tx.execute(
"UPDATE external_workdir_grants
SET cleanup_state = 'pending', cleanup_updated_at = updated_at
WHERE status IN ('revoked', 'expired')",
[],
)?;
}
tx.execute(
"INSERT INTO __yoi_schema_migrations (version, name) VALUES (?1, ?2)",
params![65, EXTERNAL_WORKDIR_CLEANUP_MIGRATION_NAME],
)?;
tx.commit()?;
Ok(())
}
fn create_latest_workspace_schema(conn: &Connection) -> Result<()> {
conn.execute_batch(include_str!("latest_schema.sql"))?;
Ok(())
@@ -12049,6 +12152,45 @@ mod tests {
)
.unwrap()
);
let cleanup_after_fence = store
.with_conn(|conn| {
conn.query_row(
"SELECT cleanup_state, cleanup_error FROM external_workdir_grants WHERE workspace_id = 'workspace-a' AND grant_id = 'grant-a'",
[],
|row| Ok((row.get::<_, String>(0)?, row.get::<_, Option<String>>(1)?)),
)
.map_err(Error::from)
})
.unwrap();
assert_eq!(cleanup_after_fence, ("pending".to_string(), None));
assert!(
store
.update_external_workdir_cleanup_state(
"workspace-a",
"grant-a",
"retry_required",
Some("attachment_release_failed"),
"2026-01-01T00:00:08Z",
)
.unwrap()
);
let cleanup_retry = store
.with_conn(|conn| {
conn.query_row(
"SELECT cleanup_state, cleanup_error FROM external_workdir_grants WHERE workspace_id = 'workspace-a' AND grant_id = 'grant-a'",
[],
|row| Ok((row.get::<_, String>(0)?, row.get::<_, Option<String>>(1)?)),
)
.map_err(Error::from)
})
.unwrap();
assert_eq!(
cleanup_retry,
(
"retry_required".to_string(),
Some("attachment_release_failed".to_string())
)
);
assert!(
!store
.activate_external_workdir_grant(
@@ -12928,6 +13070,10 @@ mod tests {
version: 64,
name: EXTERNAL_WORKDIR_GRANTS_MIGRATION_NAME.to_string(),
},
WorkspaceSchemaMigrationStep {
version: 65,
name: EXTERNAL_WORKDIR_CLEANUP_MIGRATION_NAME.to_string(),
},
]
);
@@ -12981,6 +13127,7 @@ mod tests {
(62, WORKER_REGISTRY_PROJECTION_MIGRATION_NAME.to_string()),
(63, MULTI_WORKDIR_ATTACHMENTS_MIGRATION_NAME.to_string()),
(64, EXTERNAL_WORKDIR_GRANTS_MIGRATION_NAME.to_string()),
(65, EXTERNAL_WORKDIR_CLEANUP_MIGRATION_NAME.to_string()),
]
);
assert!(!table_exists(conn, "trusted_runtime_records")?);
@@ -13051,7 +13198,7 @@ mod tests {
.iter()
.map(|migration| migration.version)
.collect::<Vec<_>>(),
vec![52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64]
vec![52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65]
);
SqliteWorkspaceStore::migrate_database(&path).unwrap();
let conn = Connection::open(&path).unwrap();
@@ -13059,7 +13206,7 @@ mod tests {
current_schema_version(&conn).unwrap(),
LATEST_SCHEMA_VERSION
);
assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 15);
assert_eq!(workspace_schema_migration_history(&conn).unwrap().len(), 16);
assert!(column_exists(&conn, "worker_workdir_links", "alias").unwrap());
assert!(!column_exists(&conn, "worker_workdir_links", "role").unwrap());
}
crates/yoi/src/workdir_share.rs
+156 -49
View File
@@ -2,7 +2,7 @@ use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use std::time::{Duration, Instant};
use chrono::Utc;
use client::BackendApiClient;
@@ -17,7 +17,8 @@ use workdir::external::{
};
use workdir::http::{WorkdirTransportError, dispatch_workdir_session_operation};
use workdir::{
BoundedReadLimits, LocalWorkdirSession, Workdir, WorkdirSession, WorkdirSessionCapabilities,
BoundedReadLimits, ExternalWorkdirRoot, LocalWorkdirSession, Workdir, WorkdirSession,
WorkdirSessionCapabilities,
};
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -29,6 +30,8 @@ pub(crate) struct WorkdirShareOptions {
pub ttl: Duration,
}
const PROVIDER_OPERATION_TIMEOUT: Duration = Duration::from_secs(30);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ProviderEnd {
Disconnected,
@@ -43,6 +46,7 @@ struct OperationCompletion {
struct OperationTask {
cancelled: Arc<AtomicBool>,
thread: Option<std::thread::JoinHandle<()>>,
}
impl OperationTask {
@@ -53,6 +57,18 @@ impl OperationTask {
fn is_cancelled(&self) -> bool {
self.cancelled.load(Ordering::Acquire)
}
fn join(mut self) -> Result<bool, String> {
let cancelled = self.is_cancelled();
let thread = self
.thread
.take()
.expect("operation task must retain its executor thread");
thread
.join()
.map_err(|_| "External Workdir operation executor panicked".to_string())?;
Ok(cancelled)
}
}
impl Drop for OperationTask {
@@ -62,18 +78,10 @@ impl Drop for OperationTask {
}
pub(crate) async fn run(options: WorkdirShareOptions) -> Result<(), String> {
let metadata = std::fs::symlink_metadata(&options.path)
.map_err(|error| format!("selected path is unavailable: {error}"))?;
if metadata.file_type().is_symlink() {
return Err("selected External Workdir root must not be a symbolic link".to_string());
}
let canonical = options
.path
.canonicalize()
.map_err(|error| format!("failed to canonicalize selected directory: {error}"))?;
if !canonical.is_dir() {
return Err("selected External Workdir path must be a directory".to_string());
}
// Pin before creating remote authority: the Backend grant can never outlive
// a failed local approval/open race, and subsequent path replacement cannot
// redirect the provider.
let pinned_root = ExternalWorkdirRoot::pin(&options.path).map_err(|error| error.to_string())?;
let client = BackendApiClient::from_stored_token(&options.backend_url)
.map_err(|error| error.to_string())?;
@@ -105,9 +113,9 @@ pub(crate) async fn run(options: WorkdirShareOptions) -> Result<(), String> {
.map_err(|error| format!("Backend returned an invalid External Workdir grant: {error}"))?;
let limits = BoundedReadLimits::EXTERNAL_DEFAULT;
let session = LocalWorkdirSession::external_read_only(
let session = LocalWorkdirSession::external_read_only_pinned(
Workdir::new(&grant.working_directory_id),
&canonical,
pinned_root,
limits,
)
.map_err(|error| error.to_string())?;
@@ -297,49 +305,63 @@ async fn serve_provider_connection(
tokio::select! {
signal = tokio::signal::ctrl_c() => {
signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?;
abort_operations(&mut operations);
stop_operations(&mut operations, &mut completions).await?;
revoke(client, &options.workspace_id, &grant.grant_id).await?;
return Ok(ProviderEnd::Interrupted);
}
completion = completions.recv() => {
let Some(completion) = completion else {
abort_operations(&mut operations);
cancel_operations(&operations);
return Err("External Workdir operation executor closed".to_string());
};
let Some(operation) = operations.remove(completion.operation_id.as_str()) else {
continue;
};
if operation.is_cancelled() {
continue;
}
let cancelled = match operation.join() {
Ok(cancelled) => cancelled,
Err(error) => {
stop_operations(&mut operations, &mut completions).await?;
return Err(error);
}
};
let outcome = if cancelled {
ExternalWorkdirOperationOutcome::Cancelled
} else {
completion.outcome
};
let result = ExternalWorkdirProviderFrame::current(
ExternalWorkdirProviderMessage::OperationResult {
generation,
operation_id: completion.operation_id,
outcome: completion.outcome,
outcome,
},
);
if send_provider_frame(&mut socket, &result).await.is_err() {
abort_operations(&mut operations);
stop_operations(&mut operations, &mut completions).await?;
return Ok(ProviderEnd::Disconnected);
}
}
incoming = socket.next() => {
let Some(incoming) = incoming else {
abort_operations(&mut operations);
stop_operations(&mut operations, &mut completions).await?;
return Ok(ProviderEnd::Disconnected);
};
let incoming = match incoming {
Ok(incoming) => incoming,
Err(_) => {
abort_operations(&mut operations);
stop_operations(&mut operations, &mut completions).await?;
return Ok(ProviderEnd::Disconnected);
}
};
match incoming {
Message::Text(text) => {
let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text)
.map_err(|error| format!("Backend returned an invalid External Workdir frame: {error}"))?;
let frame = match serde_json::from_str::<ExternalWorkdirServerFrame>(&text) {
Ok(frame) => frame,
Err(error) => {
stop_operations(&mut operations, &mut completions).await?;
return Err(format!("Backend returned an invalid External Workdir frame: {error}"));
}
};
match frame.message {
ExternalWorkdirServerMessage::Operation {
generation: frame_generation,
@@ -347,26 +369,54 @@ async fn serve_provider_connection(
operation,
} if frame_generation == generation => {
if operations.len() >= 16 || operations.contains_key(operation_id.as_str()) {
abort_operations(&mut operations);
stop_operations(&mut operations, &mut completions).await?;
return Err("Backend exceeded External Workdir operation bounds".to_string());
}
let key = operation_id.as_str().to_string();
let completion_sender = completion_sender.clone();
let session = session.clone();
let task_operation_id = operation_id.clone();
let cancelled = Arc::new(AtomicBool::new(false));
let guarded_session = session.with_operation_guard(
cancelled.clone(),
Instant::now() + PROVIDER_OPERATION_TIMEOUT,
);
let operation = operation.into_inner();
std::thread::Builder::new()
let spawned = std::thread::Builder::new()
.name("external-workdir-operation".to_string())
.spawn(move || {
let outcome = futures::executor::block_on(execute_operation(&session, operation));
let outcome = std::panic::catch_unwind(
std::panic::AssertUnwindSafe(|| {
futures::executor::block_on(execute_operation(
&guarded_session,
operation,
))
}),
)
.unwrap_or_else(|_| ExternalWorkdirOperationOutcome::Failed {
error: WorkdirTransportError {
code: workdir::http::WorkdirTransportErrorCode::Internal,
message: "External Workdir operation executor failed".to_string(),
},
});
let _ = completion_sender.blocking_send(OperationCompletion {
operation_id: task_operation_id,
outcome,
});
})
.map_err(|error| format!("failed to start External Workdir operation: {error}"))?;
operations.insert(key, OperationTask { cancelled });
});
let thread = match spawned {
Ok(thread) => thread,
Err(error) => {
stop_operations(&mut operations, &mut completions).await?;
return Err(format!("failed to start External Workdir operation: {error}"));
}
};
operations.insert(
key,
OperationTask {
cancelled,
thread: Some(thread),
},
);
}
ExternalWorkdirServerMessage::Heartbeat {
generation: frame_generation,
@@ -375,29 +425,26 @@ async fn serve_provider_connection(
let heartbeat = ExternalWorkdirProviderFrame::current(
ExternalWorkdirProviderMessage::Heartbeat { generation, sequence },
);
send_provider_frame(&mut socket, &heartbeat).await?;
if send_provider_frame(&mut socket, &heartbeat).await.is_err() {
stop_operations(&mut operations, &mut completions).await?;
return Ok(ProviderEnd::Disconnected);
}
}
ExternalWorkdirServerMessage::Cancel {
generation: frame_generation,
operation_id,
} if frame_generation == generation => {
if let Some(operation) = operations.get(operation_id.as_str()) {
// The terminal Cancelled result is emitted only after the
// executor reports that filesystem work has actually stopped.
operation.cancel();
let cancelled = ExternalWorkdirProviderFrame::current(
ExternalWorkdirProviderMessage::OperationResult {
generation,
operation_id,
outcome: ExternalWorkdirOperationOutcome::Cancelled,
},
);
send_provider_frame(&mut socket, &cancelled).await?;
}
}
ExternalWorkdirServerMessage::Revoke {
generation: frame_generation,
..
} if frame_generation == generation => {
abort_operations(&mut operations);
stop_operations(&mut operations, &mut completions).await?;
let acknowledged = ExternalWorkdirProviderFrame::current(
ExternalWorkdirProviderMessage::RevokeAcknowledged { generation },
);
@@ -405,7 +452,7 @@ async fn serve_provider_connection(
return Ok(ProviderEnd::Revoked);
}
_ => {
abort_operations(&mut operations);
stop_operations(&mut operations, &mut completions).await?;
return Err("stale or unexpected External Workdir provider frame".to_string());
}
}
@@ -417,12 +464,12 @@ async fn serve_provider_connection(
)
.await;
if !matches!(pong, Ok(Ok(()))) {
abort_operations(&mut operations);
stop_operations(&mut operations, &mut completions).await?;
return Ok(ProviderEnd::Disconnected);
}
}
Message::Close(_) => {
abort_operations(&mut operations);
stop_operations(&mut operations, &mut completions).await?;
return Ok(ProviderEnd::Disconnected);
}
_ => {}
@@ -452,12 +499,28 @@ async fn execute_operation(
}
}
fn abort_operations(operations: &mut HashMap<String, OperationTask>) {
for (_, operation) in operations.drain() {
fn cancel_operations(operations: &HashMap<String, OperationTask>) {
for operation in operations.values() {
operation.cancel();
}
}
async fn stop_operations(
operations: &mut HashMap<String, OperationTask>,
completions: &mut tokio::sync::mpsc::Receiver<OperationCompletion>,
) -> Result<(), String> {
cancel_operations(operations);
while !operations.is_empty() {
let completion = completions.recv().await.ok_or_else(|| {
"External Workdir operation executor closed during cancellation".to_string()
})?;
if let Some(operation) = operations.remove(completion.operation_id.as_str()) {
operation.join()?;
}
}
Ok(())
}
async fn send_provider_frame<S>(
socket: &mut tokio_tungstenite::WebSocketStream<S>,
frame: &ExternalWorkdirProviderFrame,
@@ -556,6 +619,50 @@ fn encode_path_segment(value: &str) -> String {
mod tests {
use super::*;
#[tokio::test]
async fn operation_shutdown_waits_until_executor_has_stopped() {
let operation_id = ExternalWorkdirOperationId::new("op-cancel-test").unwrap();
let cancelled = Arc::new(AtomicBool::new(false));
let started = Arc::new(AtomicBool::new(false));
let stopped = Arc::new(AtomicBool::new(false));
let (sender, mut completions) = tokio::sync::mpsc::channel(1);
let task_cancelled = cancelled.clone();
let task_started = started.clone();
let task_stopped = stopped.clone();
let task_operation_id = operation_id.clone();
let thread = std::thread::spawn(move || {
task_started.store(true, Ordering::Release);
while !task_cancelled.load(Ordering::Acquire) {
std::thread::yield_now();
}
std::thread::sleep(Duration::from_millis(20));
task_stopped.store(true, Ordering::Release);
sender
.blocking_send(OperationCompletion {
operation_id: task_operation_id,
outcome: ExternalWorkdirOperationOutcome::Cancelled,
})
.unwrap();
});
while !started.load(Ordering::Acquire) {
tokio::task::yield_now().await;
}
let mut operations = HashMap::from([(
operation_id.as_str().to_string(),
OperationTask {
cancelled,
thread: Some(thread),
},
)]);
stop_operations(&mut operations, &mut completions)
.await
.unwrap();
assert!(operations.is_empty());
assert!(stopped.load(Ordering::Acquire));
}
#[test]
fn path_segments_are_encoded_without_exposing_structure() {
assert_eq!(encode_path_segment("workspace/a b"), "workspace%2Fa%20b");