byteor_adapters/

s3.rs

use std::collections::{HashSet, VecDeque};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use aws_config::{BehaviorVersion, Region};
use aws_sdk_s3::primitives::ByteStream;
use tokio::runtime::Runtime;

use crate::{Adapter, AdapterError, EgressAdapter, IngressAdapter};

pub(crate) struct S3IngressObject {
    pub key: String,
    pub payload: Vec<u8>,
    pub size_bytes: usize,
    pub e_tag: Option<String>,
    pub pending_after_read: usize,
}

pub(crate) struct S3PutResult {
    pub e_tag: Option<String>,
}

#[derive(Debug)]
struct S3IngressStatus {
    transfer_state: String,
    last_object_key: Option<String>,
    last_object_size_bytes: Option<u64>,
    last_object_etag: Option<String>,
    pending_objects: u64,
}

impl Default for S3IngressStatus {
    fn default() -> Self {
        Self {
            transfer_state: "idle".to_string(),
            last_object_key: None,
            last_object_size_bytes: None,
            last_object_etag: None,
            pending_objects: 0,
        }
    }
}

impl S3IngressStatus {
    fn status_fields(&self) -> serde_json::Map<String, serde_json::Value> {
        let mut details = serde_json::Map::new();
        details.insert(
            "transfer_state".to_string(),
            serde_json::Value::String(self.transfer_state.clone()),
        );
        details.insert(
            "pending_objects".to_string(),
            serde_json::Value::from(self.pending_objects),
        );
        if let Some(key) = &self.last_object_key {
            details.insert(
                "last_object_key".to_string(),
                serde_json::Value::String(key.clone()),
            );
        }
        if let Some(size_bytes) = self.last_object_size_bytes {
            details.insert(
                "last_object_size_bytes".to_string(),
                serde_json::Value::from(size_bytes),
            );
        }
        if let Some(e_tag) = &self.last_object_etag {
            details.insert(
                "last_object_etag".to_string(),
                serde_json::Value::String(e_tag.clone()),
            );
        }
        details
    }
}

#[derive(Debug)]
struct S3EgressStatus {
    transfer_state: String,
    last_object_key: Option<String>,
    last_object_size_bytes: Option<u64>,
    last_object_etag: Option<String>,
}

impl Default for S3EgressStatus {
    fn default() -> Self {
        Self {
            transfer_state: "idle".to_string(),
            last_object_key: None,
            last_object_size_bytes: None,
            last_object_etag: None,
        }
    }
}

impl S3EgressStatus {
    fn status_fields(&self) -> serde_json::Map<String, serde_json::Value> {
        let mut details = serde_json::Map::new();
        details.insert(
            "transfer_state".to_string(),
            serde_json::Value::String(self.transfer_state.clone()),
        );
        if let Some(key) = &self.last_object_key {
            details.insert(
                "last_object_key".to_string(),
                serde_json::Value::String(key.clone()),
            );
        }
        if let Some(size_bytes) = self.last_object_size_bytes {
            details.insert(
                "last_object_size_bytes".to_string(),
                serde_json::Value::from(size_bytes),
            );
        }
        if let Some(e_tag) = &self.last_object_etag {
            details.insert(
                "last_object_etag".to_string(),
                serde_json::Value::String(e_tag.clone()),
            );
        }
        details
    }
}

/// Internal S3 ingress client contract used by the adapter and crate-local tests.
pub(crate) trait S3IngressClient: Send {
    fn read_next(&mut self, timeout: Duration) -> Result<Option<S3IngressObject>, AdapterError>;
}

/// Internal S3 egress client contract used by the adapter and crate-local tests.
pub(crate) trait S3EgressClient: Send {
    fn put_object(
        &mut self,
        key: &str,
        payload: &[u8],
        content_type: Option<&str>,
    ) -> Result<S3PutResult, AdapterError>;
}

struct AsyncS3IngressClient {
    runtime: Runtime,
    client: aws_sdk_s3::Client,
    bucket: String,
    prefix: Option<String>,
    max_object_bytes: usize,
    seen_keys: HashSet<String>,
    pending: VecDeque<S3IngressObject>,
}

impl AsyncS3IngressClient {
    fn connect(
        region: &str,
        endpoint: Option<&str>,
        bucket: &str,
        prefix: Option<&str>,
        max_object_bytes: usize,
    ) -> Result<Self, AdapterError> {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .map_err(|err| AdapterError::S3 {
                op: "s3_runtime_build",
                detail: err.to_string(),
            })?;
        let client = runtime.block_on(async {
            let shared_config = aws_config::defaults(BehaviorVersion::latest())
                .region(Region::new(region.to_string()))
                .load()
                .await;
            let mut builder = aws_sdk_s3::config::Builder::from(&shared_config);
            if let Some(endpoint) = endpoint {
                builder = builder.endpoint_url(endpoint).force_path_style(true);
            }
            Ok::<_, AdapterError>(aws_sdk_s3::Client::from_conf(builder.build()))
        })?;

        Ok(Self {
            runtime,
            client,
            bucket: bucket.to_string(),
            prefix: normalize_prefix(prefix),
            max_object_bytes,
            seen_keys: HashSet::new(),
            pending: VecDeque::new(),
        })
    }

    async fn fetch_pending_objects(
        client: aws_sdk_s3::Client,
        bucket: String,
        prefix: Option<String>,
        max_object_bytes: usize,
        seen_keys: &HashSet<String>,
    ) -> Result<Vec<S3IngressObject>, AdapterError> {
        let mut continuation_token = None;
        let mut discovered = Vec::new();

        loop {
            let mut request = client.list_objects_v2().bucket(&bucket);
            if let Some(prefix) = &prefix {
                request = request.prefix(prefix);
            }
            if let Some(token) = continuation_token.take() {
                request = request.continuation_token(token);
            }

            let response = request.send().await.map_err(|err| AdapterError::S3 {
                op: "s3_list_objects_v2",
                detail: err.to_string(),
            })?;

            for object in response.contents() {
                let Some(key) = object.key() else {
                    continue;
                };
                if seen_keys.contains(key) {
                    continue;
                }
                if let Some(size) = object.size() {
                    if size < 0 {
                        return Err(AdapterError::S3 {
                            op: "s3_list_objects_v2",
                            detail: format!("object {key} returned a negative size"),
                        });
                    }
                    if size as usize > max_object_bytes {
                        return Err(AdapterError::S3 {
                            op: "s3_get_object",
                            detail: format!(
                                "object too large key={key} len={} max={max_object_bytes}",
                                size
                            ),
                        });
                    }
                }
                discovered.push(key.to_string());
            }

            if !response.is_truncated().unwrap_or(false) {
                break;
            }
            continuation_token = response.next_continuation_token().map(ToOwned::to_owned);
            if continuation_token.is_none() {
                break;
            }
        }

        discovered.sort();
        let mut fetched = Vec::with_capacity(discovered.len());
        for key in discovered {
            let response = client
                .get_object()
                .bucket(&bucket)
                .key(&key)
                .send()
                .await
                .map_err(|err| AdapterError::S3 {
                    op: "s3_get_object",
                    detail: err.to_string(),
                })?;
            if let Some(content_length) = response.content_length() {
                if content_length < 0 {
                    return Err(AdapterError::S3 {
                        op: "s3_get_object",
                        detail: format!("object {key} returned a negative content length"),
                    });
                }
                if content_length as usize > max_object_bytes {
                    return Err(AdapterError::S3 {
                        op: "s3_get_object",
                        detail: format!(
                            "object too large key={key} len={} max={max_object_bytes}",
                            content_length
                        ),
                    });
                }
            }

            let e_tag = response.e_tag().map(ToOwned::to_owned);

            let payload = response
                .body
                .collect()
                .await
                .map_err(|err| AdapterError::S3 {
                    op: "s3_collect_object",
                    detail: err.to_string(),
                })?
                .into_bytes()
                .to_vec();

            if payload.len() > max_object_bytes {
                return Err(AdapterError::S3 {
                    op: "s3_collect_object",
                    detail: format!(
                        "object too large key={key} len={} max={max_object_bytes}",
                        payload.len()
                    ),
                });
            }

            let size_bytes = payload.len();
            fetched.push(S3IngressObject {
                key,
                payload,
                size_bytes,
                e_tag,
                pending_after_read: 0,
            });
        }

        let total = fetched.len();
        for (index, object) in fetched.iter_mut().enumerate() {
            object.pending_after_read = total.saturating_sub(index + 1);
        }

        Ok(fetched)
    }
}

impl S3IngressClient for AsyncS3IngressClient {
    fn read_next(&mut self, timeout: Duration) -> Result<Option<S3IngressObject>, AdapterError> {
        if let Some(mut object) = self.pending.pop_front() {
            object.pending_after_read = self.pending.len();
            return Ok(Some(object));
        }

        let client = self.client.clone();
        let bucket = self.bucket.clone();
        let prefix = self.prefix.clone();
        let max_object_bytes = self.max_object_bytes;
        let fetched = self.runtime.block_on(Self::fetch_pending_objects(
            client,
            bucket,
            prefix,
            max_object_bytes,
            &self.seen_keys,
        ))?;

        if fetched.is_empty() {
            std::thread::sleep(timeout);
            return Ok(None);
        }

        for object in fetched {
            self.seen_keys.insert(object.key.clone());
            self.pending.push_back(object);
        }

        let Some(mut object) = self.pending.pop_front() else {
            return Ok(None);
        };
        object.pending_after_read = self.pending.len();
        Ok(Some(object))
    }
}

struct AsyncS3EgressClient {
    runtime: Runtime,
    client: aws_sdk_s3::Client,
    bucket: String,
}

impl AsyncS3EgressClient {
    fn connect(region: &str, endpoint: Option<&str>, bucket: &str) -> Result<Self, AdapterError> {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .map_err(|err| AdapterError::S3 {
                op: "s3_runtime_build",
                detail: err.to_string(),
            })?;
        let client = runtime.block_on(async {
            let shared_config = aws_config::defaults(BehaviorVersion::latest())
                .region(Region::new(region.to_string()))
                .load()
                .await;
            let mut builder = aws_sdk_s3::config::Builder::from(&shared_config);
            if let Some(endpoint) = endpoint {
                builder = builder.endpoint_url(endpoint).force_path_style(true);
            }
            Ok::<_, AdapterError>(aws_sdk_s3::Client::from_conf(builder.build()))
        })?;

        Ok(Self {
            runtime,
            client,
            bucket: bucket.to_string(),
        })
    }
}

impl S3EgressClient for AsyncS3EgressClient {
    fn put_object(
        &mut self,
        key: &str,
        payload: &[u8],
        content_type: Option<&str>,
    ) -> Result<S3PutResult, AdapterError> {
        let client = self.client.clone();
        let bucket = self.bucket.clone();
        let key = key.to_string();
        let payload = payload.to_vec();
        let content_type = content_type.map(ToOwned::to_owned);

        self.runtime.block_on(async move {
            let mut request = client
                .put_object()
                .bucket(bucket)
                .key(key)
                .body(ByteStream::from(payload));
            if let Some(content_type) = content_type {
                request = request.content_type(content_type);
            }
            let response = request.send().await.map_err(|err| AdapterError::S3 {
                op: "s3_put_object",
                detail: err.to_string(),
            })?;
            Ok(S3PutResult {
                e_tag: response.e_tag().map(ToOwned::to_owned),
            })
        })
    }
}

fn normalize_prefix(prefix: Option<&str>) -> Option<String> {
    prefix
        .map(str::trim)
        .filter(|prefix| !prefix.is_empty())
        .map(ToOwned::to_owned)
}

fn build_object_key(prefix: Option<&str>, sequence: u64) -> String {
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos();
    let leaf = format!("{nanos:020}-{sequence:020}.bin");
    match prefix
        .map(|prefix| prefix.trim_end_matches('/'))
        .filter(|prefix| !prefix.is_empty())
    {
        Some(prefix) => format!("{}/{}", prefix.trim_end_matches('/'), leaf),
        None => leaf,
    }
}

/// Minimal blocking S3 ingress adapter.
///
/// This adapter polls one bucket/prefix and exposes each unseen object body as one ingress
/// message. Credentials come from the AWS default provider chain; a custom `endpoint` can be used
/// for S3-compatible stores such as MinIO.
pub struct S3Ingress {
    name: String,
    client: Box<dyn S3IngressClient>,
    poll_interval: Duration,
    status: Arc<Mutex<S3IngressStatus>>,
}

impl S3Ingress {
    /// Create a blocking S3 ingress adapter.
    pub fn new(
        name: impl Into<String>,
        bucket: &str,
        prefix: Option<&str>,
        region: &str,
        endpoint: Option<&str>,
        poll_interval: Duration,
        max_object_bytes: usize,
    ) -> Result<Self, AdapterError> {
        let client =
            AsyncS3IngressClient::connect(region, endpoint, bucket, prefix, max_object_bytes)?;
        Ok(Self {
            name: name.into(),
            client: Box::new(client),
            poll_interval,
            status: Arc::new(Mutex::new(S3IngressStatus::default())),
        })
    }

    #[cfg(test)]
    /// Build an S3 ingress adapter around a crate-local fake client.
    pub(crate) fn with_client(
        name: impl Into<String>,
        client: impl S3IngressClient + 'static,
        poll_interval: Duration,
    ) -> Self {
        Self {
            name: name.into(),
            client: Box::new(client),
            poll_interval,
            status: Arc::new(Mutex::new(S3IngressStatus::default())),
        }
    }
}

impl Adapter for S3Ingress {
    fn name(&self) -> &str {
        &self.name
    }

    fn transport_kind(&self) -> &str {
        "s3"
    }

    fn status_fields(&self) -> serde_json::Map<String, serde_json::Value> {
        self.status.lock().unwrap().status_fields()
    }
}

impl IngressAdapter for S3Ingress {
    fn read_next(&mut self, out: &mut [u8]) -> Result<Option<usize>, AdapterError> {
        self.status.lock().unwrap().transfer_state = "polling".to_string();
        let Some(object) = self.client.read_next(self.poll_interval)? else {
            self.status.lock().unwrap().transfer_state = "idle".to_string();
            return Ok(None);
        };
        let payload = object.payload;
        if payload.len() > out.len() {
            return Err(AdapterError::S3 {
                op: "s3_read_next",
                detail: format!("payload too large len={} buf={}", payload.len(), out.len()),
            });
        }
        {
            let mut status = self.status.lock().unwrap();
            status.transfer_state = "downloaded".to_string();
            status.pending_objects = object.pending_after_read as u64;
            if !object.key.is_empty() {
                status.last_object_key = Some(object.key);
            }
            status.last_object_size_bytes = Some(object.size_bytes as u64);
            status.last_object_etag = object.e_tag;
        }
        out[..payload.len()].copy_from_slice(&payload);
        Ok(Some(payload.len()))
    }
}

/// Minimal blocking S3 egress adapter.
///
/// Each `write_msg()` uploads one object under the configured bucket/prefix. Keys are generated as
/// `<prefix>/<unix-nanos>-<sequence>.bin`.
pub struct S3Egress {
    name: String,
    client: Box<dyn S3EgressClient>,
    prefix: Option<String>,
    content_type: Option<String>,
    sequence: u64,
    status: Arc<Mutex<S3EgressStatus>>,
}

impl S3Egress {
    /// Create a blocking S3 egress adapter.
    pub fn new(
        name: impl Into<String>,
        bucket: &str,
        prefix: Option<&str>,
        region: &str,
        endpoint: Option<&str>,
        content_type: Option<&str>,
    ) -> Result<Self, AdapterError> {
        let client = AsyncS3EgressClient::connect(region, endpoint, bucket)?;
        Ok(Self {
            name: name.into(),
            client: Box::new(client),
            prefix: normalize_prefix(prefix),
            content_type: content_type.map(ToOwned::to_owned),
            sequence: 0,
            status: Arc::new(Mutex::new(S3EgressStatus::default())),
        })
    }

    #[cfg(test)]
    /// Build an S3 egress adapter around a crate-local fake client.
    pub(crate) fn with_client(
        name: impl Into<String>,
        client: impl S3EgressClient + 'static,
        prefix: Option<&str>,
        content_type: Option<&str>,
    ) -> Self {
        Self {
            name: name.into(),
            client: Box::new(client),
            prefix: normalize_prefix(prefix),
            content_type: content_type.map(ToOwned::to_owned),
            sequence: 0,
            status: Arc::new(Mutex::new(S3EgressStatus::default())),
        }
    }
}

impl Adapter for S3Egress {
    fn name(&self) -> &str {
        &self.name
    }

    fn transport_kind(&self) -> &str {
        "s3"
    }

    fn status_fields(&self) -> serde_json::Map<String, serde_json::Value> {
        self.status.lock().unwrap().status_fields()
    }
}

impl EgressAdapter for S3Egress {
    fn write_msg(&mut self, msg: &[u8]) -> Result<(), AdapterError> {
        self.sequence = self.sequence.saturating_add(1);
        let key = build_object_key(self.prefix.as_deref(), self.sequence);
        self.status.lock().unwrap().transfer_state = "uploading".to_string();
        let result = self
            .client
            .put_object(&key, msg, self.content_type.as_deref())?;
        let mut status = self.status.lock().unwrap();
        status.transfer_state = "uploaded".to_string();
        status.last_object_key = Some(key);
        status.last_object_size_bytes = Some(msg.len() as u64);
        status.last_object_etag = result.e_tag;
        Ok(())
    }

    fn flush(&mut self) -> Result<(), AdapterError> {
        Ok(())
    }
}