use std::{ any::Any, future::{Ready, ready}, }; use actix_web::{ FromRequest, HttpRequest, HttpResponse, ResponseError, dev::Payload, http::StatusCode, web, }; use remora_attempt::{ ActivationMutationOutcome, AttemptAuthority, AttemptError, AttemptMutationOutcome, AttemptService, CheckpointProjection, CheckpointRecord, ClaimOutcome, DurableDescriptor, LeaseProjection, MutationRejection, OutcomePublication, ProgressRecord, RenewalMutationOutcome, }; use remora_contracts::{ ActivationResponseDto, ApiErrorCodeDto, ApiErrorDto, AttemptActivateRequestDto, AttemptAuthorityDto, AttemptCheckpointRequestDto, AttemptOutcomeDto, AttemptOutcomeRequestDto, AttemptProgressRequestDto, AttemptRenewRequestDto, AttemptResolvedSourceRequestDto, CheckpointDescriptorKindDto, CheckpointDescriptorRefDto, ClaimRequestDto, ClaimResponseDto, CredentialGrantVersionDto, DateTimeDto, DurableDescriptorKindDto, DurableDescriptorRefDto, ExecutionEnvelopeDto, JobPhaseProjectionDto, LeaseProjectionDto, NullableDto, PositiveIntegerDto, ResolvedSourceDescriptorKindDto, ResolvedSourceDescriptorRefDto, SafeIntegerDto, SourceProjectionDto, StagingTargetDto, TargetDto, }; use remora_domain::OpaqueId; use remora_result::PublicationInput; use remora_target::{ContentSha256, ObjectIdentity, RESERVED_NAMESPACE_SEGMENT, StoredObject}; use sha2::{Digest, Sha256}; use time::{OffsetDateTime, format_description::well_known::Rfc3339}; #[derive(Clone, Debug, Eq, PartialEq)] pub struct AuthenticatedWorker { pub identity: OpaqueId, } pub fn authenticate_rustls_worker( connection: &dyn Any, extensions: &mut actix_web::dev::Extensions, ) { type Stream = actix_tls::accept::rustls_0_23::TlsStream; let Some(stream) = connection.downcast_ref::() else { return; }; let Some(certificate) = stream .get_ref() .1 .peer_certificates() .and_then(|chain| chain.first()) else { return; }; let digest = Sha256::digest(certificate.as_ref()); let mut value = String::from("w"); for byte in digest { use std::fmt::Write as _; write!(value, "{byte:02x}").expect("writing to String cannot fail"); } extensions.insert(AuthenticatedWorker { identity: OpaqueId::parse(value).expect("certificate digest is a valid identity"), }); } impl FromRequest for AuthenticatedWorker { type Error = WorkerApiError; type Future = Ready>; fn from_request(request: &HttpRequest, _payload: &mut Payload) -> Self::Future { ready( request .conn_data::() .cloned() .ok_or_else(WorkerApiError::unauthenticated), ) } } #[derive(Clone)] pub struct WorkerApiState { attempts: AttemptService, } impl WorkerApiState { #[must_use] pub fn new(attempts: AttemptService) -> Self { Self { attempts } } } pub fn configure_worker_api(config: &mut web::ServiceConfig) { config.service( web::scope("/internal/v1") .app_data( web::JsonConfig::default() .limit(1024 * 1024) .error_handler(|_, _| WorkerApiError::invalid_request().into()), ) .route("/jobs/{job_id}/claim", web::post().to(claim)) .route("/attempts/{attempt_id}/renew", web::post().to(renew)) .route("/attempts/{attempt_id}/activate", web::post().to(activate)) .route( "/attempts/{attempt_id}/resolved-source", web::post().to(pin_resolved_source), ) .route( "/attempts/{attempt_id}/checkpoint", web::post().to(persist_checkpoint), ) .route( "/attempts/{attempt_id}/progress", web::post().to(report_progress), ) .route( "/attempts/{attempt_id}/outcome", web::post().to(report_outcome), ), ); } async fn claim( _: AuthenticatedWorker, state: web::Data, path: web::Path, request: web::Json, ) -> Result { let job_id = parse_id(path.into_inner())?; let outcome = state .attempts .claim(&job_id, &request.worker_incarnation, &request.operation_id) .await .map_err(WorkerApiError::from_attempt)?; let ClaimOutcome::Claimed { authority, .. } = outcome else { return Err(WorkerApiError::precondition_failed()); }; let execution = load_execution_envelope(&state.attempts, &authority).await?; Ok(HttpResponse::Ok().json(ClaimResponseDto { authority: authority_dto(&authority)?, execution, })) } async fn renew( _: AuthenticatedWorker, state: web::Data, path: web::Path, request: web::Json, ) -> Result { let authority = request_authority( &state, path.into_inner(), request.fence.get(), &request.worker_incarnation, ) .await?; match state .attempts .renew_operation(&authority, &request.operation_id) .await .map_err(WorkerApiError::from_attempt)? { RenewalMutationOutcome::Renewed { lease, .. } => { Ok(HttpResponse::Ok().json(lease_dto(lease)?)) } RenewalMutationOutcome::AuthorityLost { .. } => Err(WorkerApiError::precondition_failed()), } } async fn activate( _: AuthenticatedWorker, state: web::Data, path: web::Path, request: web::Json, ) -> Result { let authority = request_authority( &state, path.into_inner(), request.fence.get(), &request.worker_incarnation, ) .await?; let lease = match state .attempts .activate_operation(&authority, &request.operation_id) .await .map_err(WorkerApiError::from_attempt)? { ActivationMutationOutcome::Activated { lease, .. } => lease, ActivationMutationOutcome::PreconditionFailed => { return Err(WorkerApiError::precondition_failed()); } }; let execution = load_execution_envelope(&state.attempts, &authority).await?; Ok(HttpResponse::Ok().json(ActivationResponseDto { authority: authority_dto(&authority)?, execution, lease: lease_dto(lease)?, })) } async fn pin_resolved_source( _: AuthenticatedWorker, state: web::Data, path: web::Path, request: web::Json, ) -> Result { let authority = request_authority( &state, path.into_inner(), request.fence.get(), &request.worker_incarnation, ) .await?; accepted( state .attempts .pin_resolved_source( &authority, &request.operation_id, descriptor_from_resolved(&request.resolved_source_ref)?, ) .await .map_err(WorkerApiError::from_attempt)?, ) } async fn persist_checkpoint( _: AuthenticatedWorker, state: web::Data, path: web::Path, request: web::Json, ) -> Result { let authority = request_authority( &state, path.into_inner(), request.fence.get(), &request.worker_incarnation, ) .await?; accepted( state .attempts .persist_checkpoint( &authority, &request.operation_id, CheckpointRecord { descriptor: descriptor_from_checkpoint(&request.checkpoint_ref)?, resolved_source_id: request.resolved_source_id.clone(), validated_bytes: request.validated_bytes.get(), validated_members: request.validated_members.get(), }, ) .await .map_err(WorkerApiError::from_attempt)?, ) } async fn report_progress( _: AuthenticatedWorker, state: web::Data, path: web::Path, request: web::Json, ) -> Result { let authority = request_authority( &state, path.into_inner(), request.fence.get(), &request.worker_incarnation, ) .await?; accepted( state .attempts .report_progress( &authority, &request.operation_id, ProgressRecord { completed_bytes: request.progress.completed_bytes.get(), total_bytes: request.progress.total_bytes.map(SafeIntegerDto::get), completed_members: request.progress.completed_members.get(), total_members: request.progress.total_members.map(SafeIntegerDto::get), }, ) .await .map_err(WorkerApiError::from_attempt)?, ) } async fn report_outcome( _: AuthenticatedWorker, state: web::Data, path: web::Path, request: web::Json, ) -> Result { let authority = request_authority( &state, path.into_inner(), request.fence.get(), &request.worker_incarnation, ) .await?; let outcome = match &request.outcome { AttemptOutcomeDto::Success(success) => { let result = &success.result; state .attempts .publish_success_operation( &authority, &request.operation_id, PublicationInput { attempt_id: authority.attempt_id.clone(), manifest_id: result.manifest_id.clone(), schema_version: u32::try_from(result.schema_version.get()) .map_err(|_| WorkerApiError::invalid_request())?, manifest: StoredObject { identity: ObjectIdentity { target_store_id: result.manifest.store.clone(), bucket: result.manifest.bucket.clone(), key: result.manifest.key.clone(), version_id: result.manifest.version_id.clone(), }, content_sha256: ContentSha256::parse(&result.manifest.checksum) .map_err(|_| WorkerApiError::invalid_request())?, size: result.manifest.size.get(), }, member_count: result.member_count.get(), total_bytes: result.total_bytes.get(), }, ) .await } AttemptOutcomeDto::Failure(failure) => { state .attempts .publish_failure_operation(&authority, &request.operation_id, failure.failure.class) .await } } .map_err(WorkerApiError::from_attempt)?; let phase = match outcome { OutcomePublication::Accepted { phase, .. } => phase, OutcomePublication::Rejected { reason, .. } => { return Err(WorkerApiError::from_rejection(reason)); } }; Ok(HttpResponse::Ok().json(JobPhaseProjectionDto { phase })) } fn accepted(outcome: AttemptMutationOutcome) -> Result { match outcome { AttemptMutationOutcome::Accepted(_) => Ok(HttpResponse::NoContent().finish()), AttemptMutationOutcome::Rejected { reason, .. } => { Err(WorkerApiError::from_rejection(reason)) } } } async fn request_authority( state: &WorkerApiState, attempt_id: String, fence: u64, worker: &OpaqueId, ) -> Result { let attempt_id = parse_id(attempt_id)?; state .attempts .historical_authority(&attempt_id) .await .map_err(WorkerApiError::from_attempt)? .filter(|authority| authority.fence == fence && authority.worker_incarnation == *worker) .ok_or_else(WorkerApiError::precondition_failed) } async fn load_execution_envelope( attempts: &AttemptService, authority: &AttemptAuthority, ) -> Result { let row = attempts .execution_envelope(&authority.attempt_id) .await .map_err(WorkerApiError::from_attempt)? .ok_or_else(WorkerApiError::precondition_failed)?; if row.job_id != authority.job_id || row.attempt_id != authority.attempt_id || row.fence != authority.fence || row.worker_incarnation != authority.worker_incarnation { return Err(WorkerApiError::precondition_failed()); } let source: SourceProjectionDto = serde_json::from_value(row.source).map_err(|_| WorkerApiError::unavailable())?; let target_prefix = row.target_prefix.join("/"); let resolved_source_ref = row.resolved_source.map(|descriptor| { NullableDto(Some(descriptor_ref( descriptor, DurableDescriptorKindDto::ResolvedSource, ))) }); let checkpoint_ref = row.checkpoint.and_then(|checkpoint| match checkpoint { CheckpointProjection::Reusable { checkpoint } => Some(NullableDto(Some(descriptor_ref( checkpoint.descriptor, DurableDescriptorKindDto::Checkpoint, )))), CheckpointProjection::Complete { .. } => None, }); Ok(ExecutionEnvelopeDto { id: row.id, job_id: authority.job_id.clone(), attempt_id: authority.attempt_id.clone(), fence: positive(authority.fence)?, worker_incarnation: authority.worker_incarnation.clone(), source, sensitive_source_ref: row.sensitive_source_ref, credential_version: NullableDto(match (row.credential_id, row.credential_version_id) { (Some(credential_id), Some(version_id)) => Some(CredentialGrantVersionDto { credential_id, version_id, }), (None, None) => None, _ => return Err(WorkerApiError::unavailable()), }), resolved_source_ref: resolved_source_ref.unwrap_or(NullableDto(None)), target: TargetDto { store: row.target_store.clone(), bucket: row.target_bucket.clone(), prefix: target_prefix, }, staging_target: StagingTargetDto { store: row.target_store, bucket: row.target_bucket, prefix: format!( "{RESERVED_NAMESPACE_SEGMENT}/a/{}/", authority.attempt_id.as_str() ), }, execution_grant_ref: row.execution_grant_id, checkpoint_ref: checkpoint_ref.unwrap_or(NullableDto(None)), }) } fn descriptor_ref( descriptor: DurableDescriptor, kind: DurableDescriptorKindDto, ) -> DurableDescriptorRefDto { DurableDescriptorRefDto { id: descriptor.id, kind, schema_version: positive(descriptor.schema_version.into()).expect("validated descriptor"), byte_length: safe(descriptor.byte_length).expect("validated descriptor"), item_count: descriptor .item_count .map(|value| safe(value).expect("validated descriptor")), checksum: descriptor.checksum, } } fn descriptor_from_resolved( value: &ResolvedSourceDescriptorRefDto, ) -> Result { let ResolvedSourceDescriptorKindDto::ResolvedSource = value.kind; Ok(DurableDescriptor { id: value.id.clone(), schema_version: u32::try_from(value.schema_version.get()) .map_err(|_| WorkerApiError::invalid_request())?, byte_length: value.byte_length.get(), item_count: value.item_count.map(SafeIntegerDto::get), checksum: value.checksum.clone(), }) } fn descriptor_from_checkpoint( value: &CheckpointDescriptorRefDto, ) -> Result { let CheckpointDescriptorKindDto::Checkpoint = value.kind; Ok(DurableDescriptor { id: value.id.clone(), schema_version: u32::try_from(value.schema_version.get()) .map_err(|_| WorkerApiError::invalid_request())?, byte_length: value.byte_length.get(), item_count: value.item_count.map(SafeIntegerDto::get), checksum: value.checksum.clone(), }) } fn authority_dto(authority: &AttemptAuthority) -> Result { Ok(AttemptAuthorityDto { attempt_id: authority.attempt_id.clone(), fence: positive(authority.fence)?, worker_incarnation: authority.worker_incarnation.clone(), }) } fn lease_dto(lease: LeaseProjection) -> Result { let observed = timestamp(lease.observed_at_ms)?; let expires = timestamp(lease.expires_at_ms)?; let renew_at = observed + (expires - observed) / 2; Ok(LeaseProjectionDto { expires_at: datetime(expires)?, renew_after: datetime(renew_at)?, }) } fn timestamp(milliseconds: u64) -> Result { let nanoseconds = i128::from(milliseconds) * 1_000_000; OffsetDateTime::from_unix_timestamp_nanos(nanoseconds) .map_err(|_| WorkerApiError::unavailable()) } fn datetime(value: OffsetDateTime) -> Result { let value = value .format(&Rfc3339) .map_err(|_| WorkerApiError::unavailable())?; DateTimeDto::parse(value).map_err(|_| WorkerApiError::unavailable()) } fn positive(value: u64) -> Result { PositiveIntegerDto::new(value).map_err(|_| WorkerApiError::unavailable()) } fn safe(value: u64) -> Result { SafeIntegerDto::new(value).map_err(|_| WorkerApiError::unavailable()) } fn parse_id(value: impl Into) -> Result { OpaqueId::parse(value).map_err(|_| WorkerApiError::invalid_request()) } #[derive(Debug, thiserror::Error)] #[error("worker API request failed")] pub struct WorkerApiError { status: StatusCode, code: ApiErrorCodeDto, message: &'static str, } impl WorkerApiError { fn invalid_request() -> Self { Self::new( StatusCode::BAD_REQUEST, ApiErrorCodeDto::InvalidRequest, "invalid request", ) } fn unauthenticated() -> Self { Self::new( StatusCode::UNAUTHORIZED, ApiErrorCodeDto::Unauthenticated, "authentication required", ) } fn precondition_failed() -> Self { Self::new( StatusCode::CONFLICT, ApiErrorCodeDto::PreconditionFailed, "precondition failed", ) } fn unavailable() -> Self { Self::new( StatusCode::SERVICE_UNAVAILABLE, ApiErrorCodeDto::Unavailable, "service unavailable", ) } fn new(status: StatusCode, code: ApiErrorCodeDto, message: &'static str) -> Self { Self { status, code, message, } } fn from_attempt(error: AttemptError) -> Self { match error { AttemptError::Invalid(_) => Self::invalid_request(), AttemptError::OperationConflict => Self::precondition_failed(), AttemptError::Coordination(_) | AttemptError::Durable(_) => Self::unavailable(), } } fn from_rejection(_reason: MutationRejection) -> Self { Self::precondition_failed() } } impl ResponseError for WorkerApiError { fn status_code(&self) -> StatusCode { self.status } fn error_response(&self) -> HttpResponse { HttpResponse::build(self.status) .content_type("application/problem+json") .json(ApiErrorDto { code: self.code, message: self.message.into(), request_id: request_id(), fields: None, }) } } fn request_id() -> OpaqueId { let mut bytes = [0_u8; 16]; getrandom::fill(&mut bytes).expect("operating system randomness is required"); let mut value = String::from("r"); for byte in bytes { use std::fmt::Write as _; write!(value, "{byte:02x}").expect("writing to String cannot fail"); } OpaqueId::parse(value).expect("generated request identifier is valid") }