use std::{net::TcpListener, sync::Arc, time::Duration}; use actix_web::{App, HttpResponse, HttpServer, web}; use rcgen::{ BasicConstraints, Certificate, CertificateParams, ExtendedKeyUsagePurpose, IsCa, Issuer, KeyPair, KeyUsagePurpose, }; use remora_admission::{AdmissionConfig, AdmissionDecision, AdmissionService}; use remora_api::{ AuthenticatedWorker, WorkerApiState, authenticate_rustls_worker, configure_worker_api, }; use remora_attempt::{AttemptConfig, AttemptService, RecoveryConfig}; use remora_contracts::{ ApiErrorDto, AttemptActivateRequestDto, AttemptCheckpointRequestDto, AttemptFailureOutcomeDto, AttemptOutcomeDto, AttemptOutcomeRequestDto, AttemptProgressRequestDto, AttemptRenewRequestDto, AttemptResolvedSourceRequestDto, AttemptRuntimeFailureDto, AttemptStagingMultipartRequestDto, AttemptStagingObjectRequestDto, AttemptStagingWriteIntentRequestDto, AttemptStopRequestDto, CheckpointDescriptorKindDto, CheckpointDescriptorRefDto, ClaimRequestDto, ClaimResponseDto, HttpSourceSubmissionDto, JobPhaseProjectionDto, JobSubmissionDto, PositiveIntegerDto, ProgressValueDto, ResolvedSourceDescriptorKindDto, ResolvedSourceDescriptorRefDto, ResultDescriptorDto, ResultKindDto, SafeIntegerDto, SourceSubmissionDto, StagingMultipartIdentityDto, StagingObjectIdentityDto, StagingWriteIntentDto, }; use remora_coordination::{CoordinationStore, StoreConfig}; use remora_domain::{ BucketName, JobPhase, OpaqueId, RuntimeFailureClass, StorageScope, TargetPrefix, TargetStoreId, }; use remora_durable::{PoolConfig, connect}; use remora_identity::{DefaultTarget, Principal}; use remora_schema::migrate as migrate_schema; use remora_target::attempt_manifest_key_in; use reqwest::{Certificate as ClientCertificate, Client, Identity}; use rustls::{ RootCertStore, ServerConfig, pki_types::{PrivateKeyDer, PrivatePkcs8KeyDer}, server::WebPkiClientVerifier, }; use sqlx::PgPool; struct Harness { pool: PgPool, coordination: CoordinationStore, } impl Harness { async fn new(name: &str) -> Option { let database_url = required_env("DATABASE_URL")?; let coordination_url = required_env("REMORA_COORDINATION_URL")?; let pool = connect( &database_url, PoolConfig { max_connections: 16, acquire_timeout: Duration::from_secs(10), idle_timeout: Some(Duration::from_secs(30)), }, ) .await .unwrap(); migrate_schema(&pool).await.unwrap(); sqlx::query("TRUNCATE TABLE jobs CASCADE") .execute(&pool) .await .unwrap(); let coordination = CoordinationStore::connect( &coordination_url, format!("{}wp18:{name}:{}:", coordination_namespace(), nonce()), StoreConfig { global_limit: 16, lease_ttl: Duration::from_secs(10), reserve_ttl: Duration::from_secs(10), progress_ttl: Duration::from_secs(10), cancel_ttl: None, }, ) .await .unwrap(); Some(Self { pool, coordination }) } fn service(&self) -> AttemptService { AttemptService::new( self.pool.clone(), self.coordination.clone(), AttemptConfig { max_attempts: 3 }, ) .unwrap() } async fn job(&self) -> OpaqueId { let store = TargetStoreId::parse("primary").unwrap(); let bucket = BucketName::parse("remora-results").unwrap(); let prefix = TargetPrefix::parse("users/alice", ["_remora"]).unwrap(); let principal = Principal { id: unique_id("owner"), issuer: "https://issuer.example".into(), subject: "subject".into(), administrator: false, scopes: [StorageScope { store: store.clone(), bucket: bucket.clone(), prefix: prefix.clone(), }] .into_iter() .collect(), }; let service = AdmissionService::new( self.pool.clone(), AdmissionConfig { default_target: DefaultTarget { store, bucket, prefix, }, reserved_target_segments: vec!["_remora".into()], policy_version: "worker-api-v1".into(), }, ) .unwrap(); let outcome = service .submit( &principal, unique_id("submit").as_str(), JobSubmissionDto { source: SourceSubmissionDto::Http(HttpSourceSubmissionDto { url: "https://example.invalid/artifact".into(), result_path: "artifact.bin".into(), credential_id: None, }), target: None, label: Some("worker-api-test".into()), }, ) .await .unwrap(); let AdmissionDecision::Accepted(job) = outcome.decision else { panic!("job rejected") }; self.coordination .enqueue(job.job_id.as_str()) .await .unwrap(); job.job_id } async fn seal_result(&self, attempt_id: &OpaqueId) -> ResultDescriptorDto { let manifest_id = unique_id("manifest"); let checksum = format!("sha256:{}", "11".repeat(32)); let key = attempt_manifest_key_in("users/alice", attempt_id, &manifest_id); sqlx::query( "INSERT INTO result_inventories ( attempt_id, state, manifest_schema_version, manifest_id, manifest_store, manifest_bucket, manifest_key, manifest_version_id, manifest_sha256, manifest_size, member_count, total_bytes, staging_origin_attempt_id, target_prefix ) VALUES ($1, 'sealed', 1, $2, 'primary', 'remora-results', $3, 'version-1', $4, 128, 1, 4096, $1, 'users/alice')", ) .bind(attempt_id.as_str()) .bind(manifest_id.as_str()) .bind(&key) .bind(&checksum) .execute(&self.pool) .await .unwrap(); ResultDescriptorDto { kind: ResultKindDto::Manifest, manifest_id, schema_version: positive(1), manifest: remora_contracts::ObjectIdentityDto { store: "primary".into(), bucket: "remora-results".into(), key, version_id: Some("version-1".into()), checksum, size: safe(128), }, member_count: positive(1), total_bytes: safe(4096), } } } #[tokio::test] async fn worker_operations_are_authenticated_fenced_replayable_and_redacted() { let Some(harness) = Harness::new("flow").await else { return; }; let state = WorkerApiState::new(harness.service()); let (base, server, client) = start_authenticated_server(state).await; let job_id = harness.job().await; let worker = unique_id("worker"); let claim_operation = unique_id("operation"); let claim = ClaimRequestDto { operation_id: claim_operation.clone(), worker_incarnation: worker.clone(), }; let unauthenticated = actix_web::test::init_service(App::new().route( "/", web::post().to(|_: AuthenticatedWorker| async { HttpResponse::NoContent().finish() }), )) .await; let response = actix_web::test::call_service( &unauthenticated, actix_web::test::TestRequest::post().uri("/").to_request(), ) .await; assert_eq!(response.status(), 401); assert_eq!( response.headers().get("content-type").unwrap(), "application/problem+json" ); let claim_url = format!("{base}/internal/v1/jobs/{job_id}/claim"); let response = client.post(&claim_url).json(&claim).send().await.unwrap(); assert_eq!(response.status(), 200); let claimed: ClaimResponseDto = response.json().await.unwrap(); assert_eq!(claimed.authority.worker_incarnation, worker); assert_eq!(claimed.execution.attempt_id, claimed.authority.attempt_id); assert_eq!(claimed.execution.job_id, job_id); assert!( claimed .execution .staging_target .prefix .contains(claimed.authority.attempt_id.as_str()) ); let changed_claim = client .post(&claim_url) .json(&ClaimRequestDto { operation_id: claim_operation, worker_incarnation: unique_id("worker"), }) .send() .await .unwrap(); assert_eq!(changed_claim.status(), 409); assert_eq!( changed_claim.headers().get("content-type").unwrap(), "application/problem+json" ); let attempt_url = format!( "{base}/internal/v1/attempts/{}", claimed.authority.attempt_id ); let activation_operation = unique_id("operation"); let activation = AttemptActivateRequestDto { operation_id: activation_operation, fence: claimed.authority.fence, worker_incarnation: worker.clone(), }; let response = client .post(format!("{attempt_url}/activate")) .json(&activation) .send() .await .unwrap(); assert_eq!(response.status(), 200); let activation_body: serde_json::Value = response.json().await.unwrap(); let activated: remora_contracts::ActivationResponseDto = serde_json::from_value(activation_body.clone()).unwrap(); assert!(activated.lease.renew_after.as_str() < activated.lease.expires_at.as_str()); let renewal_operation = unique_id("operation"); let renewal = AttemptRenewRequestDto { operation_id: renewal_operation.clone(), fence: claimed.authority.fence, worker_incarnation: worker.clone(), }; let first_renewal = client .post(format!("{attempt_url}/renew")) .json(&renewal) .send() .await .unwrap(); assert_eq!(first_renewal.status(), 200); let first_lease: serde_json::Value = first_renewal.json().await.unwrap(); let replayed_lease: serde_json::Value = client .post(format!("{attempt_url}/renew")) .json(&renewal) .send() .await .unwrap() .json() .await .unwrap(); assert_eq!(replayed_lease, first_lease); let replayed_activation: serde_json::Value = client .post(format!("{attempt_url}/activate")) .json(&activation) .send() .await .unwrap() .json() .await .unwrap(); assert_eq!(replayed_activation, activation_body); let authority_mismatches = [ ( format!("{base}/internal/v1/attempts/{}/renew", unique_id("attempt")), claimed.authority.fence, worker.clone(), ), ( format!("{attempt_url}/renew"), positive(claimed.authority.fence.get() + 1), worker.clone(), ), ( format!("{attempt_url}/renew"), claimed.authority.fence, unique_id("worker"), ), ]; for (url, fence, worker_incarnation) in authority_mismatches { let response = client .post(url) .json(&AttemptRenewRequestDto { operation_id: unique_id("operation"), fence, worker_incarnation, }) .send() .await .unwrap(); assert_eq!(response.status(), 409); } let staging_key = format!( "{}o/{}", claimed.execution.staging_target.prefix, unique_id("object") ); let staging_intent = AttemptStagingWriteIntentRequestDto { operation_id: unique_id("operation"), fence: claimed.authority.fence, worker_incarnation: worker.clone(), intent: StagingWriteIntentDto { store: claimed.execution.target.store.clone(), bucket: claimed.execution.target.bucket.clone(), key: staging_key.clone(), }, }; for _ in 0..2 { assert_eq!( client .post(format!("{attempt_url}/staging/write-intent")) .json(&staging_intent) .send() .await .unwrap() .status(), 204 ); } let staging_multipart = AttemptStagingMultipartRequestDto { operation_id: unique_id("operation"), fence: claimed.authority.fence, worker_incarnation: worker.clone(), multipart: StagingMultipartIdentityDto { store: claimed.execution.target.store.clone(), bucket: claimed.execution.target.bucket.clone(), key: staging_key.clone(), upload_id: "upload-identity".into(), }, }; assert_eq!( client .post(format!("{attempt_url}/staging/multipart")) .json(&staging_multipart) .send() .await .unwrap() .status(), 204 ); let staging_operation = unique_id("operation"); let staging_object = AttemptStagingObjectRequestDto { operation_id: staging_operation.clone(), fence: claimed.authority.fence, worker_incarnation: worker.clone(), object: StagingObjectIdentityDto { store: claimed.execution.target.store.clone(), bucket: claimed.execution.target.bucket.clone(), key: staging_key.clone(), version_id: None, }, completed_upload_id: Some("upload-identity".into()), }; for _ in 0..2 { assert_eq!( client .post(format!("{attempt_url}/staging/object")) .json(&staging_object) .send() .await .unwrap() .status(), 204 ); } let mut changed_staging = staging_object; changed_staging.object.key.push_str("-changed"); assert_eq!( client .post(format!("{attempt_url}/staging/object")) .json(&changed_staging) .send() .await .unwrap() .status(), 409 ); let resolved_id = unique_id("resolved"); let resolved = AttemptResolvedSourceRequestDto { operation_id: unique_id("operation"), fence: claimed.authority.fence, worker_incarnation: worker.clone(), resolved_source_ref: ResolvedSourceDescriptorRefDto { id: resolved_id.clone(), kind: ResolvedSourceDescriptorKindDto::ResolvedSource, schema_version: positive(1), byte_length: safe(64), item_count: Some(safe(1)), checksum: format!("sha256:{}", "22".repeat(32)), }, }; assert_eq!( client .post(format!("{attempt_url}/resolved-source")) .json(&resolved) .send() .await .unwrap() .status(), 204 ); let checkpoint = AttemptCheckpointRequestDto { operation_id: unique_id("operation"), fence: claimed.authority.fence, worker_incarnation: worker.clone(), checkpoint_ref: CheckpointDescriptorRefDto { id: unique_id("checkpoint"), kind: CheckpointDescriptorKindDto::Checkpoint, schema_version: positive(1), byte_length: safe(32), item_count: Some(safe(1)), checksum: format!("sha256:{}", "33".repeat(32)), }, resolved_source_id: resolved_id, validated_bytes: safe(1024), validated_members: safe(1), staging_multipart: None, }; assert_eq!( client .post(format!("{attempt_url}/checkpoint")) .json(&checkpoint) .send() .await .unwrap() .status(), 204 ); let progress = AttemptProgressRequestDto { operation_id: unique_id("operation"), fence: claimed.authority.fence, worker_incarnation: worker.clone(), progress: ProgressValueDto { completed_bytes: safe(2048), total_bytes: Some(safe(4096)), completed_members: safe(1), total_members: Some(safe(1)), }, }; assert_eq!( client .post(format!("{attempt_url}/progress")) .json(&progress) .send() .await .unwrap() .status(), 204 ); let replayed_claim: ClaimResponseDto = client .post(&claim_url) .json(&claim) .send() .await .unwrap() .json() .await .unwrap(); assert_eq!( serde_json::to_value(replayed_claim).unwrap(), serde_json::to_value(&claimed).unwrap() ); let failure = AttemptOutcomeRequestDto { operation_id: unique_id("operation"), fence: claimed.authority.fence, worker_incarnation: worker.clone(), outcome: AttemptOutcomeDto::Failure(AttemptFailureOutcomeDto { failure: AttemptRuntimeFailureDto { class: RuntimeFailureClass::TransientSource, message: "secret-value-must-not-escape".into(), safe_details: None, }, }), }; let outcome_url = format!("{attempt_url}/outcome"); let response = client .post(&outcome_url) .json(&failure) .send() .await .unwrap(); assert_eq!(response.status(), 200); let bytes = response.bytes().await.unwrap(); assert!(!String::from_utf8_lossy(&bytes).contains("secret-value-must-not-escape")); let projection: JobPhaseProjectionDto = serde_json::from_slice(&bytes).unwrap(); assert_eq!(projection.phase, JobPhase::RetryWait); let report = harness .service() .reconcile_once(&RecoveryConfig { retry_delay: Duration::ZERO, cleanup_claim_ttl: Duration::from_secs(1), batch_size: 128, }) .await .unwrap(); assert_eq!(report.retries_scheduled, 1); let replay = client .post(&outcome_url) .json(&failure) .send() .await .unwrap(); assert_eq!(replay.status(), 200); assert_eq!(replay.bytes().await.unwrap(), bytes); let cross_kind = client .post(format!("{attempt_url}/progress")) .json(&AttemptProgressRequestDto { operation_id: renewal_operation, fence: claimed.authority.fence, worker_incarnation: worker.clone(), progress: ProgressValueDto { completed_bytes: safe(2048), total_bytes: Some(safe(4096)), completed_members: safe(1), total_members: Some(safe(1)), }, }) .send() .await .unwrap(); assert_eq!(cross_kind.status(), 409); assert_eq!( cross_kind.headers().get("content-type").unwrap(), "application/problem+json" ); let error: ApiErrorDto = cross_kind.json().await.unwrap(); assert_eq!(error.message, "precondition failed"); let stop = AttemptStopRequestDto { operation_id: unique_id("operation"), fence: claimed.authority.fence, worker_incarnation: worker, }; for _ in 0..2 { assert_eq!( client .post(format!("{attempt_url}/stop")) .json(&stop) .send() .await .unwrap() .status(), 204 ); } server.stop(true).await; } #[tokio::test] async fn success_outcome_publishes_the_exact_manifest_and_replays() { let Some(harness) = Harness::new("success").await else { return; }; let state = WorkerApiState::new(harness.service()); let (base, server, client) = start_authenticated_server(state).await; let job_id = harness.job().await; let worker = unique_id("worker"); let claim: ClaimResponseDto = client .post(format!("{base}/internal/v1/jobs/{job_id}/claim")) .json(&ClaimRequestDto { operation_id: unique_id("operation"), worker_incarnation: worker.clone(), }) .send() .await .unwrap() .json() .await .unwrap(); let attempt_url = format!("{base}/internal/v1/attempts/{}", claim.authority.attempt_id); assert_eq!( client .post(format!("{attempt_url}/activate")) .json(&AttemptActivateRequestDto { operation_id: unique_id("operation"), fence: claim.authority.fence, worker_incarnation: worker.clone(), }) .send() .await .unwrap() .status(), 200 ); let result = harness.seal_result(&claim.authority.attempt_id).await; let operation = unique_id("operation"); let outcome = AttemptOutcomeRequestDto { operation_id: operation, fence: claim.authority.fence, worker_incarnation: worker, outcome: AttemptOutcomeDto::Success(remora_contracts::AttemptSuccessOutcomeDto { result }), }; let outcome_url = format!("{attempt_url}/outcome"); for _ in 0..2 { let response = client .post(&outcome_url) .json(&outcome) .send() .await .unwrap(); assert_eq!(response.status(), 200); assert_eq!( response .json::() .await .unwrap() .phase, JobPhase::Succeeded ); } server.stop(true).await; } #[tokio::test] async fn mutual_tls_requires_a_trusted_client_certificate() { let certificates = certificates(); let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let address = listener.local_addr().unwrap(); let server = HttpServer::new(|| { App::new().route( "/identity", web::get().to(|identity: AuthenticatedWorker| async move { HttpResponse::Ok().body(identity.identity.to_string()) }), ) }) .on_connect(authenticate_rustls_worker) .listen_rustls_0_23(listener, certificates.server_config) .unwrap() .run(); let handle = server.handle(); tokio::spawn(server); let root = ClientCertificate::from_der(certificates.ca.der()).unwrap(); let untrusted = Client::builder() .add_root_certificate(root.clone()) .build() .unwrap(); assert!( untrusted .get(format!("https://localhost:{}/identity", address.port())) .send() .await .is_err() ); let mut client_pem = certificates.client.pem(); client_pem.push_str(&certificates.client_key.serialize_pem()); let trusted = Client::builder() .add_root_certificate(root) .identity(Identity::from_pem(client_pem.as_bytes()).unwrap()) .build() .unwrap(); let response = trusted .get(format!("https://localhost:{}/identity", address.port())) .send() .await .unwrap(); assert_eq!(response.status(), 200); let expected = certificate_identity(certificates.client.der().as_ref()); assert_eq!(response.text().await.unwrap(), expected); handle.stop(true).await; } async fn start_authenticated_server( state: WorkerApiState, ) -> (String, actix_web::dev::ServerHandle, Client) { let certificates = certificates(); let client = trusted_client(&certificates); let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let address = listener.local_addr().unwrap(); let server = HttpServer::new(move || { App::new() .app_data(web::Data::new(state.clone())) .configure(configure_worker_api) }) .on_connect(authenticate_rustls_worker) .listen_rustls_0_23(listener, certificates.server_config) .unwrap() .run(); let handle = server.handle(); tokio::spawn(server); ( format!("https://localhost:{}", address.port()), handle, client, ) } fn trusted_client(certificates: &Certificates) -> Client { let root = ClientCertificate::from_der(certificates.ca.der()).unwrap(); let mut client_pem = certificates.client.pem(); client_pem.push_str(&certificates.client_key.serialize_pem()); Client::builder() .add_root_certificate(root) .identity(Identity::from_pem(client_pem.as_bytes()).unwrap()) .build() .unwrap() } fn certificate_identity(certificate: &[u8]) -> String { use sha2::{Digest, Sha256}; let mut value = String::from("w"); for byte in Sha256::digest(certificate) { use std::fmt::Write as _; write!(value, "{byte:02x}").unwrap(); } value } struct Certificates { ca: Certificate, client: Certificate, client_key: KeyPair, server_config: ServerConfig, } fn certificates() -> Certificates { let _ = rustls::crypto::ring::default_provider().install_default(); let mut ca_params = CertificateParams::new(Vec::new()).unwrap(); ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); ca_params.key_usages = vec![ KeyUsagePurpose::DigitalSignature, KeyUsagePurpose::KeyCertSign, KeyUsagePurpose::CrlSign, ]; let ca_key = KeyPair::generate().unwrap(); let ca = ca_params.self_signed(&ca_key).unwrap(); let issuer = Issuer::new(ca_params, ca_key); let server_key = KeyPair::generate().unwrap(); let mut server_params = CertificateParams::new(vec!["localhost".into()]).unwrap(); server_params.key_usages = vec![KeyUsagePurpose::DigitalSignature]; server_params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth]; let server = server_params.signed_by(&server_key, &issuer).unwrap(); let client_key = KeyPair::generate().unwrap(); let mut client_params = CertificateParams::new(Vec::new()).unwrap(); client_params.key_usages = vec![KeyUsagePurpose::DigitalSignature]; client_params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ClientAuth]; let client = client_params.signed_by(&client_key, &issuer).unwrap(); let mut roots = RootCertStore::empty(); roots.add(ca.der().clone()).unwrap(); let verifier = WebPkiClientVerifier::builder(Arc::new(roots)) .build() .unwrap(); let server_key_der = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(server_key.serialize_der())); let server_config = ServerConfig::builder() .with_client_cert_verifier(verifier) .with_single_cert(vec![server.der().clone()], server_key_der) .unwrap(); Certificates { ca, client, client_key, server_config, } } fn required_env(name: &str) -> Option { match std::env::var(name) { Ok(value) if !value.is_empty() => Some(value), _ if std::env::var_os("REMORA_WORKER_API_REQUIRE").is_some() => { panic!("{name} is required") } _ => None, } } fn coordination_namespace() -> String { std::env::var("REMORA_COORDINATION_NAMESPACE").unwrap_or_else(|_| "remora:test:".into()) } fn nonce() -> u128 { use std::sync::atomic::{AtomicU64, Ordering}; static COUNTER: AtomicU64 = AtomicU64::new(0); (u128::from(std::process::id()) << 64) | u128::from(COUNTER.fetch_add(1, Ordering::Relaxed)) } fn unique_id(prefix: &str) -> OpaqueId { OpaqueId::parse(format!("{prefix}{:032x}", nonce())).unwrap() } fn positive(value: u64) -> PositiveIntegerDto { PositiveIntegerDto::new(value).unwrap() } fn safe(value: u64) -> SafeIntegerDto { SafeIntegerDto::new(value).unwrap() }