spiffe 0.18.0

use super::builder::{ReconnectConfig, ResourceLimits};
use super::errors::{MetricsErrorKind, X509SourceError};
use super::limits::validate_context;
use super::metrics::MetricsRecorder;
use crate::prelude::{debug, info, warn};
use crate::workload_api::error::WorkloadApiError;
use crate::workload_api::supervisor_common::{
    self, ErrorKey, ErrorTracker, StreamPhase, MAX_CONSECUTIVE_SAME_ERROR,
};
use crate::workload_api::x509_context::X509Context;
use crate::workload_api::WorkloadApiClient;
use crate::x509_source::source::{ApplyUpdateResult, Inner};
use crate::x509_source::types::{ClientFactory, SvidPicker};
use crate::X509Svid;
use futures::StreamExt as _;
use std::sync::Arc;
use std::time::Duration;
use tokio_util::sync::CancellationToken;

/// Attempts to create a Workload API client.
///
/// Records metrics and logs errors. Does not modify backoff (caller manages backoff progression).
///
/// On error, records `ClientCreation` metric. The caller should call `record_reconnect`
/// for steady-state reconnections (not for initial sync).
pub(super) async fn try_create_client(
    make_client: &ClientFactory,
    backoff: Duration,
    error_tracker: &mut ErrorTracker,
    metrics: Option<&dyn MetricsRecorder>,
) -> Result<WorkloadApiClient, WorkloadApiError> {
    match (make_client)().await {
        Ok(c) => {
            // Only log recovery if there were significant consecutive failures (>= 3) of the same type
            if error_tracker.last_error_kind() == Some(ErrorKey::ClientCreation)
                && error_tracker.consecutive_count() >= 3
            {
                debug!(
                    "Client creation recovered after {} consecutive failures",
                    error_tracker.consecutive_count()
                );
            }
            // Only reset tracker if the last error was client creation (actual recovery).
            // Don't reset if tracking stream connection errors (e.g., NoIdentityIssued).
            if error_tracker.last_error_kind() == Some(ErrorKey::ClientCreation) {
                error_tracker.reset();
            }
            Ok(c)
        }
        Err(e) => {
            let error_kind = ErrorKey::ClientCreation;
            let should_warn = error_tracker.record_error(error_kind);

            if should_warn {
                warn!(
                    "Failed to create WorkloadApiClient; retrying: error={}, backoff_ms={}",
                    e,
                    backoff.as_millis()
                );
            } else {
                debug!(
                    "Failed to create WorkloadApiClient (repeated); retrying: error={}, backoff_ms={}, consecutive_failures={}",
                    e,
                    backoff.as_millis(),
                    error_tracker.consecutive_count()
                );
            }

            if let Some(m) = metrics {
                m.record_error(MetricsErrorKind::ClientCreation);
            }
            Err(e)
        }
    }
}

/// Attempts to connect to the X.509 context stream.
///
/// Records metrics and logs errors. Does not modify backoff (caller manages backoff progression).
///
/// On error, records `StreamConnect` metric. The caller should call `record_reconnect`
/// for steady-state reconnections (not for initial sync).
pub(super) async fn try_connect_stream(
    client: &WorkloadApiClient,
    backoff: Duration,
    error_tracker: &mut ErrorTracker,
    metrics: Option<&dyn MetricsRecorder>,
    phase: StreamPhase,
    supervisor_id: Option<u64>,
) -> Result<
    impl futures::Stream<Item = Result<X509Context, WorkloadApiError>> + Send + 'static + use<>,
    WorkloadApiError,
> {
    match client.stream_x509_contexts().await {
        Ok(s) => {
            let id_suffix = supervisor_id.map_or_else(String::new, |id| format!(", id={id}"));

            // Only log recovery if the last error was a stream connection failure
            if error_tracker.last_error_kind() == Some(ErrorKey::StreamConnect)
                && error_tracker.consecutive_count() > 0
            {
                info!(
                    "Stream connection recovered after {} consecutive failures (phase={:?}{})",
                    error_tracker.consecutive_count(),
                    phase,
                    id_suffix
                );
            }
            error_tracker.reset();
            info!(
                "Connected to Workload API X509 context stream (phase={:?}{})",
                phase, id_suffix
            );
            Ok(s)
        }
        Err(e) => {
            // Handle "no identity issued" as a distinct transient state
            if matches!(e, WorkloadApiError::NoIdentityIssued) {
                let error_kind = ErrorKey::NoIdentityIssued;
                let should_warn = error_tracker.record_error(error_kind);

                if should_warn {
                    warn!("No identity issued yet; waiting before retry");
                } else {
                    debug!(
                        "No identity issued yet (repeated); waiting before retry: consecutive_failures={}",
                        error_tracker.consecutive_count()
                    );
                }

                // Don't record this as a stream error metric (it's expected/transient)
                return Err(e);
            }

            let error_kind = ErrorKey::StreamConnect;
            let should_warn = error_tracker.record_error(error_kind);

            if should_warn {
                warn!(
                    "Failed to connect to Workload API stream; retrying: error={}, backoff_ms={}",
                    e,
                    backoff.as_millis()
                );
            } else {
                debug!(
                    "Failed to connect to Workload API stream (repeated); retrying: error={}, backoff_ms={}, consecutive_failures={}",
                    e,
                    backoff.as_millis(),
                    error_tracker.consecutive_count()
                );
            }

            if let Some(m) = metrics {
                m.record_error(MetricsErrorKind::StreamConnect);
            }
            Err(e)
        }
    }
}

pub(super) async fn initial_sync_with_retry(
    make_client: &ClientFactory,
    picker: Option<&dyn SvidPicker>,
    cancel: &CancellationToken,
    reconnect: ReconnectConfig,
    limits: ResourceLimits,
    metrics: Option<&dyn MetricsRecorder>,
) -> Result<(Arc<X509Context>, Arc<X509Svid>), X509SourceError> {
    let mut backoff = reconnect.min_backoff;
    let mut error_tracker = ErrorTracker::new(MAX_CONSECUTIVE_SAME_ERROR);

    loop {
        if cancel.is_cancelled() {
            return Err(X509SourceError::Closed);
        }

        match try_sync_once(
            make_client,
            picker,
            limits,
            metrics,
            backoff,
            &mut error_tracker,
        )
        .await
        {
            Ok(v) => return Ok(v),
            Err(e) => {
                // Record InitialSyncFailed as an umbrella metric for any failed attempt.
                // Specific metrics (ClientCreation, StreamConnect, StreamError, StreamEnded,
                // LimitMaxSvids, LimitMaxBundles, LimitMaxBundleDerBytes, NoSuitableSvid)
                // are already recorded in try_sync_once(). Detailed logs are also produced
                // by try_create_client/try_connect_stream/stream read, so we avoid duplicate
                // outer logs here to reduce noise.
                if let Some(m) = metrics {
                    m.record_error(MetricsErrorKind::InitialSyncFailed);
                }
                // INVALID_ARGUMENT means the Workload API rejected the request itself
                // (non-conforming server, a proxy mangling gRPC metadata, protocol
                // mismatch, etc.). Retrying with backoff would just repeat the same
                // rejection forever, hanging construction indefinitely unless the caller
                // opted into `initial_sync_timeout`. Fail fast instead. Steady-state
                // reconnects (after a successful initial sync) are intentionally left
                // retrying as before, since we already hold valid material at that point.
                if let X509SourceError::Source(inner) = &e {
                    if inner.is_invalid_argument() {
                        warn!(
                            "Initial sync: Workload API rejected the request as invalid \
                             (INVALID_ARGUMENT); not retrying: error={e}"
                        );
                        return Err(e);
                    }
                }
                if sleep_or_cancel(cancel, backoff).await {
                    return Err(X509SourceError::Closed);
                }
                // Choose backoff policy based on error type
                match &e {
                    X509SourceError::Source(WorkloadApiError::NoIdentityIssued) => {
                        // Use slower backoff for "no identity issued" (expected transient state)
                        backoff = next_backoff_for_no_identity(backoff, reconnect.max_backoff);
                        warn!(
                            "Initial sync: no identity issued, using backoff_ms={}",
                            backoff.as_millis()
                        );
                    }
                    _ => {
                        // Use standard exponential backoff for other errors
                        backoff = next_backoff(backoff, reconnect.max_backoff);
                    }
                }
            }
        }
    }
}

async fn try_sync_once(
    make_client: &ClientFactory,
    picker: Option<&dyn SvidPicker>,
    limits: ResourceLimits,
    metrics: Option<&dyn MetricsRecorder>,
    backoff: Duration,
    error_tracker: &mut ErrorTracker,
) -> Result<(Arc<X509Context>, Arc<X509Svid>), X509SourceError> {
    // Use shared client creation logic (records ClientCreation metric on failure).
    // Initial sync does not record reconnect metrics (it's not a reconnect).
    let client = match try_create_client(make_client, backoff, error_tracker, metrics).await {
        Ok(c) => c,
        Err(e) => {
            // Error already logged and metric recorded by try_create_client.
            return Err(X509SourceError::Source(e));
        }
    };

    // Use shared stream connection logic (records StreamConnect metric on failure).
    let mut stream = match try_connect_stream(
        &client,
        backoff,
        error_tracker,
        metrics,
        StreamPhase::InitialSync,
        None,
    )
    .await
    {
        Ok(s) => s,
        Err(e) => {
            // Error already logged and metric recorded by try_connect_stream.
            return Err(X509SourceError::Source(e));
        }
    };

    match stream.next().await {
        Some(Ok(ctx)) => {
            let selected_svid =
                validate_context(&ctx, picker, limits, metrics).inspect_err(|e| {
                    warn!("Initial X509 context rejected; will retry: error={e}");
                })?;

            Ok((Arc::new(ctx), selected_svid))
        }
        Some(Err(e)) => {
            // Record StreamError for stream read errors.
            warn!("Initial sync: Workload API stream error; will retry: error={e}",);
            if let Some(m) = metrics {
                m.record_error(MetricsErrorKind::StreamError);
            }
            Err(X509SourceError::Source(e))
        }
        None => {
            // Record StreamEnded for empty stream.
            warn!("Initial sync: Workload API stream ended immediately; will retry");
            if let Some(m) = metrics {
                m.record_error(MetricsErrorKind::StreamEnded);
            }
            Err(X509SourceError::StreamEnded)
        }
    }
}

pub(super) use supervisor_common::{next_backoff, next_backoff_for_no_identity, sleep_or_cancel};

/// Result of processing a stream of updates.
struct StreamResult {
    /// Whether the cancellation token was triggered.
    cancelled: bool,
    /// Whether at least one valid item was received from the stream.
    had_valid_item: bool,
}

impl Inner {
    pub(super) async fn run_update_supervisor(&self, cancellation_token: CancellationToken) {
        let supervisor_id = fastrand::u64(..);
        info!("Starting update supervisor: id={}", supervisor_id);

        let mut backoff = self.reconnect().min_backoff;
        let mut error_tracker = ErrorTracker::new(MAX_CONSECUTIVE_SAME_ERROR);

        loop {
            if cancellation_token.is_cancelled() {
                debug!("Cancellation signal received; stopping updates");
                return;
            }

            let Ok(client) = try_create_client(
                self.make_client(),
                backoff,
                &mut error_tracker,
                self.metrics(),
            )
            .await
            else {
                if self
                    .backoff_and_maybe_cancel(&cancellation_token, backoff)
                    .await
                {
                    return;
                }
                backoff = next_backoff(backoff, self.reconnect().max_backoff);
                continue;
            };

            match try_connect_stream(
                &client,
                backoff,
                &mut error_tracker,
                self.metrics(),
                StreamPhase::Supervisor,
                Some(supervisor_id),
            )
            .await
            {
                Ok(mut stream) => {
                    let result = self
                        .process_stream_updates(&mut stream, &cancellation_token, supervisor_id)
                        .await;
                    if result.cancelled {
                        return;
                    }

                    // Reset backoff only if the stream delivered valid data before failing.
                    if result.had_valid_item {
                        backoff = self.reconnect().min_backoff;
                    }

                    // Stream ended or errored. Sleep/backoff before retrying.
                    if self
                        .backoff_and_maybe_cancel(&cancellation_token, backoff)
                        .await
                    {
                        return;
                    }
                    if !result.had_valid_item {
                        backoff = next_backoff(backoff, self.reconnect().max_backoff);
                    }
                }
                Err(stream_err) => {
                    // Choose backoff policy based on error type
                    match stream_err {
                        WorkloadApiError::NoIdentityIssued => {
                            // Use slower backoff for "no identity issued" (expected transient state)
                            backoff =
                                next_backoff_for_no_identity(backoff, self.reconnect().max_backoff);
                            warn!(
                                "No identity issued: using backoff_ms={}",
                                backoff.as_millis()
                            );
                        }
                        _ => {
                            // Use standard exponential backoff for other errors
                            backoff = next_backoff(backoff, self.reconnect().max_backoff);
                        }
                    }

                    if self
                        .backoff_and_maybe_cancel(&cancellation_token, backoff)
                        .await
                    {
                        return;
                    }
                }
            }
        }
    }

    async fn backoff_and_maybe_cancel(&self, token: &CancellationToken, backoff: Duration) -> bool {
        // Single place where we record reconnect for steady-state reconnect loops.
        if let Some(metrics) = self.metrics() {
            metrics.record_reconnect();
        }
        sleep_or_cancel(token, backoff).await
    }

    /// Processes stream updates until the stream ends, errors, or cancellation is requested.
    async fn process_stream_updates(
        &self,
        stream: &mut (impl futures::Stream<Item = Result<X509Context, WorkloadApiError>>
                  + Unpin
                  + Send
                  + 'static),
        cancellation_token: &CancellationToken,
        supervisor_id: u64,
    ) -> StreamResult {
        let mut update_rejection_tracker = ErrorTracker::new(MAX_CONSECUTIVE_SAME_ERROR);
        let mut had_valid_item = false;

        loop {
            let item = tokio::select! {
                () = cancellation_token.cancelled() => {
                    debug!("Cancellation signal received; stopping update loop");
                    return StreamResult { cancelled: true, had_valid_item };
                }
                v = stream.next() => v,
            };

            match item {
                Some(Ok(ctx)) => {
                    match self.apply_update(Arc::new(ctx)) {
                        Ok(ApplyUpdateResult::Applied) => {
                            had_valid_item = true;
                            if update_rejection_tracker.consecutive_count() > 0 {
                                info!(
                                    "Update validation recovered after {} consecutive failures",
                                    update_rejection_tracker.consecutive_count(),
                                );
                                update_rejection_tracker.reset();
                            }
                            info!("X509 context updated");
                        }
                        Ok(ApplyUpdateResult::Unchanged) => {
                            had_valid_item = true;
                            if update_rejection_tracker.consecutive_count() > 0 {
                                debug!(
                                    "Update validation recovered with unchanged X509 context after {} consecutive failures",
                                    update_rejection_tracker.consecutive_count(),
                                );
                                update_rejection_tracker.reset();
                            }
                            debug!("X509 context unchanged; skipping update notification");
                        }
                        Err(e) => {
                            let should_warn =
                                update_rejection_tracker.record_error(ErrorKey::UpdateRejected);

                            if should_warn {
                                warn!("Rejected X509 context update: error={e}");
                            } else {
                                debug!(
                                    "Rejected X509 context update (repeated): error={}, consecutive_rejections={}",
                                    e,
                                    update_rejection_tracker.consecutive_count()
                                );
                            }
                            // Metrics already recorded by apply_update(); do not double-count.
                        }
                    }
                }
                Some(Err(e)) => {
                    warn!(
                        "Workload API stream error; reconnecting: id={}, error={}",
                        supervisor_id, e
                    );
                    self.record_error(MetricsErrorKind::StreamError);
                    return StreamResult {
                        cancelled: false,
                        had_valid_item,
                    };
                }
                None => {
                    warn!("Workload API stream ended; reconnecting: id={supervisor_id}");
                    self.record_error(MetricsErrorKind::StreamEnded);
                    return StreamResult {
                        cancelled: false,
                        had_valid_item,
                    };
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::transport::TransportError;
    use crate::workload_api::pb::workload::{X509svid, X509svidRequest, X509svidResponse};
    use crate::{TrustDomain, X509Bundle, X509BundleSet, X509Source};
    use std::sync::atomic::{AtomicUsize, Ordering};
    use tonic::codegen::http;

    fn invalid_argument_error() -> WorkloadApiError {
        WorkloadApiError::Transport(TransportError::Status(tonic::Status::invalid_argument(
            "bad request",
        )))
    }

    fn context_with_bundle(authority: &[u8]) -> Arc<X509Context> {
        let cert = include_bytes!("../../tests/testdata/svid/x509/1-svid-chain.der");
        let key = include_bytes!("../../tests/testdata/svid/x509/1-key.der");
        let svid = Arc::new(X509Svid::parse_from_der(cert, key).unwrap());
        let trust_domain = TrustDomain::new("example.org").unwrap();
        let bundle = X509Bundle::from_x509_authorities(trust_domain, &[authority]).unwrap();
        let mut bundle_set = X509BundleSet::new();
        bundle_set.add_bundle(bundle);
        Arc::new(X509Context::new([svid], Arc::new(bundle_set)))
    }

    #[tokio::test]
    async fn initial_sync_fails_fast_on_invalid_argument() {
        // A gRPC INVALID_ARGUMENT response indicates the Workload API rejected the
        // request itself (non-conforming server, metadata-mangling proxy, etc.), so
        // retrying with backoff would just repeat the same rejection. Initial sync must
        // return promptly instead of retrying indefinitely.
        let attempts = Arc::new(AtomicUsize::new(0));
        let attempts_clone = Arc::clone(&attempts);
        let make_client: ClientFactory = Arc::new(move || {
            attempts_clone.fetch_add(1, Ordering::SeqCst);
            Box::pin(async { Err(invalid_argument_error()) })
        });

        let cancel = CancellationToken::new();
        // Deliberately large backoff: if the fast-fail path regresses and the retry loop
        // sleeps before giving up, the outer timeout below will catch it rather than
        // hanging the test suite.
        let reconnect = ReconnectConfig::new(Duration::from_secs(30), Duration::from_secs(60));

        let result = tokio::time::timeout(
            Duration::from_millis(500),
            initial_sync_with_retry(
                &make_client,
                None,
                &cancel,
                reconnect,
                ResourceLimits::default(),
                None,
            ),
        )
        .await
        .expect("initial sync should fail fast instead of retrying on INVALID_ARGUMENT");

        assert!(
            matches!(
                &result,
                Err(X509SourceError::Source(e)) if e.is_invalid_argument()
            ),
            "expected an invalid-argument Source error, got: {result:?}"
        );
        assert_eq!(
            attempts.load(Ordering::SeqCst),
            1,
            "must not retry client creation after an INVALID_ARGUMENT response"
        );
    }

    fn x509_svid_response(bundle: &[u8]) -> X509svidResponse {
        let cert = include_bytes!("../../tests/testdata/svid/x509/1-svid-chain.der");
        let key = include_bytes!("../../tests/testdata/svid/x509/1-key.der");
        X509svidResponse {
            svids: vec![X509svid {
                spiffe_id: "spiffe://example.org/service".into(),
                x509_svid: cert.to_vec().into(),
                x509_svid_key: key.to_vec().into(),
                bundle: bundle.to_vec().into(),
                hint: String::new(),
            }],
            ..X509svidResponse::default()
        }
    }

    /// Serves `FetchX509SVID` over loopback TCP. The Nth call streams `responses[N]`
    /// (the last response is repeated once exhausted). Returns the endpoint and call count.
    async fn spawn_fetch_x509_svid_server(
        responses: Vec<X509svidResponse>,
    ) -> (String, Arc<AtomicUsize>) {
        let calls = Arc::new(AtomicUsize::new(0));
        let calls_clone = Arc::clone(&calls);
        let responses = Arc::new(responses);

        let svc = tower::service_fn(move |req: http::Request<tonic::body::Body>| {
            let calls = Arc::clone(&calls_clone);
            let responses = Arc::clone(&responses);
            async move {
                let handler =
                    tower::service_fn(move |_: tonic::Request<X509svidRequest>| {
                        let call = calls.fetch_add(1, Ordering::SeqCst);
                        let response = responses
                            .get(call)
                            .or_else(|| responses.last())
                            .cloned()
                            .expect("at least one response");
                        async move {
                            Ok::<_, tonic::Status>(tonic::Response::new(futures::stream::iter([
                                Ok::<_, tonic::Status>(response),
                            ])))
                        }
                    });
                let mut grpc = tonic::server::Grpc::new(tonic_prost::ProstCodec::default());
                Ok::<_, std::convert::Infallible>(grpc.server_streaming(handler, req).await)
            }
        });

        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(
            tonic::transport::Server::builder()
                .serve_with_incoming(svc, tonic::transport::server::TcpIncoming::from(listener)),
        );

        (format!("tcp://{addr}"), calls)
    }

    #[tokio::test]
    async fn initial_sync_retries_malformed_response_until_valid() {
        let authority = include_bytes!("../../tests/testdata/bundle/x509/cert1.der");
        let (endpoint, calls) = spawn_fetch_x509_svid_server(vec![
            x509_svid_response(&[]),
            x509_svid_response(authority),
        ])
        .await;
        let make_client: ClientFactory = Arc::new(move || {
            let endpoint = endpoint.clone();
            Box::pin(async move { WorkloadApiClient::connect_to(endpoint).await })
        });
        let cancel = CancellationToken::new();
        let reconnect = ReconnectConfig::new(Duration::from_millis(10), Duration::from_millis(10));

        let (ctx, _) = tokio::time::timeout(
            Duration::from_secs(5),
            initial_sync_with_retry(
                &make_client,
                None,
                &cancel,
                reconnect,
                ResourceLimits::default(),
                None,
            ),
        )
        .await
        .expect("initial sync should reconnect after a malformed response")
        .expect("initial sync should accept the subsequent valid response");

        assert_eq!(calls.load(Ordering::SeqCst), 2);
        let bundle = ctx
            .bundle_set()
            .get(&TrustDomain::new("example.org").unwrap())
            .expect("local bundle from the valid response");
        let authorities: Vec<&[u8]> = bundle.authorities().iter().map(AsRef::as_ref).collect();
        assert_eq!(authorities, [authority.as_slice()]);
    }

    #[tokio::test]
    async fn malformed_update_retains_snapshot_without_notification_and_next_update_recovers() {
        let old_authority = include_bytes!("../../tests/testdata/bundle/x509/cert1.der");
        let new_authority = include_bytes!("../../tests/testdata/bundle/x509/cert2.der");
        let initial = context_with_bundle(old_authority);
        let recovered = context_with_bundle(new_authority);
        let source = X509Source::new_for_test(
            Arc::clone(&initial),
            ReconnectConfig::default(),
            ResourceLimits::default(),
            None,
            None,
        );
        let cancel = CancellationToken::new();

        let mut malformed_stream =
            futures::stream::iter([Err(WorkloadApiError::MissingRequiredField {
                field: "X509SVID.bundle",
            })]);
        source
            .inner_for_test()
            .process_stream_updates(&mut malformed_stream, &cancel, 1)
            .await;

        assert_eq!(source.x509_context().unwrap().as_ref(), initial.as_ref());
        assert_eq!(source.updated().last(), 0);

        let mut recovered_stream = futures::stream::iter([Ok(recovered.as_ref().clone())]);
        source
            .inner_for_test()
            .process_stream_updates(&mut recovered_stream, &cancel, 2)
            .await;

        assert_eq!(source.x509_context().unwrap().as_ref(), recovered.as_ref());
        assert_eq!(source.updated().last(), 1);
    }
}