saddle-framework 0.2.0

use std::{
    marker::PhantomData,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, Ordering},
    },
    time::Duration,
};

use saddle_admission::{EntryIoAuditPlan, ProcessLedger, StartupPublishToken};
use saddle_core::{ComponentLifecycle, ErrorKind, LifecycleFuture, SaddleError};
use saddle_runtime::{
    compiled_route::{
        OfficialCompiledDriverFinalizer, OfficialHttpRouteCoordinator, OfficialHttpTcpOutcome,
        OfficialTcpAttemptProfile, WaitingHttpTcp,
    },
    post_driver::MustSubmitDriverFinalizer,
    startup_assembly::VerifiedTransportRuntimeProfile,
};
use saddle_service::internal::RegisteredRouteExecutionProof;
use tokio::{
    io::AsyncWriteExt,
    net::{TcpListener, TcpStream},
    task::JoinHandle,
};

use super::{
    bootstrap::{
        GeneratedApplicationBootstrap, GeneratedBootstrapParts, GeneratedCompiledAdapter,
        GeneratedContextFactory, GeneratedRuntimeAdapter, ProductionServiceBundle,
    },
    framing::{Http1Framing, Http1ResponsePlan},
    transport::{
        Http1EntryError, Http1EntryOutcome, attempt_connection, cancel_waiting, retry_waiting,
    },
};

// Private candidate machine inputs; final Gate/configuration freeze may change
// the values without changing the generated bootstrap ownership contract.
const ACTIVE: usize = 1;
const CONNECTIONS: usize = 2;
const IO_BYTES: usize = 256;
const TICK: Duration = Duration::from_millis(1);

struct Running {
    stop: Arc<AtomicBool>,
    worker: JoinHandle<OfficialCompiledDriverFinalizer>,
}

pub(crate) struct ProductionHttp1<A, C, B>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    listener: Mutex<Option<TcpListener>>,
    assembly: Mutex<Option<Assembly<A, C>>>,
    publish: Mutex<Option<StartupPublishToken<B>>>,
    finalizer: Mutex<Option<MustSubmitDriverFinalizer>>,
    running: Mutex<Option<Running>>,
    types: PhantomData<fn() -> (A, C)>,
    _candidate: crate::startup_candidate::CommittedStartupCandidateOwner,
}

impl<A, C, B> ProductionHttp1<A, C, B>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    pub(crate) fn from_staged(
        listener: TcpListener,
        assembly: Assembly<A, C>,
        publish: StartupPublishToken<B>,
        finalizer: MustSubmitDriverFinalizer,
        candidate: crate::startup_candidate::CommittedStartupCandidateOwner,
    ) -> Self {
        Self {
            listener: Mutex::new(Some(listener)),
            assembly: Mutex::new(Some(assembly)),
            publish: Mutex::new(Some(publish)),
            finalizer: Mutex::new(Some(finalizer)),
            running: Mutex::new(None),
            types: PhantomData,
            _candidate: candidate,
        }
    }
}

impl<A, C, B> ComponentLifecycle for ProductionHttp1<A, C, B>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
    B: Send + 'static,
{
    fn name(&self) -> &'static str {
        "http1-production"
    }

    fn start(&self) -> LifecycleFuture<'_> {
        Box::pin(async move {
            if self.running.lock().unwrap().is_some() {
                return Err(component_error("http1.already_started"));
            }
            let publish = self
                .publish
                .lock()
                .unwrap()
                .take()
                .ok_or_else(|| component_error("http1.publish_token_consumed"))?;
            drop(publish.into_continuation());
            let (coordinator, adapter, head_deadline, attempt) = self
                .assembly
                .lock()
                .unwrap()
                .take()
                .ok_or_else(|| component_error("http1.staged_server_consumed"))?;
            let listener = self
                .listener
                .lock()
                .unwrap()
                .take()
                .ok_or_else(|| component_error("http1.listener_consumed"))?;
            let stop = Arc::new(AtomicBool::new(false));
            let worker = tokio::spawn(run_driver(
                listener,
                coordinator,
                adapter,
                head_deadline,
                attempt,
                Arc::clone(&stop),
            ));
            *self.running.lock().unwrap() = Some(Running { stop, worker });
            Ok(())
        })
    }

    fn shutdown(&self) -> LifecycleFuture<'_> {
        Box::pin(async move {
            let running = self
                .running
                .lock()
                .unwrap()
                .take()
                .ok_or_else(|| component_error("http1.not_started"))?;
            running.stop.store(true, Ordering::Release);
            let finalizer = running
                .worker
                .await
                .map_err(|_| component_error("http1.driver_failed"))?;
            self.finalizer
                .lock()
                .unwrap()
                .take()
                .ok_or_else(|| component_error("http1.finalizer_consumed"))?
                .submit(finalizer);
            Ok(())
        })
    }
}

type RuntimeAdapter<A, C> = GeneratedRuntimeAdapter<A, C>;
type Coordinator<A, C> = OfficialHttpRouteCoordinator<RuntimeAdapter<A, C>, Http1Framing>;
type Assembly<A, C> = (
    Coordinator<A, C>,
    Arc<RuntimeAdapter<A, C>>,
    Duration,
    OfficialTcpAttemptProfile,
);

/// Fallible, resource-free transport preflight. It owns the checked generated
/// graph but no Runtime, ledger/domain, DB pool, socket, task, or listener.
pub(crate) struct StagedHttp1<A, C>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    adapter: GeneratedRuntimeAdapter<A, C>,
    expected_generation: u64,
    expected_build_identity: [u8; 32],
    audit: EntryIoAuditPlan,
}

impl<A, C> StagedHttp1<A, C>
where
    A: ProductionServiceBundle,
    C: GeneratedContextFactory<A::Context>,
{
    pub(crate) fn service_production_fact_input(
        &self,
    ) -> Result<
        saddle_service::internal::ServiceProductionFactInput,
        saddle_service::internal::ServiceProductionFactError,
    > {
        self.adapter.production_fact_input()
    }
}

pub(crate) struct PreparedHttp1<A, C>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    adapter: GeneratedRuntimeAdapter<A, C>,
    head_deadline: Duration,
    attempt_deadline: Duration,
    finalization_termination_bound: Duration,
    task_storage_bound: usize,
    audit: EntryIoAuditPlan,
}

pub(crate) struct BoundPreparedHttp1<A, C>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    listener: TcpListener,
    prepared: PreparedHttp1<A, C>,
}

impl<A, C> PreparedHttp1<A, C>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    pub(crate) fn bind_owned(self, listener: TcpListener) -> BoundPreparedHttp1<A, C> {
        BoundPreparedHttp1 {
            listener,
            prepared: self,
        }
    }
}

impl<A, C> BoundPreparedHttp1<A, C>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    pub(crate) fn into_parts(self) -> (TcpListener, PreparedHttp1<A, C>) {
        (self.listener, self.prepared)
    }
}

pub(crate) struct GeneratedTransportBinding {
    pub generation: u64,
    pub build_identity: [u8; 32],
    pub route_set_identity: [u8; 32],
    pub static_layout_identity: [u8; 32],
}

impl<A, C> StagedHttp1<A, C>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    pub(crate) fn compose_generated<T>(
        self,
        runtime: saddle_runtime::capacity_leaf::VerifiedRuntimeBuildCapacityLeaf,
        calibration: saddle_runtime::capacity_leaf::VerifiedRuntimeBuildCalibrationOwner,
        termination: T,
    ) -> Result<
        (Self, saddle_admission::ComposedGeneratedStartupFactsOwner),
        saddle_admission::GeneratedStartupComposerError,
    >
    where
        A: ProductionServiceBundle,
        T: saddle_admission::GeneratedTerminationTopologyWorkOwner,
    {
        let (adapter, generated) =
            self.adapter
                .compose_generated(runtime, calibration, termination)?;
        Ok((Self { adapter, ..self }, generated))
    }

    /// Fallibly pairs the generated graph with Runtime's one-shot profile.
    /// This runs during transactional prepare, before commit or listener bind.
    pub(crate) fn prepare(
        self,
        profile: VerifiedTransportRuntimeProfile,
    ) -> Result<PreparedHttp1<A, C>, SaddleError> {
        let (plan, generation, build, paired_routes) = profile.identities();
        if !runtime_profile_identity_matches(
            plan,
            generation,
            build,
            paired_routes,
            self.expected_generation,
            self.expected_build_identity,
        ) {
            return Err(component_error("http1.runtime_profile_foreign"));
        }
        let (head, attempt, finalization, task_storage) = profile.into_transport_parts();
        if head.is_zero() || attempt.is_zero() || finalization.is_zero() || task_storage == 0 {
            return Err(component_error("http1.runtime_profile_insufficient"));
        }
        Ok(PreparedHttp1 {
            adapter: self.adapter,
            head_deadline: head,
            attempt_deadline: attempt,
            finalization_termination_bound: finalization,
            task_storage_bound: task_storage,
            audit: self.audit,
        })
    }
}

fn runtime_profile_identity_matches(
    plan: [u8; 32],
    generation: u64,
    build: [u8; 32],
    paired_routes: [u8; 32],
    expected_generation: u64,
    expected_build: [u8; 32],
) -> bool {
    // `paired_routes` is Runtime's domain-separated attestation of the paired
    // plan. `stage()` already checked the generated raw route-set, static
    // layout and every route proof; comparing the two domains is invalid.
    ![plan, build, paired_routes].contains(&[0; 32])
        && generation == expected_generation
        && build == expected_build
}

impl<A, C> PreparedHttp1<A, C>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    /// Infallibly attaches the exact owners after every profile check passed.
    pub(crate) fn attach(
        self,
        database: Option<saddle_db::Database>,
        ledger: ProcessLedger,
        domain: saddle_admission::OfficialTokioDomain,
        allocation: saddle_admission::ProcessAllocationProfile,
        db: Option<saddle_admission::DbPermitDomain>,
    ) -> Assembly<A, C> {
        let adapter = Arc::new(self.adapter.attach_database(database));
        let coordinator = OfficialHttpRouteCoordinator::new(
            Arc::clone(&adapter),
            Arc::new(Http1Framing),
            ledger,
            ACTIVE,
            domain,
            allocation,
            db,
        );
        (
            coordinator,
            adapter,
            self.head_deadline,
            OfficialTcpAttemptProfile {
                task_storage_bound: self.task_storage_bound,
                read_audit_plan: self.audit,
                write_audit_plan: self.audit,
                expected_body_bytes: 0,
                deadline: self.attempt_deadline,
                termination_bound: self.finalization_termination_bound,
            },
        )
    }
}

pub(crate) fn stage<A, C, const ROUTES: usize>(
    bootstrap: GeneratedApplicationBootstrap<A, C, ROUTES>,
    binding: GeneratedTransportBinding,
) -> Result<StagedHttp1<A, C>, SaddleError>
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    let GeneratedBootstrapParts {
        adapter,
        routes,
        proofs,
        static_facts,
    } = bootstrap.into_parts(None);
    if static_facts.route_count != ROUTES
        || static_facts.route_set_identity != binding.route_set_identity
        || static_facts.static_layout_identity != binding.static_layout_identity
        || binding.generation == 0
        || [
            binding.build_identity,
            binding.route_set_identity,
            binding.static_layout_identity,
        ]
        .contains(&[0; 32])
    {
        return Err(component_error("http1.generated_static_profile_foreign"));
    }
    for (route, proof) in routes.iter().zip(proofs) {
        let derived = adapter
            .lookup_generated(route.route(), route.declared_length())
            .map_err(|_| component_error("http1.route_proof_invalid"))?;
        if derived.identity() != proof.identity() {
            return Err(component_error("http1.route_proof_foreign"));
        }
    }
    let audit = EntryIoAuditPlan::locked_linux_x86_64_tokio_1_53_1(IO_BYTES, IO_BYTES, &[])
        .map_err(|_| component_error("http1.audit_profile_failed"))?;
    Ok(StagedHttp1 {
        adapter,
        expected_generation: binding.generation,
        expected_build_identity: binding.build_identity,
        audit,
    })
}

async fn run_driver<A, C>(
    listener: TcpListener,
    mut coordinator: Coordinator<A, C>,
    adapter: Arc<RuntimeAdapter<A, C>>,
    head_deadline: Duration,
    attempt: OfficialTcpAttemptProfile,
    stop: Arc<AtomicBool>,
) -> OfficialCompiledDriverFinalizer
where
    A: GeneratedCompiledAdapter,
    C: GeneratedContextFactory<A::Context>,
{
    let mut active = 0_usize;
    let mut waiting: Option<WaitingHttpTcp<RegisteredRouteExecutionProof, (), Http1ResponsePlan>> =
        None;
    loop {
        if coordinator.reap_finished().await {
            active = active.saturating_sub(1);
        }
        if stop.load(Ordering::Acquire) {
            coordinator.stop_accepting();
            if let Some(owner) = waiting.take() {
                let _ = cancel_waiting(&coordinator, owner);
            }
            return coordinator.shutdown(false).await;
        }
        if active == 0 {
            if let Some(owner) = waiting.take() {
                if let Some(outcome) = retry_waiting(&mut coordinator, owner, attempt) {
                    apply_outcome(outcome, &mut active, &mut waiting, attempt.deadline).await;
                }
                continue;
            }
        }
        if active + usize::from(waiting.is_some()) >= CONNECTIONS {
            tokio::time::sleep(TICK).await;
            continue;
        }
        tokio::select! {
            accepted = listener.accept() => {
                let Ok((socket, _)) = accepted else {
                    tokio::time::sleep(TICK).await;
                    continue;
                };
                match attempt_connection(
                    &mut coordinator,
                    adapter.as_ref(),
                    (),
                    socket,
                    head_deadline,
                    attempt,
                ).await {
                    Http1EntryOutcome::Runtime(outcome) => {
                        apply_outcome(outcome, &mut active, &mut waiting, attempt.deadline).await;
                    }
                    Http1EntryOutcome::Rejected(error, socket) => {
                        let response = match error {
                            Http1EntryError::Parse(_) => BAD_REQUEST,
                            Http1EntryError::UnknownRoute => NOT_FOUND,
                            Http1EntryError::HeadDeadline | Http1EntryError::Disconnect => b"",
                        };
                        write_pre_admission(socket, response, attempt.deadline).await;
                    }
                }
            }
            () = tokio::time::sleep(TICK) => {}
        }
    }
}

async fn apply_outcome<C>(
    outcome: OfficialHttpTcpOutcome<RegisteredRouteExecutionProof, C, Http1ResponsePlan>,
    active: &mut usize,
    waiting: &mut Option<WaitingHttpTcp<RegisteredRouteExecutionProof, C, Http1ResponsePlan>>,
    deadline: Duration,
) where
    C: Copy,
{
    match outcome {
        OfficialHttpTcpOutcome::Accepted(_) => *active += 1,
        OfficialHttpTcpOutcome::Registered(owner) => *waiting = Some(owner),
        OfficialHttpTcpOutcome::Reject(_, socket, _, _, _) => {
            write_pre_admission(socket, UNAVAILABLE, deadline).await;
        }
        OfficialHttpTcpOutcome::Stop(_, socket, _, _, _) => {
            if let Some(socket) = socket {
                write_pre_admission(socket, UNAVAILABLE, deadline).await;
            }
        }
    }
}

const BAD_REQUEST: &[u8] =
    b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
const NOT_FOUND: &[u8] =
    b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
const UNAVAILABLE: &[u8] =
    b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";

async fn write_pre_admission(mut socket: TcpStream, response: &[u8], deadline: Duration) {
    let _ = tokio::time::timeout(deadline, async {
        if !response.is_empty() {
            socket.write_all(response).await?;
        }
        socket.shutdown().await
    })
    .await;
}

fn component_error(code: &'static str) -> SaddleError {
    SaddleError::new(
        ErrorKind::Infrastructure,
        code,
        "production HTTP/1.1 component failed",
    )
}

#[cfg(test)]
mod tests {
    use super::runtime_profile_identity_matches;

    #[test]
    fn paired_route_attestation_is_not_compared_to_raw_route_identity() {
        assert!(runtime_profile_identity_matches(
            [1; 32], 7, [2; 32], [3; 32], 7, [2; 32],
        ));
    }

    #[test]
    fn foreign_or_missing_runtime_profile_identity_is_rejected() {
        assert!(!runtime_profile_identity_matches(
            [1; 32], 8, [2; 32], [3; 32], 7, [2; 32],
        ));
        assert!(!runtime_profile_identity_matches(
            [1; 32], 7, [4; 32], [3; 32], 7, [2; 32],
        ));
        assert!(!runtime_profile_identity_matches(
            [1; 32], 7, [2; 32], [0; 32], 7, [2; 32],
        ));
    }
}