Add authenticated internal worker API

crates/remora-api/src/worker.rs

@@ -0,0 +1,613 @@
+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<actix_web::rt::net::TcpStream>;
+    let Some(stream) = connection.downcast_ref::<Stream>() 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<Result<Self, Self::Error>>;
+
+    fn from_request(request: &HttpRequest, _payload: &mut Payload) -> Self::Future {
+        ready(
+            request
+                .conn_data::<Self>()
+                .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<WorkerApiState>,
+    path: web::Path<String>,
+    request: web::Json<ClaimRequestDto>,
+) -> Result<HttpResponse, WorkerApiError> {
+    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<WorkerApiState>,
+    path: web::Path<String>,
+    request: web::Json<AttemptRenewRequestDto>,
+) -> Result<HttpResponse, WorkerApiError> {
+    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<WorkerApiState>,
+    path: web::Path<String>,
+    request: web::Json<AttemptActivateRequestDto>,
+) -> Result<HttpResponse, WorkerApiError> {
+    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<WorkerApiState>,
+    path: web::Path<String>,
+    request: web::Json<AttemptResolvedSourceRequestDto>,
+) -> Result<HttpResponse, WorkerApiError> {
+    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<WorkerApiState>,
+    path: web::Path<String>,
+    request: web::Json<AttemptCheckpointRequestDto>,
+) -> Result<HttpResponse, WorkerApiError> {
+    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<WorkerApiState>,
+    path: web::Path<String>,
+    request: web::Json<AttemptProgressRequestDto>,
+) -> Result<HttpResponse, WorkerApiError> {
+    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<WorkerApiState>,
+    path: web::Path<String>,
+    request: web::Json<AttemptOutcomeRequestDto>,
+) -> Result<HttpResponse, WorkerApiError> {
+    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<HttpResponse, WorkerApiError> {
+    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<AttemptAuthority, WorkerApiError> {
+    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<ExecutionEnvelopeDto, WorkerApiError> {
+    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<DurableDescriptor, WorkerApiError> {
+    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<DurableDescriptor, WorkerApiError> {
+    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<AttemptAuthorityDto, WorkerApiError> {
+    Ok(AttemptAuthorityDto {
+        attempt_id: authority.attempt_id.clone(),
+        fence: positive(authority.fence)?,
+        worker_incarnation: authority.worker_incarnation.clone(),
+    })
+}
+
+fn lease_dto(lease: LeaseProjection) -> Result<LeaseProjectionDto, WorkerApiError> {
+    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<OffsetDateTime, WorkerApiError> {
+    let nanoseconds = i128::from(milliseconds) * 1_000_000;
+    OffsetDateTime::from_unix_timestamp_nanos(nanoseconds)
+        .map_err(|_| WorkerApiError::unavailable())
+}
+
+fn datetime(value: OffsetDateTime) -> Result<DateTimeDto, WorkerApiError> {
+    let value = value
+        .format(&Rfc3339)
+        .map_err(|_| WorkerApiError::unavailable())?;
+    DateTimeDto::parse(value).map_err(|_| WorkerApiError::unavailable())
+}
+
+fn positive(value: u64) -> Result<PositiveIntegerDto, WorkerApiError> {
+    PositiveIntegerDto::new(value).map_err(|_| WorkerApiError::unavailable())
+}
+
+fn safe(value: u64) -> Result<SafeIntegerDto, WorkerApiError> {
+    SafeIntegerDto::new(value).map_err(|_| WorkerApiError::unavailable())
+}
+
+fn parse_id(value: impl Into<String>) -> Result<OpaqueId, WorkerApiError> {
+    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")
+}