remora/crates/remora-admission/src/service.rs

use remora_contracts::{JobSubmissionDto, TargetDto};
use remora_credentials::{
    CredentialCaller, CredentialError, CredentialGrantVersion, CredentialRegistry,
};
use remora_domain::{OpaqueId, Target, fingerprint, ports::MutationScope};
use remora_durable::{
    DurableError, MutationDisposition, MutationStatus, PgTx, enqueue, execute_mutation,
    with_transaction,
};
use remora_identity::{
    DefaultTarget, Principal, TargetOverride, may_request_execution_grant_revocation,
    resolve_authorized_target,
};
use serde_json::json;
use sqlx::PgPool;
use time::OffsetDateTime;

use crate::source::{ValidatedSource, validate_source};
use crate::{AcceptedJob, AdmissionDecision, AdmissionError, AdmissionRejection, SubmitOutcome};

#[derive(Clone, Debug)]
pub struct AdmissionConfig {
    pub default_target: DefaultTarget,
    pub reserved_target_segments: Vec<String>,
    pub policy_version: String,
}

#[derive(Clone, Debug)]
pub struct AdmissionService {
    pool: PgPool,
    config: AdmissionConfig,
}

impl AdmissionService {
    pub fn new(pool: PgPool, config: AdmissionConfig) -> Result<Self, AdmissionError> {
        if config.policy_version.is_empty() || config.policy_version.len() > 128 {
            return Err(AdmissionError::Config(
                "policy version length must be between 1 and 128".into(),
            ));
        }
        Ok(Self { pool, config })
    }

    pub async fn submit(
        &self,
        principal: &Principal,
        idempotency_key: &str,
        request: JobSubmissionDto,
    ) -> Result<SubmitOutcome, AdmissionError> {
        if idempotency_key.is_empty() || idempotency_key.len() > 256 {
            return Err(AdmissionError::Invalid(
                "idempotency key length must be between 1 and 256".into(),
            ));
        }
        let request_fingerprint =
            fingerprint(&request).map_err(|error| AdmissionError::Invalid(error.to_string()))?;
        let subject = principal.id.clone();
        let principal = principal.clone();
        let service = self.clone();
        let outcome = execute_mutation(
            &self.pool,
            MutationScope {
                subject: &subject,
                operation: "job.submit",
                key: idempotency_key,
            },
            "admission_submit_job",
            &[],
            request_fingerprint,
            move |tx| {
                Box::pin(async move {
                    let decision = service.submit_in_tx(tx, &principal, request).await?;
                    let status = match decision {
                        AdmissionDecision::Accepted(_) => MutationStatus::Accepted,
                        AdmissionDecision::Rejected { .. } => MutationStatus::Rejected,
                    };
                    let response = serde_json::to_value(decision)
                        .map_err(|error| DurableError::Effect(error.to_string()))?;
                    Ok((status, response))
                })
            },
        )
        .await?;
        if outcome.disposition == MutationDisposition::Conflict {
            return Err(AdmissionError::IdempotencyConflict);
        }
        let decision = serde_json::from_value(outcome.response)
            .map_err(|error| AdmissionError::Config(error.to_string()))?;
        Ok(SubmitOutcome {
            disposition: outcome.disposition,
            decision,
        })
    }

    pub async fn submit_in_tx(
        &self,
        tx: &mut PgTx<'_>,
        principal: &Principal,
        request: JobSubmissionDto,
    ) -> Result<AdmissionDecision, DurableError> {
        decide_submission(tx, principal, &self.config, request).await
    }

    pub async fn revoke_execution_grant(
        &self,
        principal: &Principal,
        grant_id: &OpaqueId,
    ) -> Result<(), AdmissionError> {
        if !may_request_execution_grant_revocation(principal) {
            return Err(AdmissionError::Unauthorized);
        }
        let principal = principal.clone();
        let grant_id = grant_id.clone();
        with_transaction(&self.pool, move |tx| {
            Box::pin(
                async move { Self::revoke_execution_grant_in_tx(tx, &principal, &grant_id).await },
            )
        })
        .await
        .map_err(|error| match &error {
            DurableError::Effect(message) if message == "grant_not_found" => {
                AdmissionError::GrantNotFound
            }
            _ => AdmissionError::Durable(error),
        })
    }

    pub async fn revoke_execution_grant_in_tx(
        tx: &mut PgTx<'_>,
        principal: &Principal,
        grant_id: &OpaqueId,
    ) -> Result<(), DurableError> {
        if !may_request_execution_grant_revocation(principal) {
            return Err(DurableError::Effect("grant_unauthorized".into()));
        }
        let row = sqlx::query_as::<_, (String, Option<String>)>(
            "SELECT status, credential_id
             FROM execution_grants
             WHERE grant_id = $1
             FOR UPDATE",
        )
        .bind(grant_id.as_str())
        .fetch_optional(tx.as_mut())
        .await?
        .ok_or_else(|| DurableError::Effect("grant_not_found".into()))?;
        if row.0 == "revoked" {
            return Ok(());
        }
        let now = OffsetDateTime::now_utc();
        sqlx::query(
            "UPDATE execution_grants
             SET status = 'revoked', revoked_at = $2
             WHERE grant_id = $1 AND status = 'active'",
        )
        .bind(grant_id.as_str())
        .bind(now)
        .execute(tx.as_mut())
        .await?;
        if row.1.is_some() {
            let changed = sqlx::query(
                "UPDATE execution_grant_credential_pins
                 SET status = 'revoked', revoked_at = COALESCE(revoked_at, $2)
                 WHERE grant_id = $1 AND status = 'active'",
            )
            .bind(grant_id.as_str())
            .bind(now)
            .execute(tx.as_mut())
            .await?;
            if changed.rows_affected() != 1 {
                return Err(DurableError::Effect(
                    "active credential grant is missing its pin".into(),
                ));
            }
        }
        enqueue(
            tx,
            "execution_grant.revoked",
            json!({ "grant_id": grant_id }),
        )
        .await?;
        Ok(())
    }
}

async fn decide_submission(
    tx: &mut remora_durable::PgTx<'_>,
    principal: &Principal,
    config: &AdmissionConfig,
    request: JobSubmissionDto,
) -> Result<AdmissionDecision, DurableError> {
    if request
        .label
        .as_ref()
        .is_some_and(|label| label.is_empty() || label.len() > 200)
    {
        return Ok(rejected(AdmissionRejection::InvalidLabel));
    }
    let target_override = request
        .target
        .map(|target| TargetOverride {
            bucket: target.bucket,
            prefix: target.prefix,
        })
        .unwrap_or_default();
    let reserved = config
        .reserved_target_segments
        .iter()
        .map(String::as_str)
        .collect::<Vec<_>>();
    let target = match resolve_authorized_target(
        principal,
        &config.default_target,
        &target_override,
        &reserved,
    ) {
        Ok(target) => target,
        Err(remora_identity::IdentityError::UnauthorizedTarget) => {
            return Ok(rejected(AdmissionRejection::UnauthorizedTarget));
        }
        Err(_) => return Ok(rejected(AdmissionRejection::UnauthorizedTarget)),
    };
    let source = match validate_source(request.source) {
        Ok(source) => source,
        Err(reason) => return Ok(rejected(reason)),
    };
    let caller = CredentialCaller {
        principal_id: principal.id.clone(),
        administrator: principal.administrator,
    };
    let credential = match pin_credential(tx, &caller, &source).await {
        Ok(pin) => pin,
        Err(reason) => return Ok(rejected(reason)),
    };

    let job_id = new_id('j')?;
    let grant_id = new_id('g')?;
    let sensitive_source_ref = new_id('s')?;
    let now = OffsetDateTime::now_utc();
    insert_accepted(
        tx,
        AcceptedInsert {
            job_id: &job_id,
            grant_id: &grant_id,
            sensitive_source_ref: &sensitive_source_ref,
            principal,
            config,
            label: request.label.as_deref(),
            source: &source,
            target: &target,
            credential: credential.as_ref(),
            now,
        },
    )
    .await?;
    if let Some(pin) = &credential {
        CredentialRegistry::bind_execution_grant_in_tx(tx, &grant_id, pin, now)
            .await
            .map_err(DurableError::effect)?;
    }
    enqueue(tx, "job.dispatch", json!({ "job_id": job_id })).await?;
    enqueue(tx, "job.updated", json!({ "job_id": job_id })).await?;

    Ok(AdmissionDecision::Accepted(Box::new(AcceptedJob {
        job_id,
        execution_grant_id: grant_id,
        source: source.projection,
        target: target_dto(&target),
    })))
}

async fn pin_credential(
    tx: &mut remora_durable::PgTx<'_>,
    caller: &CredentialCaller,
    source: &ValidatedSource,
) -> Result<Option<CredentialGrantVersion>, AdmissionRejection> {
    let Some(credential_id) = &source.credential_id else {
        return Ok(None);
    };
    let (pin, kind) =
        CredentialRegistry::authorize_and_pin_current_version_in_tx(tx, caller, credential_id)
            .await
            .map_err(|error| match error {
                CredentialError::Durable(_) => AdmissionRejection::CredentialUnavailable,
                _ => AdmissionRejection::CredentialUnavailable,
            })?;
    if !source.allowed_credential_kinds.contains(&kind) {
        return Err(AdmissionRejection::CredentialKindMismatch);
    }
    Ok(Some(pin))
}

struct AcceptedInsert<'a> {
    job_id: &'a OpaqueId,
    grant_id: &'a OpaqueId,
    sensitive_source_ref: &'a OpaqueId,
    principal: &'a Principal,
    config: &'a AdmissionConfig,
    label: Option<&'a str>,
    source: &'a ValidatedSource,
    target: &'a Target,
    credential: Option<&'a CredentialGrantVersion>,
    now: OffsetDateTime,
}

async fn insert_accepted(
    tx: &mut remora_durable::PgTx<'_>,
    accepted: AcceptedInsert<'_>,
) -> Result<(), DurableError> {
    let source_projection = serde_json::to_value(&accepted.source.projection)
        .map_err(|error| DurableError::Effect(error.to_string()))?;
    let target_prefix = serde_json::to_value(accepted.target.prefix.segments())
        .map_err(|error| DurableError::Effect(error.to_string()))?;
    sqlx::query(
        "INSERT INTO sensitive_source_inputs (sensitive_source_ref, source_input, created_at)
         VALUES ($1, $2, $3)",
    )
    .bind(accepted.sensitive_source_ref.as_str())
    .bind(&accepted.source.private_input)
    .bind(accepted.now)
    .execute(tx.as_mut())
    .await?;
    sqlx::query(
        "INSERT INTO execution_grants (
            grant_id, owner_id, target_store, target_bucket, target_prefix,
            policy_version, credential_id, credential_version_id, status, created_at, revoked_at
         ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, 'active', $9, NULL)",
    )
    .bind(accepted.grant_id.as_str())
    .bind(accepted.principal.id.as_str())
    .bind(accepted.target.store.as_str())
    .bind(accepted.target.bucket.as_str())
    .bind(&target_prefix)
    .bind(&accepted.config.policy_version)
    .bind(accepted.credential.map(|pin| pin.credential_id.as_str()))
    .bind(accepted.credential.map(|pin| pin.version_id.as_str()))
    .bind(accepted.now)
    .execute(tx.as_mut())
    .await?;
    sqlx::query(
        "INSERT INTO jobs (
            job_id, owner_id, label, revision, phase, lifecycle, source_projection,
            sensitive_source_ref, target_store, target_bucket, target_prefix,
            execution_grant_id, dispatch_needed, created_at, updated_at
         ) VALUES ($1, $2, $3, 1, 'queued', 'visible', $4, $5, $6, $7, $8, $9, TRUE, $10, $10)",
    )
    .bind(accepted.job_id.as_str())
    .bind(accepted.principal.id.as_str())
    .bind(accepted.label)
    .bind(source_projection)
    .bind(accepted.sensitive_source_ref.as_str())
    .bind(accepted.target.store.as_str())
    .bind(accepted.target.bucket.as_str())
    .bind(target_prefix)
    .bind(accepted.grant_id.as_str())
    .bind(accepted.now)
    .execute(tx.as_mut())
    .await?;
    Ok(())
}

fn rejected(reason: AdmissionRejection) -> AdmissionDecision {
    AdmissionDecision::Rejected { reason }
}

fn target_dto(target: &Target) -> TargetDto {
    TargetDto {
        store: target.store.as_str().to_owned(),
        bucket: target.bucket.as_str().to_owned(),
        prefix: target.prefix.canonical(),
    }
}

fn new_id(prefix: char) -> Result<OpaqueId, DurableError> {
    let mut bytes = [0u8; 16];
    getrandom::fill(&mut bytes).map_err(|error| {
        DurableError::Config(format!("failed to draw identifier entropy: {error}"))
    })?;
    let mut value = String::with_capacity(33);
    value.push(prefix);
    for byte in bytes {
        use std::fmt::Write as _;
        write!(value, "{byte:02x}").expect("writing to String cannot fail");
    }
    OpaqueId::parse(value).map_err(DurableError::effect)
}