edgesentry_rs/ingest/

storage.rs

use std::collections::HashMap;

use thiserror::Error;
use tracing::{debug, error, info, instrument, warn};

use crate::integrity::compute_payload_hash;
use crate::record::AuditRecord;
use super::policy::IntegrityPolicyGate;
use super::verify::IngestError;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IngestDecision {
    Accepted,
    Rejected,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OperationLogEntry {
    pub decision: IngestDecision,
    pub device_id: String,
    pub sequence: u64,
    pub message: String,
}

pub trait RawDataStore {
    type Error: std::error::Error;

    fn put(&mut self, object_ref: &str, payload: &[u8]) -> Result<(), Self::Error>;
}

pub trait AuditLedger {
    type Error: std::error::Error;

    fn append(&mut self, record: AuditRecord) -> Result<(), Self::Error>;
}

pub trait OperationLogStore {
    type Error: std::error::Error;

    fn write(&mut self, entry: OperationLogEntry) -> Result<(), Self::Error>;
}

#[derive(Debug, Error)]
pub enum IngestServiceError {
    #[error("ingest verification failed: {0}")]
    Verify(#[from] IngestError),
    #[error("payload hash mismatch for device={device_id} sequence={sequence}")]
    PayloadHashMismatch { device_id: String, sequence: u64 },
    #[error("raw data store error: {0}")]
    RawDataStore(String),
    #[error("audit ledger error: {0}")]
    AuditLedger(String),
    #[error("operation log error: {0}")]
    OperationLog(String),
}

pub struct IngestService<R, L, O>
where
    R: RawDataStore,
    L: AuditLedger,
    O: OperationLogStore,
{
    policy: IntegrityPolicyGate,
    raw_data_store: R,
    audit_ledger: L,
    operation_log: O,
}

impl<R, L, O> IngestService<R, L, O>
where
    R: RawDataStore,
    L: AuditLedger,
    O: OperationLogStore,
{
    pub fn new(policy: IntegrityPolicyGate, raw_data_store: R, audit_ledger: L, operation_log: O) -> Self {
        Self {
            policy,
            raw_data_store,
            audit_ledger,
            operation_log,
        }
    }

    pub fn register_device(&mut self, device_id: impl Into<String>, key: ed25519_dalek::VerifyingKey) {
        self.policy.register_device(device_id, key);
    }

    #[instrument(skip(self, raw_payload), fields(
        device_id = %record.device_id,
        sequence  = record.sequence,
        object_ref = %record.object_ref,
    ))]
    pub fn ingest(&mut self, record: AuditRecord, raw_payload: &[u8], cert_identity: Option<&str>) -> Result<(), IngestServiceError> {
        debug!(payload_bytes = raw_payload.len(), "ingest started");

        let payload_hash = compute_payload_hash(raw_payload);
        if payload_hash != record.payload_hash {
            warn!("payload hash mismatch — record rejected");
            self.log_rejection(&record, "payload hash mismatch");
            return Err(IngestServiceError::PayloadHashMismatch {
                device_id: record.device_id,
                sequence: record.sequence,
            });
        }

        if let Err(error) = self.policy.enforce(&record, cert_identity) {
            warn!(reason = %error, "integrity policy rejected record");
            self.log_rejection(&record, &error.to_string());
            return Err(IngestServiceError::Verify(error));
        }

        self.raw_data_store
            .put(&record.object_ref, raw_payload)
            .map_err(|e| {
                error!(error = %e, "raw data store write failed");
                IngestServiceError::RawDataStore(e.to_string())
            })?;

        self.audit_ledger
            .append(record.clone())
            .map_err(|e| {
                error!(error = %e, "audit ledger append failed");
                IngestServiceError::AuditLedger(e.to_string())
            })?;

        self.operation_log
            .write(OperationLogEntry {
                decision: IngestDecision::Accepted,
                device_id: record.device_id,
                sequence: record.sequence,
                message: "ingest accepted".to_string(),
            })
            .map_err(|e| {
                error!(error = %e, "operation log write failed");
                IngestServiceError::OperationLog(e.to_string())
            })?;

        info!("record accepted");
        Ok(())
    }

    pub fn raw_data_store(&self) -> &R {
        &self.raw_data_store
    }

    pub fn audit_ledger(&self) -> &L {
        &self.audit_ledger
    }

    pub fn operation_log(&self) -> &O {
        &self.operation_log
    }

    fn log_rejection(&mut self, record: &AuditRecord, reason: &str) {
        let _ = self.operation_log.write(OperationLogEntry {
            decision: IngestDecision::Rejected,
            device_id: record.device_id.clone(),
            sequence: record.sequence,
            message: reason.to_string(),
        });
    }
}

#[derive(Debug, Error)]
#[error("in-memory store error: {message}")]
pub struct InMemoryStoreError {
    message: String,
}

#[derive(Default)]
pub struct InMemoryRawDataStore {
    objects: HashMap<String, Vec<u8>>,
}

impl InMemoryRawDataStore {
    pub fn get(&self, object_ref: &str) -> Option<&[u8]> {
        self.objects.get(object_ref).map(Vec::as_slice)
    }
}

impl RawDataStore for InMemoryRawDataStore {
    type Error = InMemoryStoreError;

    fn put(&mut self, object_ref: &str, payload: &[u8]) -> Result<(), Self::Error> {
        self.objects.insert(object_ref.to_string(), payload.to_vec());
        Ok(())
    }
}

#[derive(Default)]
pub struct InMemoryAuditLedger {
    records: Vec<AuditRecord>,
}

impl InMemoryAuditLedger {
    pub fn records(&self) -> &[AuditRecord] {
        &self.records
    }
}

impl AuditLedger for InMemoryAuditLedger {
    type Error = InMemoryStoreError;

    fn append(&mut self, record: AuditRecord) -> Result<(), Self::Error> {
        self.records.push(record);
        Ok(())
    }
}

#[derive(Default)]
pub struct InMemoryOperationLog {
    entries: Vec<OperationLogEntry>,
}

impl InMemoryOperationLog {
    pub fn entries(&self) -> &[OperationLogEntry] {
        &self.entries
    }
}

impl OperationLogStore for InMemoryOperationLog {
    type Error = InMemoryStoreError;

    fn write(&mut self, entry: OperationLogEntry) -> Result<(), Self::Error> {
        self.entries.push(entry);
        Ok(())
    }
}

#[cfg(feature = "postgres")]
#[derive(Debug, Error)]
pub enum PostgresStoreError {
    #[error("postgres error: {0}")]
    Postgres(#[from] ::postgres::Error),
}

#[cfg(feature = "postgres")]
pub struct PostgresAuditLedger {
    client: ::postgres::Client,
}

#[cfg(feature = "postgres")]
impl PostgresAuditLedger {
    pub fn connect(url: &str) -> Result<Self, ::postgres::Error> {
        let client = ::postgres::Client::connect(url, ::postgres::NoTls)?;
        Ok(Self { client })
    }
}

#[cfg(feature = "postgres")]
impl AuditLedger for PostgresAuditLedger {
    type Error = PostgresStoreError;

    fn append(&mut self, record: AuditRecord) -> Result<(), Self::Error> {
        self.client.execute(
            "INSERT INTO audit_records \
             (device_id, sequence, timestamp_ms, payload_hash, signature, prev_record_hash, object_ref) \
             VALUES ($1, $2, $3, $4, $5, $6, $7)",
            &[
                &record.device_id,
                &(record.sequence as i64),
                &(record.timestamp_ms as i64),
                &record.payload_hash.as_ref(),
                &record.signature.as_ref(),
                &record.prev_record_hash.as_ref(),
                &record.object_ref,
            ],
        )?;
        Ok(())
    }
}

#[cfg(feature = "postgres")]
pub struct PostgresOperationLog {
    client: ::postgres::Client,
}

#[cfg(feature = "postgres")]
impl PostgresOperationLog {
    pub fn connect(url: &str) -> Result<Self, ::postgres::Error> {
        let client = ::postgres::Client::connect(url, ::postgres::NoTls)?;
        Ok(Self { client })
    }

    pub fn reset(&mut self) -> Result<(), ::postgres::Error> {
        self.client.batch_execute(
            "TRUNCATE TABLE operation_logs, audit_records RESTART IDENTITY;",
        )
    }
}

#[cfg(feature = "postgres")]
impl OperationLogStore for PostgresOperationLog {
    type Error = ::postgres::Error;

    fn write(&mut self, entry: OperationLogEntry) -> Result<(), Self::Error> {
        let decision = match entry.decision {
            IngestDecision::Accepted => "Accepted",
            IngestDecision::Rejected => "Rejected",
        };
        self.client.execute(
            "INSERT INTO operation_logs (decision, device_id, sequence, message) \
             VALUES ($1, $2, $3, $4)",
            &[
                &decision,
                &entry.device_id,
                &(entry.sequence as i64),
                &entry.message,
            ],
        )?;
        Ok(())
    }
}

#[cfg(feature = "s3")]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum S3Backend {
    AwsS3,
    Minio,
}

#[cfg(feature = "s3")]
#[derive(Debug, Clone)]
pub struct S3ObjectStoreConfig {
    pub backend: S3Backend,
    pub bucket: String,
    pub region: String,
    pub endpoint: Option<String>,
    pub access_key_id: Option<String>,
    pub secret_access_key: Option<String>,
}

#[cfg(feature = "s3")]
impl S3ObjectStoreConfig {
    pub fn for_aws_s3(bucket: impl Into<String>, region: impl Into<String>) -> Self {
        Self {
            backend: S3Backend::AwsS3,
            bucket: bucket.into(),
            region: region.into(),
            endpoint: None,
            access_key_id: None,
            secret_access_key: None,
        }
    }

    pub fn for_minio(
        bucket: impl Into<String>,
        region: impl Into<String>,
        endpoint: impl Into<String>,
        access_key_id: impl Into<String>,
        secret_access_key: impl Into<String>,
    ) -> Self {
        Self {
            backend: S3Backend::Minio,
            bucket: bucket.into(),
            region: region.into(),
            endpoint: Some(endpoint.into()),
            access_key_id: Some(access_key_id.into()),
            secret_access_key: Some(secret_access_key.into()),
        }
    }
}

#[cfg(feature = "s3")]
#[derive(Debug, Error)]
pub enum S3StoreError {
    #[error("invalid config: {0}")]
    InvalidConfig(String),
    #[error("runtime initialization failed: {0}")]
    Runtime(String),
    #[error("s3 put failed: {0}")]
    Put(String),
}

#[cfg(feature = "s3")]
pub struct S3CompatibleRawDataStore {
    runtime: tokio::runtime::Runtime,
    client: aws_sdk_s3::Client,
    bucket: String,
}

#[cfg(feature = "s3")]
impl S3CompatibleRawDataStore {
    pub fn new(config: S3ObjectStoreConfig) -> Result<Self, S3StoreError> {
        use aws_config::BehaviorVersion;
        use aws_config::Region;
        use aws_credential_types::Credentials;

        let runtime = tokio::runtime::Runtime::new()
            .map_err(|e| S3StoreError::Runtime(e.to_string()))?;

        let mut loader = aws_config::defaults(BehaviorVersion::latest())
            .region(Region::new(config.region.clone()));

        match config.backend {
            S3Backend::AwsS3 => {
                if let (Some(access_key_id), Some(secret_access_key)) =
                    (config.access_key_id.clone(), config.secret_access_key.clone())
                {
                    let creds = Credentials::new(access_key_id, secret_access_key, None, None, "static");
                    loader = loader.credentials_provider(creds);
                }
            }
            S3Backend::Minio => {
                let endpoint = config.endpoint.clone().ok_or_else(|| {
                    S3StoreError::InvalidConfig("endpoint is required for MinIO backend".to_string())
                })?;
                let access_key_id = config.access_key_id.clone().ok_or_else(|| {
                    S3StoreError::InvalidConfig("access_key_id is required for MinIO backend".to_string())
                })?;
                let secret_access_key = config.secret_access_key.clone().ok_or_else(|| {
                    S3StoreError::InvalidConfig(
                        "secret_access_key is required for MinIO backend".to_string(),
                    )
                })?;

                let creds = Credentials::new(access_key_id, secret_access_key, None, None, "static");
                loader = loader.endpoint_url(endpoint).credentials_provider(creds);
            }
        }

        let shared = runtime.block_on(loader.load());
        let mut s3_config_builder = aws_sdk_s3::config::Builder::from(&shared);

        if config.backend == S3Backend::Minio {
            s3_config_builder = s3_config_builder.force_path_style(true);
        }

        let client = aws_sdk_s3::Client::from_conf(s3_config_builder.build());

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

#[cfg(feature = "s3")]
impl S3CompatibleRawDataStore {
    /// Extract the object key from an `object_ref` that may be a full `s3://bucket/key` URI
    /// or a plain key path.  For example:
    ///   `s3://bucket/lift-01/1.bin`  →  `lift-01/1.bin`
    ///   `lift-01/1.bin`              →  `lift-01/1.bin`
    fn object_key<'a>(&self, object_ref: &'a str) -> &'a str {
        if let Some(rest) = object_ref.strip_prefix("s3://") {
            if let Some(slash) = rest.find('/') {
                return &rest[slash + 1..];
            }
        }
        object_ref
    }
}

#[cfg(feature = "s3")]
impl RawDataStore for S3CompatibleRawDataStore {
    type Error = S3StoreError;

    fn put(&mut self, object_ref: &str, payload: &[u8]) -> Result<(), Self::Error> {
        let key = self.object_key(object_ref);
        let stream = aws_sdk_s3::primitives::ByteStream::from(payload.to_vec());
        self.runtime
            .block_on(
                self.client
                    .put_object()
                    .bucket(&self.bucket)
                    .key(key)
                    .body(stream)
                    .send(),
            )
            .map_err(|e| S3StoreError::Put(e.to_string()))?;
        Ok(())
    }
}

// ── Async ingest traits and implementations ───────────────────────────────────

#[cfg(feature = "async-ingest")]
use std::sync::Arc;
#[cfg(feature = "async-ingest")]
use tokio::sync::Mutex as AsyncMutex;

/// Async variant of [`RawDataStore`] for use in async ingest pipelines.
///
/// Implementations must be `Send + Sync` so they can be shared across tasks.
/// The `+ Send` bound on the returned future is required for use with
/// multi-threaded tokio runtimes.
#[cfg(feature = "async-ingest")]
pub trait AsyncRawDataStore: Send + Sync {
    type Error: std::error::Error + Send;

    fn put(
        &self,
        object_ref: &str,
        payload: &[u8],
    ) -> impl std::future::Future<Output = Result<(), Self::Error>> + Send;
}

/// Async variant of [`AuditLedger`] for use in async ingest pipelines.
#[cfg(feature = "async-ingest")]
pub trait AsyncAuditLedger: Send + Sync {
    type Error: std::error::Error + Send;

    fn append(
        &self,
        record: AuditRecord,
    ) -> impl std::future::Future<Output = Result<(), Self::Error>> + Send;
}

/// Async variant of [`OperationLogStore`] for use in async ingest pipelines.
#[cfg(feature = "async-ingest")]
pub trait AsyncOperationLogStore: Send + Sync {
    type Error: std::error::Error + Send;

    fn write(
        &self,
        entry: OperationLogEntry,
    ) -> impl std::future::Future<Output = Result<(), Self::Error>> + Send;
}

// ── In-memory async implementations (for testing) ────────────────────────────

/// Async in-memory raw data store backed by a `tokio::sync::Mutex`.
#[cfg(feature = "async-ingest")]
#[derive(Default, Clone)]
pub struct AsyncInMemoryRawDataStore {
    objects: Arc<AsyncMutex<HashMap<String, Vec<u8>>>>,
}

#[cfg(feature = "async-ingest")]
impl AsyncInMemoryRawDataStore {
    pub async fn get(&self, object_ref: &str) -> Option<Vec<u8>> {
        self.objects.lock().await.get(object_ref).cloned()
    }
}

#[cfg(feature = "async-ingest")]
impl AsyncRawDataStore for AsyncInMemoryRawDataStore {
    type Error = InMemoryStoreError;

    async fn put(&self, object_ref: &str, payload: &[u8]) -> Result<(), Self::Error> {
        self.objects
            .lock()
            .await
            .insert(object_ref.to_string(), payload.to_vec());
        Ok(())
    }
}

/// Async in-memory audit ledger backed by a `tokio::sync::Mutex`.
#[cfg(feature = "async-ingest")]
#[derive(Default, Clone)]
pub struct AsyncInMemoryAuditLedger {
    records: Arc<AsyncMutex<Vec<AuditRecord>>>,
}

#[cfg(feature = "async-ingest")]
impl AsyncInMemoryAuditLedger {
    pub async fn records(&self) -> Vec<AuditRecord> {
        self.records.lock().await.clone()
    }
}

#[cfg(feature = "async-ingest")]
impl AsyncAuditLedger for AsyncInMemoryAuditLedger {
    type Error = InMemoryStoreError;

    async fn append(&self, record: AuditRecord) -> Result<(), Self::Error> {
        self.records.lock().await.push(record);
        Ok(())
    }
}

/// Async in-memory operation log backed by a `tokio::sync::Mutex`.
#[cfg(feature = "async-ingest")]
#[derive(Default, Clone)]
pub struct AsyncInMemoryOperationLog {
    entries: Arc<AsyncMutex<Vec<OperationLogEntry>>>,
}

#[cfg(feature = "async-ingest")]
impl AsyncInMemoryOperationLog {
    pub async fn entries(&self) -> Vec<OperationLogEntry> {
        self.entries.lock().await.clone()
    }
}

#[cfg(feature = "async-ingest")]
impl AsyncOperationLogStore for AsyncInMemoryOperationLog {
    type Error = InMemoryStoreError;

    async fn write(&self, entry: OperationLogEntry) -> Result<(), Self::Error> {
        self.entries.lock().await.push(entry);
        Ok(())
    }
}

// ── S3 async implementation ───────────────────────────────────────────────────

/// `S3CompatibleRawDataStore` also implements `AsyncRawDataStore` when both
/// `s3` and `async-ingest` features are active.  The async path calls `.await`
/// directly on the SDK future, bypassing the embedded synchronous runtime.
#[cfg(all(feature = "s3", feature = "async-ingest"))]
impl AsyncRawDataStore for S3CompatibleRawDataStore {
    type Error = S3StoreError;

    async fn put(&self, object_ref: &str, payload: &[u8]) -> Result<(), Self::Error> {
        let key = self.object_key(object_ref);
        let stream = aws_sdk_s3::primitives::ByteStream::from(payload.to_vec());
        self.client
            .put_object()
            .bucket(&self.bucket)
            .key(key)
            .body(stream)
            .send()
            .await
            .map_err(|e| S3StoreError::Put(e.to_string()))?;
        Ok(())
    }
}

// ── AsyncIngestService ────────────────────────────────────────────────────────

/// Async orchestration service for tamper-evident record ingestion.
///
/// Drop-in async counterpart to [`IngestService`].  All storage calls are
/// `await`-ed so the calling thread is never blocked.  The policy gate is
/// wrapped in a `tokio::sync::Mutex` to allow `&self` on `ingest`, enabling
/// the service to be shared across tasks via `Arc`.
#[cfg(feature = "async-ingest")]
pub struct AsyncIngestService<R, L, O>
where
    R: AsyncRawDataStore,
    L: AsyncAuditLedger,
    O: AsyncOperationLogStore,
{
    policy: AsyncMutex<super::policy::IntegrityPolicyGate>,
    raw_data_store: R,
    audit_ledger: L,
    operation_log: O,
}

#[cfg(feature = "async-ingest")]
impl<R, L, O> AsyncIngestService<R, L, O>
where
    R: AsyncRawDataStore,
    L: AsyncAuditLedger,
    O: AsyncOperationLogStore,
{
    pub fn new(
        policy: super::policy::IntegrityPolicyGate,
        raw_data_store: R,
        audit_ledger: L,
        operation_log: O,
    ) -> Self {
        Self {
            policy: AsyncMutex::new(policy),
            raw_data_store,
            audit_ledger,
            operation_log,
        }
    }

    pub async fn register_device(
        &self,
        device_id: impl Into<String>,
        key: ed25519_dalek::VerifyingKey,
    ) {
        self.policy.lock().await.register_device(device_id, key);
    }

    #[tracing::instrument(skip(self, raw_payload), fields(
        device_id = %record.device_id,
        sequence  = record.sequence,
        object_ref = %record.object_ref,
    ))]
    pub async fn ingest(
        &self,
        record: AuditRecord,
        raw_payload: &[u8],
        cert_identity: Option<&str>,
    ) -> Result<(), IngestServiceError> {
        debug!(payload_bytes = raw_payload.len(), "async ingest started");

        let payload_hash = compute_payload_hash(raw_payload);
        if payload_hash != record.payload_hash {
            warn!("payload hash mismatch — record rejected");
            let _ = self
                .operation_log
                .write(OperationLogEntry {
                    decision: IngestDecision::Rejected,
                    device_id: record.device_id.clone(),
                    sequence: record.sequence,
                    message: "payload hash mismatch".to_string(),
                })
                .await;
            return Err(IngestServiceError::PayloadHashMismatch {
                device_id: record.device_id,
                sequence: record.sequence,
            });
        }

        // Acquire the policy lock, run the check, then release before any I/O.
        let policy_result = {
            let mut policy = self.policy.lock().await;
            policy.enforce(&record, cert_identity)
        };
        if let Err(error) = policy_result {
            warn!(reason = %error, "integrity policy rejected record");
            let _ = self
                .operation_log
                .write(OperationLogEntry {
                    decision: IngestDecision::Rejected,
                    device_id: record.device_id.clone(),
                    sequence: record.sequence,
                    message: error.to_string(),
                })
                .await;
            return Err(IngestServiceError::Verify(error));
        }

        self.raw_data_store
            .put(&record.object_ref, raw_payload)
            .await
            .map_err(|e| {
                error!(error = %e, "raw data store write failed");
                IngestServiceError::RawDataStore(e.to_string())
            })?;

        self.audit_ledger
            .append(record.clone())
            .await
            .map_err(|e| {
                error!(error = %e, "audit ledger append failed");
                IngestServiceError::AuditLedger(e.to_string())
            })?;

        self.operation_log
            .write(OperationLogEntry {
                decision: IngestDecision::Accepted,
                device_id: record.device_id,
                sequence: record.sequence,
                message: "ingest accepted".to_string(),
            })
            .await
            .map_err(|e| {
                error!(error = %e, "operation log write failed");
                IngestServiceError::OperationLog(e.to_string())
            })?;

        info!("record accepted");
        Ok(())
    }

    pub fn raw_data_store(&self) -> &R {
        &self.raw_data_store
    }

    pub fn audit_ledger(&self) -> &L {
        &self.audit_ledger
    }

    pub fn operation_log(&self) -> &O {
        &self.operation_log
    }
}