remora/crates/remora-worker/src/descriptor.rs

use remora_contracts::{
    CheckpointDescriptorKindDto, CheckpointDescriptorRefDto, DurableDescriptorKindDto,
    DurableDescriptorRefDto, PositiveIntegerDto, ResolvedSourceDescriptorKindDto,
    ResolvedSourceDescriptorRefDto, SafeIntegerDto,
};
use remora_domain::OpaqueId;
use serde::Serialize;
use serde_json::Value;
use sha2::{Digest, Sha256};
use sqlx::PgPool;
use time::OffsetDateTime;

use crate::WorkerError;

#[derive(Clone)]
pub struct DescriptorStore {
    pool: PgPool,
}

impl DescriptorStore {
    #[must_use]
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }

    pub async fn store_resolved<T: Serialize>(
        &self,
        attempt_id: &OpaqueId,
        payload: &T,
        item_count: Option<u64>,
    ) -> Result<ResolvedSourceDescriptorRefDto, WorkerError> {
        let stored = self
            .store(attempt_id, "resolved_source", payload, item_count)
            .await?;
        Ok(ResolvedSourceDescriptorRefDto {
            id: stored.id,
            kind: ResolvedSourceDescriptorKindDto::ResolvedSource,
            schema_version: stored.schema_version,
            byte_length: stored.byte_length,
            item_count: stored.item_count,
            checksum: stored.checksum,
        })
    }

    pub async fn store_checkpoint<T: Serialize>(
        &self,
        attempt_id: &OpaqueId,
        payload: &T,
        item_count: Option<u64>,
    ) -> Result<CheckpointDescriptorRefDto, WorkerError> {
        let stored = self
            .store(attempt_id, "checkpoint", payload, item_count)
            .await?;
        Ok(CheckpointDescriptorRefDto {
            id: stored.id,
            kind: CheckpointDescriptorKindDto::Checkpoint,
            schema_version: stored.schema_version,
            byte_length: stored.byte_length,
            item_count: stored.item_count,
            checksum: stored.checksum,
        })
    }

    pub async fn load_ref(
        &self,
        job_id: &OpaqueId,
        descriptor: &DurableDescriptorRefDto,
    ) -> Result<Option<Value>, WorkerError> {
        type Row = (String, i32, Value, i64, Option<i64>, String);
        let row = sqlx::query_as::<_, Row>(
            "SELECT d.kind, d.schema_version, d.payload, d.byte_length, d.item_count, d.checksum
             FROM worker_descriptors d
             JOIN attempts a ON a.attempt_id = d.attempt_id
             WHERE d.descriptor_id = $1 AND a.job_id = $2",
        )
        .bind(descriptor.id.as_str())
        .bind(job_id.as_str())
        .fetch_optional(&self.pool)
        .await
        .map_err(|_| WorkerError::ApiUnavailable)?;
        let Some((kind, schema_version, payload, byte_length, item_count, checksum)) = row else {
            return Ok(None);
        };
        let encoded = canonical_bytes(&payload)?;
        let expected_kind = match descriptor.kind {
            DurableDescriptorKindDto::ResolvedSource => "resolved_source",
            DurableDescriptorKindDto::Checkpoint => "checkpoint",
        };
        if kind != expected_kind
            || u64::try_from(schema_version).ok() != Some(descriptor.schema_version.get())
            || u64::try_from(byte_length).ok() != Some(descriptor.byte_length.get())
            || item_count.map(u64::try_from).transpose().ok().flatten()
                != descriptor.item_count.map(SafeIntegerDto::get)
            || usize::try_from(byte_length).ok() != Some(encoded.len())
            || checksum != descriptor.checksum
            || digest(&encoded) != checksum
        {
            return Err(WorkerError::InvalidResponse);
        }
        Ok(Some(payload))
    }

    async fn store<T: Serialize>(
        &self,
        attempt_id: &OpaqueId,
        kind: &str,
        payload: &T,
        item_count: Option<u64>,
    ) -> Result<StoredDescriptor, WorkerError> {
        let payload = serde_json::to_value(payload).map_err(|_| WorkerError::InvalidRequest)?;
        let encoded = canonical_bytes(&payload)?;
        let id = descriptor_id();
        let checksum = digest(&encoded);
        let byte_length = SafeIntegerDto::new(
            u64::try_from(encoded.len()).map_err(|_| WorkerError::InvalidRequest)?,
        )
        .map_err(|_| WorkerError::InvalidRequest)?;
        let item_count = item_count
            .map(SafeIntegerDto::new)
            .transpose()
            .map_err(|_| WorkerError::InvalidRequest)?;
        sqlx::query(
            "INSERT INTO worker_descriptors (
                descriptor_id, attempt_id, kind, schema_version, payload,
                byte_length, item_count, checksum, created_at
             ) VALUES ($1, $2, $3, 1, $4, $5, $6, $7, $8)",
        )
        .bind(id.as_str())
        .bind(attempt_id.as_str())
        .bind(kind)
        .bind(&payload)
        .bind(i64::try_from(byte_length.get()).map_err(|_| WorkerError::InvalidRequest)?)
        .bind(
            item_count
                .map(SafeIntegerDto::get)
                .map(i64::try_from)
                .transpose()
                .map_err(|_| WorkerError::InvalidRequest)?,
        )
        .bind(&checksum)
        .bind(OffsetDateTime::now_utc())
        .execute(&self.pool)
        .await
        .map_err(|_| WorkerError::ApiUnavailable)?;
        Ok(StoredDescriptor {
            id,
            schema_version: PositiveIntegerDto::new(1).expect("one is positive"),
            byte_length,
            item_count,
            checksum,
        })
    }
}

struct StoredDescriptor {
    id: OpaqueId,
    schema_version: PositiveIntegerDto,
    byte_length: SafeIntegerDto,
    item_count: Option<SafeIntegerDto>,
    checksum: String,
}

fn canonical_bytes(value: &Value) -> Result<Vec<u8>, WorkerError> {
    serde_json_canonicalizer::to_vec(value).map_err(|_| WorkerError::InvalidRequest)
}

fn digest(bytes: &[u8]) -> String {
    let mut value = String::from("sha256:");
    for byte in Sha256::digest(bytes) {
        use std::fmt::Write as _;
        write!(value, "{byte:02x}").expect("writing to String cannot fail");
    }
    value
}

fn descriptor_id() -> OpaqueId {
    let mut bytes = [0_u8; 16];
    getrandom::fill(&mut bytes).expect("operating system randomness is required");
    let mut value = String::from("d");
    for byte in bytes {
        use std::fmt::Write as _;
        write!(value, "{byte:02x}").expect("writing to String cannot fail");
    }
    OpaqueId::parse(value).expect("generated descriptor identifier is valid")
}