agentos_vfs/s3/

object_backend.rs

use super::block_store::{collect_body, s3_error};
use async_trait::async_trait;
use aws_sdk_s3::primitives::ByteStream;
use aws_sdk_s3::Client;
use std::collections::HashMap;
use vfs::engine::{InodeType, ObjectBackend, ObjectEntry, ObjectMeta, Timespec, VfsResult};

#[derive(Debug, Clone, Default)]
pub struct S3ObjectBackendOptions {
    pub prefix: String,
}

#[derive(Debug, Clone)]
pub struct S3ObjectBackend {
    client: Client,
    bucket: String,
    options: S3ObjectBackendOptions,
}

impl S3ObjectBackend {
    pub fn new(client: Client, bucket: impl Into<String>) -> Self {
        Self::with_options(client, bucket, S3ObjectBackendOptions::default())
    }

    pub fn with_options(
        client: Client,
        bucket: impl Into<String>,
        options: S3ObjectBackendOptions,
    ) -> Self {
        Self {
            client,
            bucket: bucket.into(),
            options,
        }
    }

    fn key_for(&self, key: &str) -> String {
        format!("{}{}", self.options.prefix, key)
    }

    fn strip_prefix<'a>(&self, key: &'a str) -> &'a str {
        key.strip_prefix(&self.options.prefix).unwrap_or(key)
    }
}

#[async_trait]
impl ObjectBackend for S3ObjectBackend {
    async fn list(&self, prefix: &str) -> VfsResult<Vec<ObjectEntry>> {
        let s3_prefix = self.key_for(prefix);
        let mut continuation = None;
        let mut entries = Vec::new();
        loop {
            let response = self
                .client
                .list_objects_v2()
                .bucket(&self.bucket)
                .prefix(&s3_prefix)
                .delimiter("/")
                .set_continuation_token(continuation)
                .send()
                .await
                .map_err(|err| s3_error(format!("list s3 prefix '{s3_prefix}': {err}")))?;

            for object in response.contents() {
                let Some(key) = object.key() else {
                    continue;
                };
                entries.push(ObjectEntry {
                    name: self.strip_prefix(key).to_string(),
                    size: object.size().unwrap_or(0).max(0) as u64,
                    mtime: object
                        .last_modified()
                        .map(timespec_from_smithy)
                        .unwrap_or_else(Timespec::now),
                    is_prefix: false,
                });
            }

            for prefix in response.common_prefixes() {
                let Some(prefix) = prefix.prefix() else {
                    continue;
                };
                entries.push(ObjectEntry {
                    name: self.strip_prefix(prefix).to_string(),
                    size: 0,
                    mtime: Timespec::now(),
                    is_prefix: true,
                });
            }

            if response.is_truncated().unwrap_or(false) {
                continuation = response.next_continuation_token().map(ToOwned::to_owned);
            } else {
                break;
            }
        }
        Ok(entries)
    }

    async fn head(&self, key: &str) -> VfsResult<Option<ObjectMeta>> {
        let s3_key = self.key_for(key);
        let response = match self
            .client
            .head_object()
            .bucket(&self.bucket)
            .key(&s3_key)
            .send()
            .await
        {
            Ok(response) => response,
            Err(err)
                if err.as_service_error().is_some_and(|err| {
                    err.is_not_found()
                        || matches!(err.meta().code(), Some("NotFound" | "NoSuchKey"))
                }) =>
            {
                return Ok(None);
            }
            Err(err) => return Err(s3_error(format!("head s3 object '{s3_key}': {err}"))),
        };

        let metadata = response.metadata();
        let kind = metadata
            .and_then(|metadata| metadata.get("vfs-kind"))
            .map(|kind| match kind.as_str() {
                "directory" => InodeType::Directory,
                "symlink" => InodeType::Symlink,
                _ => InodeType::File,
            })
            .unwrap_or(InodeType::File);
        Ok(Some(ObjectMeta {
            size: response.content_length().unwrap_or(0).max(0) as u64,
            mtime: response
                .last_modified()
                .map(timespec_from_smithy)
                .unwrap_or_else(Timespec::now),
            mode: metadata
                .and_then(|metadata| metadata.get("vfs-mode"))
                .and_then(|mode| u32::from_str_radix(mode, 8).ok())
                .unwrap_or(0o644),
            uid: metadata
                .and_then(|metadata| metadata.get("vfs-uid"))
                .and_then(|uid| uid.parse().ok())
                .unwrap_or(0),
            gid: metadata
                .and_then(|metadata| metadata.get("vfs-gid"))
                .and_then(|gid| gid.parse().ok())
                .unwrap_or(0),
            kind,
            symlink_target: metadata
                .and_then(|metadata| metadata.get("vfs-symlink-target").cloned()),
        }))
    }

    async fn get_range(&self, key: &str, off: u64, len: u64) -> VfsResult<Vec<u8>> {
        if len == 0 {
            return Ok(Vec::new());
        }
        let s3_key = self.key_for(key);
        let end = off
            .checked_add(len)
            .and_then(|value| value.checked_sub(1))
            .ok_or_else(|| vfs::engine::VfsError::einval("invalid S3 byte range"))?;
        let response = self
            .client
            .get_object()
            .bucket(&self.bucket)
            .key(&s3_key)
            .range(format!("bytes={off}-{end}"))
            .send()
            .await
            .map_err(|err| s3_error(format!("get s3 object range '{s3_key}': {err}")))?;
        collect_body(response.body, &s3_key).await
    }

    async fn put(&self, key: &str, data: &[u8], meta: ObjectMeta) -> VfsResult<()> {
        let s3_key = self.key_for(key);
        let mut metadata = HashMap::new();
        metadata.insert("vfs-kind".to_string(), kind_name(meta.kind).to_string());
        metadata.insert("vfs-mode".to_string(), format!("{:o}", meta.mode));
        metadata.insert("vfs-uid".to_string(), meta.uid.to_string());
        metadata.insert("vfs-gid".to_string(), meta.gid.to_string());
        if let Some(target) = meta.symlink_target {
            metadata.insert("vfs-symlink-target".to_string(), target);
        }
        self.client
            .put_object()
            .bucket(&self.bucket)
            .key(&s3_key)
            .body(ByteStream::from(data.to_vec()))
            .set_metadata(Some(metadata))
            .send()
            .await
            .map_err(|err| s3_error(format!("put s3 object '{s3_key}': {err}")))?;
        Ok(())
    }

    async fn copy(&self, src: &str, dst: &str) -> VfsResult<()> {
        let src_key = self.key_for(src);
        let dst_key = self.key_for(dst);
        self.client
            .copy_object()
            .bucket(&self.bucket)
            .copy_source(format!("{}/{}", self.bucket, src_key))
            .key(&dst_key)
            .send()
            .await
            .map_err(|err| s3_error(format!("copy s3 object '{src_key}' to '{dst_key}': {err}")))?;
        Ok(())
    }

    async fn delete(&self, key: &str) -> VfsResult<()> {
        let s3_key = self.key_for(key);
        self.client
            .delete_object()
            .bucket(&self.bucket)
            .key(&s3_key)
            .send()
            .await
            .map_err(|err| s3_error(format!("delete s3 object '{s3_key}': {err}")))?;
        Ok(())
    }
}

fn kind_name(kind: InodeType) -> &'static str {
    match kind {
        InodeType::File => "file",
        InodeType::Directory => "directory",
        InodeType::Symlink => "symlink",
    }
}

fn timespec_from_smithy(time: &aws_sdk_s3::primitives::DateTime) -> Timespec {
    Timespec {
        sec: time.secs(),
        nsec: 0,
    }
}