Add atomic durable job admission

crates/remora-admission/src/error.rs

@@ -0,0 +1,30 @@
+use remora_credentials::CredentialError;
+use remora_durable::DurableError;
+use remora_identity::IdentityError;
+
+#[derive(Debug, thiserror::Error)]
+pub enum AdmissionError {
+    #[error("invalid submission: {0}")]
+    Invalid(String),
+    #[error("submission is not authorized")]
+    Unauthorized,
+    #[error("idempotency key was reused with a different request")]
+    IdempotencyConflict,
+    #[error("execution grant not found")]
+    GrantNotFound,
+    #[error("configuration error: {0}")]
+    Config(String),
+    #[error(transparent)]
+    Credential(#[from] CredentialError),
+    #[error(transparent)]
+    Durable(#[from] DurableError),
+}
+
+impl From<IdentityError> for AdmissionError {
+    fn from(error: IdentityError) -> Self {
+        match error {
+            IdentityError::UnauthorizedTarget => Self::Unauthorized,
+            other => Self::Invalid(other.to_string()),
+        }
+    }
+}

crates/remora-admission/src/lib.rs

@@ -0,0 +1,10 @@
+mod error;
+mod migrate;
+mod service;
+mod source;
+mod types;
+
+pub use error::AdmissionError;
+pub use migrate::{migrate, migrator};
+pub use service::{AdmissionConfig, AdmissionService};
+pub use types::{AcceptedJob, AdmissionDecision, AdmissionRejection, SubmitOutcome};

crates/remora-admission/src/migrate.rs

@@ -0,0 +1,17 @@
+use remora_durable::DurableError;
+use sqlx::PgPool;
+
+const MIGRATION_TABLE: &str = "_remora_admission_migrations";
+
+pub async fn migrate(pool: &PgPool) -> Result<(), DurableError> {
+    let mut migrator = sqlx::migrate!("./migrations");
+    migrator.dangerous_set_table_name(MIGRATION_TABLE);
+    migrator.run(pool).await?;
+    Ok(())
+}
+
+pub fn migrator() -> sqlx::migrate::Migrator {
+    let mut migrator = sqlx::migrate!("./migrations");
+    migrator.dangerous_set_table_name(MIGRATION_TABLE);
+    migrator
+}

crates/remora-admission/src/service.rs

@@ -0,0 +1,362 @@
+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, 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 config = self.config.clone();
+        let outcome = execute_mutation(
+            &self.pool,
+            MutationScope {
+                subject: &subject,
+                operation: "job.submit",
+                key: idempotency_key,
+            },
+            "submit_job",
+            &[],
+            request_fingerprint,
+            move |tx| {
+                Box::pin(async move {
+                    let decision = decide_submission(tx, &principal, &config, 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 revoke_execution_grant(
+        &self,
+        principal: &Principal,
+        grant_id: &OpaqueId,
+    ) -> Result<(), AdmissionError> {
+        if !may_request_execution_grant_revocation(principal) {
+            return Err(AdmissionError::Unauthorized);
+        }
+        let grant_id = grant_id.clone();
+        with_transaction(&self.pool, move |tx| {
+            Box::pin(async move {
+                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(())
+            })
+        })
+        .await
+        .map_err(|error| match &error {
+            DurableError::Effect(message) if message == "grant_not_found" => {
+                AdmissionError::GrantNotFound
+            }
+            _ => AdmissionError::Durable(error),
+        })
+    }
+}
+
+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?;
+
+    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)
+}

crates/remora-admission/src/source.rs

@@ -0,0 +1,164 @@
+use std::collections::BTreeSet;
+
+use remora_contracts::{
+    HttpSourceProjectionDto, HuggingFaceSelectionDto, HuggingFaceSourceProjectionDto,
+    ObjectStoreSourceProjectionDto, SourceProjectionDto, SourceSubmissionDto,
+};
+use remora_credentials::SourceCredentialKind;
+use remora_domain::{BucketName, MemberPath, OpaqueId, TargetStoreId};
+use serde_json::Value;
+use url::Url;
+
+use crate::AdmissionRejection;
+
+pub(crate) struct ValidatedSource {
+    pub projection: SourceProjectionDto,
+    pub private_input: Value,
+    pub credential_id: Option<OpaqueId>,
+    pub allowed_credential_kinds: &'static [SourceCredentialKind],
+}
+
+pub(crate) fn validate_source(
+    source: SourceSubmissionDto,
+) -> Result<ValidatedSource, AdmissionRejection> {
+    let private_input =
+        serde_json::to_value(&source).map_err(|_| AdmissionRejection::InvalidSource)?;
+    match source {
+        SourceSubmissionDto::Http(source) => {
+            if source.url.len() > 8192 {
+                return Err(AdmissionRejection::InvalidSource);
+            }
+            let parsed = Url::parse(&source.url).map_err(|_| AdmissionRejection::InvalidSource)?;
+            if parsed.scheme() != "https"
+                || parsed.host_str().is_none()
+                || parsed.port_or_known_default() != Some(443)
+                || !parsed.username().is_empty()
+                || parsed.password().is_some()
+                || parsed.fragment().is_some()
+            {
+                return Err(AdmissionRejection::InvalidSource);
+            }
+            let result_path = MemberPath::parse(&source.result_path)
+                .map_err(|_| AdmissionRejection::InvalidSource)?;
+            let mut display = parsed;
+            display.set_query(None);
+            display.set_fragment(None);
+            let credential_id = source.credential_id;
+            Ok(ValidatedSource {
+                projection: SourceProjectionDto::Http(HttpSourceProjectionDto {
+                    display_url: display.to_string(),
+                    result_path: result_path.canonical(),
+                    credential_id: credential_id.clone(),
+                }),
+                private_input,
+                credential_id,
+                allowed_credential_kinds: &[
+                    SourceCredentialKind::BasicAuth,
+                    SourceCredentialKind::BearerToken,
+                    SourceCredentialKind::HttpHeaders,
+                ],
+            })
+        }
+        SourceSubmissionDto::ObjectStore(source) => {
+            TargetStoreId::parse(&source.store).map_err(|_| AdmissionRejection::InvalidSource)?;
+            BucketName::parse(&source.bucket).map_err(|_| AdmissionRejection::InvalidSource)?;
+            if source.key.is_empty()
+                || source.key.len() > 4096
+                || source.key.chars().any(char::is_control)
+                || source
+                    .version_id
+                    .as_ref()
+                    .is_some_and(|version| version.is_empty() || version.len() > 1024)
+            {
+                return Err(AdmissionRejection::InvalidSource);
+            }
+            let result_path = MemberPath::parse(&source.result_path)
+                .map_err(|_| AdmissionRejection::InvalidSource)?;
+            let credential_id = source.credential_id;
+            Ok(ValidatedSource {
+                projection: SourceProjectionDto::ObjectStore(ObjectStoreSourceProjectionDto {
+                    store: source.store,
+                    bucket: source.bucket,
+                    key: source.key,
+                    version_id: source.version_id,
+                    result_path: result_path.canonical(),
+                    credential_id: credential_id.clone(),
+                }),
+                private_input,
+                credential_id,
+                allowed_credential_kinds: &[SourceCredentialKind::ObjectStore],
+            })
+        }
+        SourceSubmissionDto::HuggingFace(source) => {
+            if !valid_repository_id(&source.repository_id)
+                || source.revision.is_empty()
+                || source.revision.len() > 256
+                || source.revision.chars().any(char::is_control)
+            {
+                return Err(AdmissionRejection::InvalidSource);
+            }
+            let selection = source
+                .selection
+                .map(canonical_selection)
+                .transpose()?
+                .flatten();
+            let credential_id = source.credential_id;
+            Ok(ValidatedSource {
+                projection: SourceProjectionDto::HuggingFace(HuggingFaceSourceProjectionDto {
+                    repository_type: source.repository_type,
+                    repository_id: source.repository_id,
+                    revision: source.revision,
+                    selection,
+                    credential_id: credential_id.clone(),
+                }),
+                private_input,
+                credential_id,
+                allowed_credential_kinds: &[SourceCredentialKind::HuggingFace],
+            })
+        }
+    }
+}
+
+fn canonical_selection(
+    selection: HuggingFaceSelectionDto,
+) -> Result<Option<HuggingFaceSelectionDto>, AdmissionRejection> {
+    let paths = canonical_paths(selection.paths)?;
+    let prefixes = canonical_paths(selection.prefixes)?;
+    if paths.is_none() && prefixes.is_none() {
+        return Err(AdmissionRejection::InvalidSource);
+    }
+    Ok(Some(HuggingFaceSelectionDto { paths, prefixes }))
+}
+
+fn canonical_paths(paths: Option<Vec<String>>) -> Result<Option<Vec<String>>, AdmissionRejection> {
+    let Some(paths) = paths else {
+        return Ok(None);
+    };
+    if paths.is_empty() || paths.len() > 1024 {
+        return Err(AdmissionRejection::InvalidSource);
+    }
+    let canonical = paths
+        .into_iter()
+        .map(|path| {
+            MemberPath::parse(&path)
+                .map(|path| path.canonical())
+                .map_err(|_| AdmissionRejection::InvalidSource)
+        })
+        .collect::<Result<BTreeSet<_>, _>>()?;
+    Ok(Some(canonical.into_iter().collect()))
+}
+
+fn valid_repository_id(value: &str) -> bool {
+    let Some((owner, name)) = value.split_once('/') else {
+        return false;
+    };
+    !owner.is_empty()
+        && !name.is_empty()
+        && !name.contains('/')
+        && owner.chars().all(repository_character)
+        && name.chars().all(repository_character)
+}
+
+fn repository_character(character: char) -> bool {
+    character.is_ascii_alphanumeric() || matches!(character, '.' | '_' | '-')
+}

crates/remora-admission/src/types.rs

@@ -0,0 +1,35 @@
+use remora_contracts::{SourceProjectionDto, TargetDto};
+use remora_domain::OpaqueId;
+use remora_durable::MutationDisposition;
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum AdmissionRejection {
+    InvalidSource,
+    InvalidLabel,
+    UnauthorizedTarget,
+    CredentialUnavailable,
+    CredentialKindMismatch,
+}
+
+#[derive(Clone, Deserialize, PartialEq, Serialize)]
+pub struct AcceptedJob {
+    pub job_id: OpaqueId,
+    pub execution_grant_id: OpaqueId,
+    pub source: SourceProjectionDto,
+    pub target: TargetDto,
+}
+
+#[derive(Clone, Deserialize, PartialEq, Serialize)]
+#[serde(tag = "status", rename_all = "snake_case")]
+pub enum AdmissionDecision {
+    Accepted(Box<AcceptedJob>),
+    Rejected { reason: AdmissionRejection },
+}
+
+#[derive(Clone, PartialEq)]
+pub struct SubmitOutcome {
+    pub disposition: MutationDisposition,
+    pub decision: AdmissionDecision,
+}