Add atomic durable job admission
This commit is contained in:
parent
be651d559c
commit
b3e63b9e2b
1 changed files with 1129 additions and 0 deletions
362
crates/remora-admission/src/service.rs
Normal file
362
crates/remora-admission/src/service.rs
Normal 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)
|
||||
}
|
||||
Loading…
Reference in a new issue