codex/codex-rs/exec-server/src/server/file_system_handler.rs

use std::io;

use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use codex_exec_server_protocol::JSONRPCErrorError;
use codex_protocol::config_types::WindowsSandboxLevel;

use crate::CapabilityRootsDiscoverParams;
use crate::CapabilityRootsDiscoverResponse;
use crate::CopyOptions;
use crate::CreateDirectoryOptions;
use crate::ExecServerRuntimePaths;
use crate::ExecutorFileSystem;
use crate::GetMetadataOptions;
use crate::ReadFileOptions;
use crate::RemoveOptions;
use crate::WriteFileOptions;
use crate::file_read::FileReadHandleManager;
use crate::local_file_system::LocalFileSystem;
use crate::protocol::FS_READ_DIRECTORY_METHOD;
use crate::protocol::FS_WRITE_FILE_METHOD;
use crate::protocol::FsCanonicalizeParams;
use crate::protocol::FsCanonicalizeResponse;
use crate::protocol::FsCloseParams;
use crate::protocol::FsCloseResponse;
use crate::protocol::FsCopyParams;
use crate::protocol::FsCopyResponse;
use crate::protocol::FsCreateDirectoryParams;
use crate::protocol::FsCreateDirectoryResponse;
use crate::protocol::FsGetMetadataParams;
use crate::protocol::FsGetMetadataResponse;
use crate::protocol::FsOpenParams;
use crate::protocol::FsOpenResponse;
use crate::protocol::FsReadBlockParams;
use crate::protocol::FsReadBlockResponse;
use crate::protocol::FsReadDirectoryEntry;
use crate::protocol::FsReadDirectoryParams;
use crate::protocol::FsReadDirectoryResponse;
use crate::protocol::FsReadFileParams;
use crate::protocol::FsReadFileResponse;
use crate::protocol::FsRemoveParams;
use crate::protocol::FsRemoveResponse;
use crate::protocol::FsWalkParams;
use crate::protocol::FsWalkResponse;
use crate::protocol::FsWriteFileParams;
use crate::protocol::FsWriteFileResponse;
use crate::rpc::internal_error;
use crate::rpc::invalid_request;
use crate::rpc::not_found;

const MAX_FILE_READ_HANDLE_ID_BYTES: usize = 32;
// Each read-directory entry needs four JSON values. Keep same-version
// producers comfortably below the shared 256K-value decoder budget.
const MAX_READ_DIRECTORY_ENTRIES: usize = 50_000;

#[derive(Clone)]
pub(crate) struct FileSystemHandler {
    file_system: LocalFileSystem,
    file_reads: FileReadHandleManager,
}

impl FileSystemHandler {
    pub(crate) fn new(runtime_paths: ExecServerRuntimePaths) -> Self {
        Self {
            file_system: LocalFileSystem::with_runtime_paths(runtime_paths),
            file_reads: FileReadHandleManager::default(),
        }
    }

    pub(crate) async fn shutdown(&self) {
        self.file_reads.close_all().await;
    }

    pub(crate) async fn discover_capability_roots(
        &self,
        params: CapabilityRootsDiscoverParams,
    ) -> Result<CapabilityRootsDiscoverResponse, JSONRPCErrorError> {
        let sandbox = params
            .roots
            .first()
            .and_then(|root| root.sandbox.as_ref())
            .filter(|sandbox| {
                sandbox.should_run_in_sandbox()
                    && (!cfg!(target_os = "windows")
                        || sandbox.windows_sandbox_level != WindowsSandboxLevel::Disabled)
                    && params
                        .roots
                        .iter()
                        .all(|root| root.sandbox.as_ref() == Some(*sandbox))
            })
            .cloned();

        if let Some(sandbox) = sandbox {
            let mut batched_params = params.clone();
            for root in &mut batched_params.roots {
                root.sandbox = None;
            }
            let result = match self.file_system.sandboxed() {
                Ok(file_system) => {
                    file_system
                        .discover_capability_roots(batched_params, &sandbox)
                        .await
                }
                Err(error) => Err(error),
            };
            match result {
                Ok(response) => return Ok(response),
                Err(error) => {
                    tracing::warn!(%error, "batched capability discovery failed; retrying roots separately");
                }
            }
        }

        crate::discover_capability_roots(&self.file_system, params)
            .await
            .map_err(|error| invalid_request(error.to_string()))
    }

    pub(crate) async fn open(
        &self,
        params: FsOpenParams,
    ) -> Result<FsOpenResponse, JSONRPCErrorError> {
        validate_file_read_handle_id(&params.handle_id)?;
        let file = self
            .file_system
            .open_file_for_read(&params.path, params.sandbox.as_ref())
            .await
            .map_err(map_fs_error)?;
        let handle_id = self
            .file_reads
            .open(params.handle_id, file)
            .await
            .map_err(map_fs_error)?;
        Ok(FsOpenResponse { handle_id })
    }

    pub(crate) async fn read_block(
        &self,
        params: FsReadBlockParams,
    ) -> Result<FsReadBlockResponse, JSONRPCErrorError> {
        validate_file_read_handle_id(&params.handle_id)?;
        let block = self
            .file_reads
            .read_block(&params.handle_id, params.offset, params.len)
            .await
            .map_err(map_fs_error)?;
        Ok(FsReadBlockResponse {
            chunk: block.bytes.into(),
            eof: block.eof,
        })
    }

    pub(crate) async fn close(
        &self,
        params: FsCloseParams,
    ) -> Result<FsCloseResponse, JSONRPCErrorError> {
        validate_file_read_handle_id(&params.handle_id)?;
        self.file_reads.close(&params.handle_id).await;
        Ok(FsCloseResponse {})
    }

    pub(crate) async fn read_file(
        &self,
        params: FsReadFileParams,
    ) -> Result<FsReadFileResponse, JSONRPCErrorError> {
        let bytes = self
            .file_system
            .read_file(
                &params.path,
                ReadFileOptions {
                    follow_symlinks: params.follow_symlinks.unwrap_or(true),
                },
                params.sandbox.as_ref(),
            )
            .await
            .map_err(map_fs_error)?;
        Ok(FsReadFileResponse {
            data_base64: STANDARD.encode(bytes),
        })
    }

    pub(crate) async fn write_file(
        &self,
        params: FsWriteFileParams,
    ) -> Result<FsWriteFileResponse, JSONRPCErrorError> {
        let bytes = STANDARD.decode(params.data_base64).map_err(|err| {
            invalid_request(format!(
                "{FS_WRITE_FILE_METHOD} requires valid base64 dataBase64: {err}"
            ))
        })?;
        self.file_system
            .write_file(
                &params.path,
                bytes,
                WriteFileOptions {
                    follow_symlinks: params.follow_symlinks.unwrap_or(true),
                },
                params.sandbox.as_ref(),
            )
            .await
            .map_err(map_fs_error)?;
        Ok(FsWriteFileResponse {})
    }

    pub(crate) async fn create_directory(
        &self,
        params: FsCreateDirectoryParams,
    ) -> Result<FsCreateDirectoryResponse, JSONRPCErrorError> {
        let recursive = params.recursive.unwrap_or(true);
        self.file_system
            .create_directory(
                &params.path,
                CreateDirectoryOptions {
                    recursive,
                    follow_symlinks: params.follow_symlinks.unwrap_or(true),
                },
                params.sandbox.as_ref(),
            )
            .await
            .map_err(map_fs_error)?;
        Ok(FsCreateDirectoryResponse {})
    }

    pub(crate) async fn get_metadata(
        &self,
        params: FsGetMetadataParams,
    ) -> Result<FsGetMetadataResponse, JSONRPCErrorError> {
        let metadata = self
            .file_system
            .get_metadata(
                &params.path,
                GetMetadataOptions {
                    follow_symlinks: params.follow_symlinks.unwrap_or(true),
                },
                params.sandbox.as_ref(),
            )
            .await
            .map_err(map_fs_error)?;
        Ok(FsGetMetadataResponse {
            is_directory: metadata.is_directory,
            is_file: metadata.is_file,
            is_symlink: metadata.is_symlink,
            size: metadata.size,
            created_at_ms: metadata.created_at_ms,
            modified_at_ms: metadata.modified_at_ms,
        })
    }

    pub(crate) async fn canonicalize(
        &self,
        params: FsCanonicalizeParams,
    ) -> Result<FsCanonicalizeResponse, JSONRPCErrorError> {
        let path = self
            .file_system
            .canonicalize(&params.path, params.sandbox.as_ref())
            .await
            .map_err(map_fs_error)?;
        Ok(FsCanonicalizeResponse { path })
    }

    pub(crate) async fn read_directory(
        &self,
        params: FsReadDirectoryParams,
    ) -> Result<FsReadDirectoryResponse, JSONRPCErrorError> {
        let entries = self
            .file_system
            .read_directory(&params.path, params.sandbox.as_ref())
            .await
            .map_err(map_fs_error)?;
        let entry_count = entries.len();
        if entry_count > MAX_READ_DIRECTORY_ENTRIES {
            return Err(internal_error(format!(
                "{FS_READ_DIRECTORY_METHOD} returned {entry_count} entries; limit is {MAX_READ_DIRECTORY_ENTRIES}"
            )));
        }
        let entries = entries
            .into_iter()
            .map(|entry| FsReadDirectoryEntry {
                file_name: entry.file_name,
                is_directory: entry.is_directory,
                is_file: entry.is_file,
            })
            .collect();
        Ok(FsReadDirectoryResponse { entries })
    }

    pub(crate) async fn walk(
        &self,
        params: FsWalkParams,
    ) -> Result<FsWalkResponse, JSONRPCErrorError> {
        self.file_system
            .walk(&params.path, params.options, params.sandbox.as_ref())
            .await
            .map_err(map_fs_error)
    }

    pub(crate) async fn remove(
        &self,
        params: FsRemoveParams,
    ) -> Result<FsRemoveResponse, JSONRPCErrorError> {
        let recursive = params.recursive.unwrap_or(true);
        let force = params.force.unwrap_or(true);
        self.file_system
            .remove(
                &params.path,
                RemoveOptions {
                    recursive,
                    force,
                    follow_symlinks: params.follow_symlinks.unwrap_or(true),
                },
                params.sandbox.as_ref(),
            )
            .await
            .map_err(map_fs_error)?;
        Ok(FsRemoveResponse {})
    }

    pub(crate) async fn copy(
        &self,
        params: FsCopyParams,
    ) -> Result<FsCopyResponse, JSONRPCErrorError> {
        self.file_system
            .copy(
                &params.source_path,
                &params.destination_path,
                CopyOptions {
                    recursive: params.recursive,
                },
                params.sandbox.as_ref(),
            )
            .await
            .map_err(map_fs_error)?;
        Ok(FsCopyResponse {})
    }
}

fn validate_file_read_handle_id(handle_id: &str) -> Result<(), JSONRPCErrorError> {
    if handle_id.len() > MAX_FILE_READ_HANDLE_ID_BYTES {
        return Err(invalid_request(format!(
            "file read handle ID must not exceed {MAX_FILE_READ_HANDLE_ID_BYTES} bytes"
        )));
    }
    Ok(())
}

fn map_fs_error(err: io::Error) -> JSONRPCErrorError {
    match err.kind() {
        io::ErrorKind::NotFound => not_found(err.to_string()),
        io::ErrorKind::InvalidInput | io::ErrorKind::PermissionDenied => {
            invalid_request(err.to_string())
        }
        _ => internal_error(err.to_string()),
    }
}

#[cfg(test)]
mod tests {
    use codex_protocol::protocol::NetworkAccess;
    use codex_protocol::protocol::SandboxPolicy;
    use codex_utils_path_uri::PathUri;
    use pretty_assertions::assert_eq;

    use super::*;
    use crate::FileSystemSandboxContext;
    use crate::protocol::FsReadFileParams;
    use crate::protocol::FsWriteFileParams;

    #[tokio::test]
    async fn no_platform_sandbox_policies_do_not_require_configured_sandbox_helper() {
        let temp_dir = tempfile::tempdir().expect("tempdir");
        let runtime_paths = ExecServerRuntimePaths::new(
            std::env::current_exe().expect("current exe"),
            /*codex_linux_sandbox_exe*/ None,
        )
        .expect("runtime paths");
        let handler = FileSystemHandler::new(runtime_paths);
        let sandbox_cwd = PathUri::from_host_native_path(temp_dir.path()).expect("tempdir URI");
        let sandbox_context = |sandbox_policy| {
            FileSystemSandboxContext::from_legacy_sandbox_policy(
                sandbox_policy,
                sandbox_cwd.clone(),
            )
            .expect("sandbox context")
        };

        for (file_name, sandbox_policy) in [
            ("danger.txt", SandboxPolicy::DangerFullAccess),
            (
                "external.txt",
                SandboxPolicy::ExternalSandbox {
                    network_access: NetworkAccess::Restricted,
                },
            ),
        ] {
            let path =
                PathUri::from_host_native_path(temp_dir.path().join(file_name)).expect("path URI");

            handler
                .write_file(FsWriteFileParams {
                    path: path.clone(),
                    follow_symlinks: None,
                    data_base64: STANDARD.encode("ok"),
                    sandbox: Some(sandbox_context(sandbox_policy.clone())),
                })
                .await
                .expect("write file");

            let canonicalized = handler
                .canonicalize(FsCanonicalizeParams {
                    path: path.clone(),
                    sandbox: Some(sandbox_context(sandbox_policy.clone())),
                })
                .await
                .expect("canonicalize file");
            assert_eq!(
                canonicalized.path,
                PathUri::from_host_native_path(
                    std::fs::canonicalize(temp_dir.path().join(file_name)).expect("canonical path"),
                )
                .expect("canonical path URI"),
            );

            let response = handler
                .read_file(FsReadFileParams {
                    path,
                    follow_symlinks: None,
                    sandbox: Some(sandbox_context(sandbox_policy)),
                })
                .await
                .expect("read file");

            assert_eq!(response.data_base64, STANDARD.encode("ok"));
        }
    }
}