ms2k/remora
SHA256
Watch
0
Code Releases Activity

Add atomic durable job admission

Browse source
This commit is contained in:
ms2k 2026-07-17 03:20:57 +09:00
commit b3e63b9e2b
9 changed files with 1129 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

23
crates/remora-admission/Cargo.toml Normal file
View file

@ -0,0 +1,23 @@
[package]
name = "remora-admission"
edition.workspace = true
publish = false
rust-version.workspace = true
version.workspace = true
[dependencies]
getrandom = "0.3.4"
remora-contracts = { path = "../remora-contracts" }
remora-credentials = { path = "../remora-credentials" }
remora-domain.workspace = true
remora-durable.workspace = true
remora-identity = { path = "../remora-identity" }
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
thiserror.workspace = true
time.workspace = true
url = "2.5.8"
[dev-dependencies]
tokio.workspace = true

62
crates/remora-admission/migrations/0001_admission.sql Normal file
View file

@ -0,0 +1,62 @@
CREATE TABLE sensitive_source_inputs (
sensitive_source_ref TEXT PRIMARY KEY,
source_input JSONB NOT NULL,
created_at TIMESTAMPTZ NOT NULL
);
CREATE TABLE execution_grants (
grant_id TEXT PRIMARY KEY,
owner_id TEXT NOT NULL,
target_store TEXT NOT NULL,
target_bucket TEXT NOT NULL,
target_prefix JSONB NOT NULL,
policy_version TEXT NOT NULL,
credential_id TEXT,
credential_version_id TEXT,
status TEXT NOT NULL,
created_at TIMESTAMPTZ NOT NULL,
revoked_at TIMESTAMPTZ,
CONSTRAINT execution_grants_status_check CHECK (status IN ('active', 'revoked')),
CONSTRAINT execution_grants_revocation_check CHECK (
(status = 'active' AND revoked_at IS NULL)
OR (status = 'revoked' AND revoked_at IS NOT NULL)
),
CONSTRAINT execution_grants_credential_check CHECK (
(credential_id IS NULL) = (credential_version_id IS NULL)
),
CONSTRAINT execution_grants_id_status_unique UNIQUE (grant_id, status)
);
ALTER TABLE execution_grant_credential_pins
ADD CONSTRAINT execution_grant_pins_authoritative_grant_fk
FOREIGN KEY (grant_id, status)
REFERENCES execution_grants (grant_id, status)
DEFERRABLE INITIALLY DEFERRED;
CREATE TABLE jobs (
job_id TEXT PRIMARY KEY,
owner_id TEXT NOT NULL,
label TEXT,
revision BIGINT NOT NULL DEFAULT 1 CHECK (revision >= 1),
phase TEXT NOT NULL CHECK (phase IN ('queued', 'running', 'retry_wait', 'cancelling', 'succeeded', 'failed', 'cancelled')),
lifecycle TEXT NOT NULL CHECK (lifecycle IN ('visible', 'soft_deleted', 'hard_deleting', 'hard_deleted')),
source_projection JSONB NOT NULL,
sensitive_source_ref TEXT NOT NULL UNIQUE REFERENCES sensitive_source_inputs (sensitive_source_ref),
target_store TEXT NOT NULL,
target_bucket TEXT NOT NULL,
target_prefix JSONB NOT NULL,
execution_grant_id TEXT NOT NULL UNIQUE REFERENCES execution_grants (grant_id) DEFERRABLE INITIALLY DEFERRED,
dispatch_needed BOOLEAN NOT NULL,
created_at TIMESTAMPTZ NOT NULL,
updated_at TIMESTAMPTZ NOT NULL,
CONSTRAINT jobs_label_length CHECK (label IS NULL OR char_length(label) BETWEEN 1 AND 200)
);
ALTER TABLE execution_grants
ADD CONSTRAINT execution_grants_job_fk
FOREIGN KEY (grant_id)
REFERENCES jobs (execution_grant_id)
DEFERRABLE INITIALLY DEFERRED;
CREATE INDEX jobs_owner_created_idx ON jobs (owner_id, created_at DESC, job_id DESC);
CREATE INDEX jobs_dispatch_idx ON jobs (created_at, job_id) WHERE dispatch_needed;

30
crates/remora-admission/src/error.rs Normal file
View file

@ -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()),
}
}
}

10
crates/remora-admission/src/lib.rs Normal file
View file

@ -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};

17
crates/remora-admission/src/migrate.rs Normal file
View file

@ -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
}

362
crates/remora-admission/src/service.rs Normal file
View file

@ -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)
}

164
crates/remora-admission/src/source.rs Normal file
View file

@ -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, '.' | '_' | '-')
}

35
crates/remora-admission/src/types.rs Normal file
View file

@ -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,
}

348
crates/remora-admission/tests/substrate.rs Normal file
View file

@ -0,0 +1,348 @@
use std::time::Duration;
use remora_admission::{
AdmissionConfig, AdmissionDecision, AdmissionError, AdmissionService, SubmitOutcome, migrate,
};
use remora_contracts::{
HttpSourceSubmissionDto, JobSubmissionDto, SourceSubmissionDto, TargetOverrideDto,
};
use remora_credentials::{CredentialRegistry, migrate as credentials_migrate};
use remora_domain::{BucketName, OpaqueId, StorageScope, TargetPrefix, TargetStoreId};
use remora_durable::{MutationDisposition, PoolConfig, connect, migrate as durable_migrate};
use remora_identity::{DefaultTarget, Principal};
use sqlx::PgPool;
fn database_url() -> Option<String> {
std::env::var("DATABASE_URL")
.ok()
.filter(|value| !value.is_empty())
}
async fn pool() -> Option<PgPool> {
let Some(url) = database_url() else {
if std::env::var("REMORA_ADMISSION_REQUIRE").is_ok() {
panic!("DATABASE_URL is required when REMORA_ADMISSION_REQUIRE is set");
}
return None;
};
let pool = connect(
&url,
PoolConfig {
max_connections: 8,
acquire_timeout: Duration::from_secs(10),
idle_timeout: Some(Duration::from_secs(30)),
},
)
.await
.expect("connect postgres");
durable_migrate(&pool).await.expect("durable migrate");
credentials_migrate(&pool)
.await
.expect("credentials migrate");
migrate(&pool).await.expect("admission migrate");
Some(pool)
}
fn unique_id(prefix: &str) -> OpaqueId {
let nonce = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos();
OpaqueId::parse(format!("{prefix}{nonce:032x}")).unwrap()
}
fn target_parts() -> (TargetStoreId, BucketName, TargetPrefix) {
(
TargetStoreId::parse("primary").unwrap(),
BucketName::parse("remora-results").unwrap(),
TargetPrefix::parse("users", ["_remora"]).unwrap(),
)
}
fn principal(id: OpaqueId, authorized: bool, administrator: bool) -> Principal {
let (store, bucket, prefix) = target_parts();
Principal {
id,
issuer: "https://issuer.example".into(),
subject: "subject".into(),
administrator,
scopes: authorized
.then_some(StorageScope {
store,
bucket,
prefix,
})
.into_iter()
.collect(),
}
}
fn service(pool: PgPool) -> AdmissionService {
let (store, bucket, prefix) = target_parts();
AdmissionService::new(
pool,
AdmissionConfig {
default_target: DefaultTarget {
store,
bucket,
prefix,
},
reserved_target_segments: vec!["_remora".into()],
policy_version: "source-policy-v1".into(),
},
)
.unwrap()
}
fn request(url: String) -> JobSubmissionDto {
JobSubmissionDto {
source: SourceSubmissionDto::Http(HttpSourceSubmissionDto {
url,
result_path: "artifact.bin".into(),
credential_id: None,
}),
target: Some(TargetOverrideDto {
bucket: None,
prefix: Some("users/alice".into()),
}),
label: Some("artifact".into()),
}
}
fn accepted(outcome: SubmitOutcome) -> remora_admission::AcceptedJob {
match outcome.decision {
AdmissionDecision::Accepted(job) => *job,
AdmissionDecision::Rejected { .. } => panic!("submission rejected"),
}
}
#[tokio::test]
async fn exact_replay_and_lost_response_recovery_create_one_durable_acceptance() {
let Some(pool) = pool().await else { return };
let service = service(pool.clone());
let owner = principal(unique_id("owner"), true, false);
let key = unique_id("key").to_string();
let submission = request("https://downloads.example/archive?token=private".into());
let first = service.submit(&owner, &key, submission).await.unwrap();
assert_eq!(first.disposition, MutationDisposition::First);
let first_job = accepted(first);
let replay = service
.submit(
&owner,
&key,
request("https://downloads.example/archive?token=private".into()),
)
.await
.unwrap();
assert_eq!(replay.disposition, MutationDisposition::Replay);
let replay_job = accepted(replay);
assert_eq!(first_job.job_id, replay_job.job_id);
assert_eq!(first_job.execution_grant_id, replay_job.execution_grant_id);
let counts: (i64, i64, i64) = sqlx::query_as(
"SELECT
(SELECT count(*) FROM jobs WHERE job_id = $1),
(SELECT count(*) FROM execution_grants WHERE grant_id = $2),
(SELECT count(*) FROM outbox_messages WHERE topic = 'job.dispatch' AND payload->>'job_id' = $1)",
)
.bind(first_job.job_id.as_str())
.bind(first_job.execution_grant_id.as_str())
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(counts, (1, 1, 1));
}
#[tokio::test]
async fn changed_fingerprint_conflicts_and_current_authorization_is_rechecked_for_new_keys() {
let Some(pool) = pool().await else { return };
let service = service(pool);
let owner_id = unique_id("owner");
let owner = principal(owner_id.clone(), true, false);
let key = unique_id("key").to_string();
service
.submit(
&owner,
&key,
request("https://downloads.example/one".into()),
)
.await
.unwrap();
assert!(matches!(
service
.submit(
&owner,
&key,
request("https://downloads.example/two".into())
)
.await,
Err(AdmissionError::IdempotencyConflict)
));
let lost_authorization = principal(owner_id, false, false);
let replay = service
.submit(
&lost_authorization,
&key,
request("https://downloads.example/one".into()),
)
.await
.unwrap();
assert_eq!(replay.disposition, MutationDisposition::Replay);
let fresh = service
.submit(
&lost_authorization,
unique_id("key").as_str(),
request("https://downloads.example/one".into()),
)
.await
.unwrap();
assert!(matches!(fresh.decision, AdmissionDecision::Rejected { .. }));
}
#[tokio::test]
async fn public_projection_and_outbox_do_not_leak_sensitive_url() {
let Some(pool) = pool().await else { return };
let service = service(pool.clone());
let owner = principal(unique_id("owner"), true, false);
let secret = unique_id("secret").to_string();
let accepted = accepted(
service
.submit(
&owner,
unique_id("key").as_str(),
request(format!(
"https://downloads.example/archive?token={secret}&signature=hidden"
)),
)
.await
.unwrap(),
);
let public_text: String = sqlx::query_scalar(
"SELECT concat_ws(' ',
(SELECT source_projection::text FROM jobs WHERE job_id = $1),
(SELECT payload::text FROM outbox_messages WHERE topic = 'job.dispatch' AND payload->>'job_id' = $1),
(SELECT response_body::text FROM mutation_records WHERE response_body->'accepted'->>'job_id' = $1)
)",
)
.bind(accepted.job_id.as_str())
.fetch_one(&pool)
.await
.unwrap();
assert!(!public_text.contains(&secret));
assert!(!public_text.contains("signature"));
let private_text: String = sqlx::query_scalar(
"SELECT source_input::text FROM sensitive_source_inputs
WHERE sensitive_source_ref = (SELECT sensitive_source_ref FROM jobs WHERE job_id = $1)",
)
.bind(accepted.job_id.as_str())
.fetch_one(&pool)
.await
.unwrap();
assert!(private_text.contains(&secret));
}
#[tokio::test]
async fn credential_pin_and_accepted_grant_revoke_atomically_without_reactivation() {
let Some(pool) = pool().await else { return };
let service = service(pool.clone());
let owner_id = unique_id("owner");
let credential_id = unique_id("cred");
let version_id = unique_id("version");
let now = time::OffsetDateTime::now_utc();
let mut tx = pool.begin().await.unwrap();
sqlx::query(
"INSERT INTO source_credentials (
credential_id, owner_id, label, kind, status, active_version_id,
secret_mount, secret_path, created_at, updated_at
) VALUES ($1, $2, 'fixture', 'bearer_token', 'active', $3, 'secret', 'fixture', $4, $4)",
)
.bind(credential_id.as_str())
.bind(owner_id.as_str())
.bind(version_id.as_str())
.bind(now)
.execute(&mut *tx)
.await
.unwrap();
sqlx::query(
"INSERT INTO source_credential_versions (
version_id, credential_id, vault_kv_version, created_at, destroyed_at, secret_mount, secret_path
) VALUES ($1, $2, 1, $3, NULL, 'secret', 'fixture/version')",
)
.bind(version_id.as_str())
.bind(credential_id.as_str())
.bind(now)
.execute(&mut *tx)
.await
.unwrap();
tx.commit().await.unwrap();
let mut request = request("https://downloads.example/private".into());
let SourceSubmissionDto::Http(source) = &mut request.source else {
unreachable!()
};
source.credential_id = Some(credential_id.clone());
let accepted = accepted(
service
.submit(
&principal(owner_id, true, false),
unique_id("key").as_str(),
request,
)
.await
.unwrap(),
);
let binding: (String, String, String) = sqlx::query_as(
"SELECT credential_id, version_id, status
FROM execution_grant_credential_pins WHERE grant_id = $1",
)
.bind(accepted.execution_grant_id.as_str())
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
binding,
(
credential_id.to_string(),
version_id.to_string(),
"active".into()
)
);
let admin = principal(unique_id("admin"), false, true);
let (left, right) = tokio::join!(
service.revoke_execution_grant(&admin, &accepted.execution_grant_id),
service.revoke_execution_grant(&admin, &accepted.execution_grant_id)
);
left.unwrap();
right.unwrap();
let statuses: (String, String, i64) = sqlx::query_as(
"SELECT
(SELECT status FROM execution_grants WHERE grant_id = $1),
(SELECT status FROM execution_grant_credential_pins WHERE grant_id = $1),
(SELECT count(*) FROM outbox_messages WHERE topic = 'execution_grant.revoked' AND payload->>'grant_id' = $1)",
)
.bind(accepted.execution_grant_id.as_str())
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(statuses, ("revoked".into(), "revoked".into(), 1));
let mut tx = pool.begin().await.unwrap();
let pin = remora_credentials::CredentialGrantVersion {
credential_id,
version_id,
};
assert!(
CredentialRegistry::bind_execution_grant_in_tx(
&mut tx,
&accepted.execution_grant_id,
&pin,
time::OffsetDateTime::now_utc(),
)
.await
.is_err()
);
tx.rollback().await.unwrap();
}