road-runner-common 0.7.0

//! Live snapshot consumer (feature `authz-kafka`).
//!
//! Cold-start path: call the caller-supplied HTTP bootstrap closure (GET
//! `/internal/authz/snapshot`) so the PDP is ready before serving. Then follow the
//! compacted Kafka topic and hot-swap on every new snapshot.
//!
//! Compaction-safe cold start: we do NOT use a consumer group offset (a redeployed
//! pod could sit at a committed offset and never re-read the retained record).
//! Instead we `assign()` partition 0 at `Beginning` — on a compacted topic that is
//! the single latest snapshot — with auto-commit off.

use std::future::Future;
use std::sync::Arc;
use std::time::Duration;

use rdkafka::config::ClientConfig;
use rdkafka::consumer::{Consumer, StreamConsumer};
use rdkafka::{Message, Offset, TopicPartitionList};

use super::pdp::Pdp;
use super::verify::SnapshotVerifier;

pub struct SnapshotConsumerConfig {
    pub brokers: String,
    /// Compacted topic, e.g. `cex-policy.authz-snapshot`.
    pub topic: String,
}

/// Verify + install one envelope, respecting monotonic versioning. Returns whether
/// a new snapshot was installed.
fn apply_envelope(pdp: &Pdp, verifier: &SnapshotVerifier, bytes: &[u8]) -> bool {
    let current = pdp.version().unwrap_or(0);
    match verifier.verify(bytes, current) {
        Ok(payload) => {
            let version = payload.policy_version;
            pdp.install(payload);
            #[cfg(feature = "observability")]
            tracing::info!(version, "authz snapshot installed");
            let _ = version;
            true
        }
        Err(super::verify::VerifyError::Stale { .. }) => false,
        Err(e) => {
            #[cfg(feature = "observability")]
            tracing::warn!(error = %e, "authz snapshot rejected");
            let _ = e;
            false
        }
    }
}

/// Spawn the snapshot consumer task. `bootstrap` performs the one-shot HTTP fetch of
/// the current envelope (returning `Ok(None)` if unavailable, so the task falls back
/// to Kafka). The task runs forever; hold the returned `JoinHandle` if you want it.
pub fn spawn_snapshot_consumer<F, Fut>(
    pdp: Arc<Pdp>,
    verifier: SnapshotVerifier,
    cfg: SnapshotConsumerConfig,
    bootstrap: F,
) -> tokio::task::JoinHandle<()>
where
    F: Fn() -> Fut + Send + 'static,
    Fut: Future<Output = anyhow::Result<Option<Vec<u8>>>> + Send,
{
    tokio::spawn(async move {
        // 1) HTTP bootstrap so admin endpoints can be served ASAP. Retry a few times;
        //    Kafka is the durable source, so give up to it after a bounded effort.
        for attempt in 0..30u32 {
            match bootstrap().await {
                Ok(Some(bytes)) => {
                    if apply_envelope(&pdp, &verifier, &bytes) {
                        break;
                    }
                    break; // got a (possibly equal) snapshot; proceed to Kafka
                }
                Ok(None) => {}
                Err(e) => {
                    #[cfg(feature = "observability")]
                    tracing::warn!(error = %e, attempt, "authz bootstrap fetch failed");
                    let _ = (e, attempt);
                }
            }
            tokio::time::sleep(Duration::from_secs(2)).await;
        }

        // 2) Follow the compacted topic. Reconnect on any consumer error.
        loop {
            match build_consumer(&cfg) {
                Ok(consumer) => {
                    #[cfg(feature = "observability")]
                    tracing::info!(topic = %cfg.topic, "authz snapshot consumer connected");
                    loop {
                        match consumer.recv().await {
                            Ok(msg) => {
                                if let Some(bytes) = msg.payload() {
                                    apply_envelope(&pdp, &verifier, bytes);
                                }
                            }
                            Err(e) => {
                                #[cfg(feature = "observability")]
                                tracing::warn!(error = %e, "authz consumer recv error; reconnecting");
                                let _ = e;
                                break;
                            }
                        }
                    }
                }
                Err(e) => {
                    #[cfg(feature = "observability")]
                    tracing::warn!(error = %e, "authz consumer connect failed; retrying");
                    let _ = e;
                }
            }
            tokio::time::sleep(Duration::from_secs(3)).await;
        }
    })
}

fn build_consumer(cfg: &SnapshotConsumerConfig) -> Result<StreamConsumer, anyhow::Error> {
    // Unique group id per process so we never inherit a stale committed offset; we
    // assign manually (no subscribe/commit), so the group is effectively unused.
    let group = format!("authz-snapshot-{}", uuid::Uuid::new_v4());
    let consumer: StreamConsumer = ClientConfig::new()
        .set("bootstrap.servers", &cfg.brokers)
        .set("group.id", group)
        .set("enable.auto.commit", "false")
        .set("enable.partition.eof", "false")
        .create()?;

    let mut tpl = TopicPartitionList::new();
    tpl.add_partition_offset(&cfg.topic, 0, Offset::Beginning)?;
    consumer.assign(&tpl)?;
    Ok(consumer)
}