agent-workspace-adapter-local 0.3.0

use crate::{FilesystemSnapshotObjectStore, LocalWorkspace};
use agent_workspace_contract::*;
use anyhow::{Context, Result};
use async_trait::async_trait;
use sha2::{Digest, Sha256};
use std::path::PathBuf;
use std::sync::Arc;

/// Explicit local/test snapshot importer. The configured root is the only
/// directory under which destination workspace IDs may be created.
#[derive(Debug, Clone)]
pub struct LocalSnapshotProvider {
    destination_root: PathBuf,
    objects: FilesystemSnapshotObjectStore,
}

impl LocalSnapshotProvider {
    pub fn for_local_development(destination_root: PathBuf) -> Result<Self> {
        std::fs::create_dir_all(&destination_root)?;
        let objects = FilesystemSnapshotObjectStore::for_local_development(
            destination_root.join(".snapshot-objects"),
        )?;
        Ok(Self {
            destination_root,
            objects,
        })
    }

    fn destination(&self, workspace_id: &str) -> Result<PathBuf> {
        anyhow::ensure!(
            !workspace_id.is_empty()
                && workspace_id.len() <= 128
                && workspace_id
                    .bytes()
                    .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')),
            "invalid destination workspace id"
        );
        Ok(self.destination_root.join(workspace_id))
    }
}

async fn exact(reader: &dyn SnapshotObjectReader, offset: &mut u64, len: usize) -> Result<Vec<u8>> {
    let mut output = Vec::with_capacity(len);
    while output.len() < len {
        let bytes = reader
            .read_chunk(
                *offset,
                (len - output.len()).min(WORKSPACE_UPLOAD_CHUNK_BYTES),
            )
            .await?;
        anyhow::ensure!(!bytes.is_empty(), "truncated snapshot object");
        *offset = offset.saturating_add(bytes.len() as u64);
        output.extend_from_slice(&bytes);
    }
    Ok(output)
}

async fn integer<const N: usize>(
    reader: &dyn SnapshotObjectReader,
    offset: &mut u64,
) -> Result<[u8; N]> {
    exact(reader, offset, N)
        .await?
        .try_into()
        .map_err(|_| anyhow::anyhow!("invalid snapshot integer"))
}

#[async_trait]
impl SnapshotProvider for LocalSnapshotProvider {
    async fn export_snapshot(
        &self,
        durable_backend_id: &str,
    ) -> Result<Arc<dyn SnapshotObjectReader>> {
        let workspace = LocalWorkspace::new(PathBuf::from(durable_backend_id));
        let key = format!("exports/{}.awsnp", uuid::Uuid::new_v4().simple());
        let mut upload = self.objects.begin_upload(&key).await?;
        let mut digest = Sha256::new();
        let mut total = WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC.len();
        upload.write_chunk(WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC).await?;
        digest.update(WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC);
        for path in workspace.walk_tree(".", 64).await? {
            if workspace.is_dir(&path).await.unwrap_or(false) {
                continue;
            }
            let bytes = workspace.read_file_bytes(&path).await?;
            total = total
                .saturating_add(12)
                .saturating_add(path.len())
                .saturating_add(bytes.len());
            anyhow::ensure!(
                total <= WORKSPACE_MAX_SNAPSHOT_BYTES,
                "snapshot byte quota exceeded"
            );
            let path_len = u32::try_from(path.len())?.to_be_bytes();
            let content_len = u64::try_from(bytes.len())?.to_be_bytes();
            for chunk in [path_len.as_slice(), content_len.as_slice(), path.as_bytes()] {
                upload.write_chunk(chunk).await?;
                digest.update(chunk);
            }
            for chunk in bytes.chunks(WORKSPACE_UPLOAD_CHUNK_BYTES) {
                upload.write_chunk(chunk).await?;
                digest.update(chunk);
            }
        }
        let digest = format!("sha256:{:x}", digest.finalize());
        let object = upload.commit(&digest, total as u64).await?;
        self.objects.open(&object).await
    }

    async fn import_snapshot(
        &self,
        workspace_id: &str,
        config: &WorkspaceConfig,
        snapshot: Arc<dyn SnapshotObjectReader>,
        _idempotency_key: &str,
    ) -> Result<ProvisionedWorkspace> {
        anyhow::ensure!(
            matches!(config.backend.as_str(), "local" | "local-test"),
            "unsupported local snapshot destination backend"
        );
        anyhow::ensure!(
            snapshot.len() <= WORKSPACE_MAX_SNAPSHOT_BYTES as u64,
            "snapshot byte quota exceeded"
        );
        let mut hasher = Sha256::new();
        let mut verify_offset = 0;
        while verify_offset < snapshot.len() {
            let bytes = snapshot
                .read_chunk(verify_offset, WORKSPACE_UPLOAD_CHUNK_BYTES)
                .await?;
            anyhow::ensure!(!bytes.is_empty(), "truncated snapshot object");
            hasher.update(&bytes);
            verify_offset = verify_offset.saturating_add(bytes.len() as u64);
        }
        anyhow::ensure!(
            format!("sha256:{:x}", hasher.finalize()) == snapshot.digest(),
            "snapshot digest mismatch"
        );
        let destination = self.destination(workspace_id)?;
        let workspace = LocalWorkspace::new(destination.clone());
        let mut offset = 0;
        anyhow::ensure!(
            exact(
                snapshot.as_ref(),
                &mut offset,
                WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC.len()
            )
            .await?
                == WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC,
            "unsupported snapshot format"
        );
        while offset < snapshot.len() {
            let path_len =
                u32::from_be_bytes(integer(snapshot.as_ref(), &mut offset).await?) as usize;
            let content_len = usize::try_from(u64::from_be_bytes(
                integer(snapshot.as_ref(), &mut offset).await?,
            ))?;
            anyhow::ensure!(
                path_len <= WORKSPACE_MAX_PATH_BYTES && content_len <= WORKSPACE_MAX_FILE_BYTES,
                "snapshot entry exceeds policy"
            );
            let path = String::from_utf8(exact(snapshot.as_ref(), &mut offset, path_len).await?)?;
            let content = exact(snapshot.as_ref(), &mut offset, content_len).await?;
            workspace.write_file_bytes(&path, &content).await?;
        }
        let root = destination
            .to_str()
            .context("snapshot destination path is not UTF-8")?
            .to_owned();
        Ok(ProvisionedWorkspace {
            durable_backend_id: root.clone(),
            root,
        })
    }
}