seer_core/dns/

follow.rs

use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use tokio::sync::watch;
use tracing::{debug, instrument};

use super::records::{DnsRecord, RecordType};
use super::resolver::DnsResolver;
use crate::error::{Result, SeerError};

/// Upper bound on follow iterations. A non-interactive caller (API / Python
/// bindings) passing a huge count would otherwise schedule an effectively
/// unbounded long-running loop; the interval is already capped at 60 minutes.
const MAX_FOLLOW_ITERATIONS: usize = 10_000;

/// Configuration for DNS follow operation
#[derive(Debug, Clone)]
pub struct FollowConfig {
    /// Number of checks to perform
    pub iterations: usize,
    /// Interval between checks in seconds
    pub interval_secs: u64,
    /// Only output when records change
    pub changes_only: bool,
}

impl Default for FollowConfig {
    fn default() -> Self {
        Self {
            iterations: 10,
            interval_secs: 60,
            changes_only: false,
        }
    }
}

impl FollowConfig {
    /// Construct a new `FollowConfig`.
    ///
    /// Validates:
    /// - `iterations` must be >= 1
    /// - `interval_minutes` must be finite (not NaN / infinity)
    /// - `interval_minutes` must be non-negative
    /// - `interval_minutes` must be at most 60
    pub fn new(iterations: usize, interval_minutes: f64) -> Result<Self> {
        if iterations == 0 {
            return Err(SeerError::InvalidInput(
                "iterations must be at least 1".into(),
            ));
        }
        if iterations > MAX_FOLLOW_ITERATIONS {
            return Err(SeerError::InvalidInput(format!(
                "iterations must be at most {MAX_FOLLOW_ITERATIONS}"
            )));
        }
        if !interval_minutes.is_finite() {
            return Err(SeerError::InvalidInput(
                "interval_minutes must be a finite number".into(),
            ));
        }
        if interval_minutes < 0.0 {
            return Err(SeerError::InvalidInput(
                "interval_minutes must be non-negative".into(),
            ));
        }
        if interval_minutes > 60.0 {
            return Err(SeerError::InvalidInput(
                "interval_minutes must be at most 60".into(),
            ));
        }
        // A sub-second interval truncates to 0 seconds. For a multi-iteration
        // follow that means back-to-back live DNS queries with no spacing — a
        // self-inflicted query flood. Floor to 1s whenever more than one
        // iteration will run; a single-shot follow (iterations == 1) does no
        // looping and may keep a 0s interval.
        let mut interval_secs = (interval_minutes * 60.0) as u64;
        if iterations > 1 {
            interval_secs = interval_secs.max(1);
        }
        Ok(Self {
            iterations,
            interval_secs,
            changes_only: false,
        })
    }

    pub fn with_changes_only(mut self, changes_only: bool) -> Self {
        self.changes_only = changes_only;
        self
    }
}

/// Result of a single follow iteration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FollowIteration {
    /// Iteration number (1-based)
    pub iteration: usize,
    /// Total number of iterations
    pub total_iterations: usize,
    /// Timestamp of the check
    pub timestamp: DateTime<Utc>,
    /// Records found (or empty if error/NXDOMAIN)
    pub records: Vec<DnsRecord>,
    /// Whether records changed from previous iteration
    pub changed: bool,
    /// Values added since previous iteration
    pub added: Vec<String>,
    /// Values removed since previous iteration
    pub removed: Vec<String>,
    /// Error message if the check failed
    pub error: Option<String>,
}

impl FollowIteration {
    pub fn success(&self) -> bool {
        self.error.is_none()
    }

    pub fn record_count(&self) -> usize {
        self.records.len()
    }
}

/// Complete result of a follow operation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FollowResult {
    /// Domain that was monitored
    pub domain: String,
    /// Record type that was monitored
    pub record_type: RecordType,
    /// Nameserver used (if custom)
    pub nameserver: Option<String>,
    /// Configuration used
    pub iterations_requested: usize,
    pub interval_secs: u64,
    /// All iteration results
    pub iterations: Vec<FollowIteration>,
    /// Whether the operation was interrupted
    pub interrupted: bool,
    /// Total number of changes detected
    pub total_changes: usize,
    /// Start time
    pub started_at: DateTime<Utc>,
    /// End time
    pub ended_at: DateTime<Utc>,
}

impl FollowResult {
    pub fn completed_iterations(&self) -> usize {
        self.iterations.len()
    }

    pub fn successful_iterations(&self) -> usize {
        self.iterations.iter().filter(|i| i.success()).count()
    }

    pub fn failed_iterations(&self) -> usize {
        self.iterations.iter().filter(|i| !i.success()).count()
    }
}

/// Callback type for real-time progress updates
pub type FollowProgressCallback = Arc<dyn Fn(&FollowIteration) + Send + Sync>;

/// DNS Follower - monitors DNS records over time
#[derive(Clone)]
pub struct DnsFollower {
    resolver: DnsResolver,
}

impl Default for DnsFollower {
    fn default() -> Self {
        Self::new()
    }
}

impl DnsFollower {
    pub fn new() -> Self {
        Self {
            resolver: DnsResolver::new(),
        }
    }

    pub fn with_resolver(resolver: DnsResolver) -> Self {
        Self { resolver }
    }

    /// Follow DNS records over time
    #[instrument(skip(self, config, callback, cancel_rx))]
    pub async fn follow(
        &self,
        domain: &str,
        record_type: RecordType,
        nameserver: Option<&str>,
        config: FollowConfig,
        callback: Option<FollowProgressCallback>,
        cancel_rx: Option<watch::Receiver<bool>>,
    ) -> Result<FollowResult> {
        let domain = crate::validation::normalize_domain(domain)?;
        let started_at = Utc::now();
        let mut iterations: Vec<FollowIteration> = Vec::with_capacity(config.iterations);
        let mut previous_values: HashSet<String> = HashSet::new();
        let mut total_changes = 0;
        let mut interrupted = false;

        debug!(
            domain = %domain,
            record_type = %record_type,
            iterations = config.iterations,
            interval_secs = config.interval_secs,
            "Starting DNS follow"
        );

        for i in 0..config.iterations {
            // Check for cancellation
            if let Some(ref rx) = cancel_rx {
                if *rx.borrow() {
                    debug!("Follow operation cancelled");
                    interrupted = true;
                    break;
                }
            }

            let timestamp = Utc::now();
            let iteration_num = i + 1;

            // Perform DNS lookup
            let (records, error) = match self
                .resolver
                .resolve(&domain, record_type, nameserver)
                .await
            {
                Ok(records) => (records, None),
                Err(e) => {
                    debug!(domain = %domain, error = %e, "DNS follow query failed");
                    // Sanitized for external return; full detail logged above.
                    (Vec::new(), Some(e.sanitized_message()))
                }
            };

            // Extract record values (original casing) for this iteration.
            let current_values: Vec<String> = records.iter().map(|r| r.data.to_string()).collect();

            // Compare with previous iteration. The diff is case-insensitive so
            // a resolver applying 0x20 query-name randomization (returning the
            // same record with different casing) is not reported as a spurious
            // change; `added`/`removed` still carry the original casing.
            let (changed, added, removed) = if i == 0 {
                // First iteration - no previous to compare
                (false, Vec::new(), Vec::new())
            } else {
                diff_record_values(&current_values, &previous_values)
            };

            if changed {
                total_changes += 1;
            }

            let iteration = FollowIteration {
                iteration: iteration_num,
                total_iterations: config.iterations,
                timestamp,
                records,
                changed,
                added,
                removed,
                error,
            };

            // Call progress callback
            if let Some(ref cb) = callback {
                // Only call if not changes_only mode, or if this is first iteration or changed
                if !config.changes_only || iteration_num == 1 || changed {
                    cb(&iteration);
                }
            }

            iterations.push(iteration);
            // Store case-folded keys for the next iteration's comparison.
            previous_values = current_values
                .iter()
                .map(|v| v.to_ascii_lowercase())
                .collect();

            // Sleep before next iteration (unless this is the last one)
            if i < config.iterations - 1 {
                let sleep_duration = Duration::from_secs(config.interval_secs);

                // Use interruptible sleep
                if let Some(ref rx) = cancel_rx {
                    let mut rx_clone = rx.clone();
                    tokio::select! {
                        _ = tokio::time::sleep(sleep_duration) => {}
                        _ = rx_clone.changed() => {
                            if *rx_clone.borrow() {
                                debug!("Follow operation cancelled during sleep");
                                interrupted = true;
                                break;
                            }
                        }
                    }
                } else {
                    tokio::time::sleep(sleep_duration).await;
                }
            }
        }

        let ended_at = Utc::now();

        Ok(FollowResult {
            domain: domain.to_string(),
            record_type,
            nameserver: nameserver.map(|s| s.to_string()),
            iterations_requested: config.iterations,
            interval_secs: config.interval_secs,
            iterations,
            interrupted,
            total_changes,
            started_at,
            ended_at,
        })
    }

    /// Simple follow without callback or cancellation
    #[instrument(skip(self, config), fields(domain = %domain, record_type = ?record_type))]
    pub async fn follow_simple(
        &self,
        domain: &str,
        record_type: RecordType,
        nameserver: Option<&str>,
        config: FollowConfig,
    ) -> Result<FollowResult> {
        self.follow(domain, record_type, nameserver, config, None, None)
            .await
    }
}

/// Diffs the current iteration's record values against the previous
/// iteration's case-folded key set. Returns `(changed, added, removed)` where
/// `added`/`removed` preserve the original casing for display, but membership
/// is decided on the lowercased key — so a record that reappears with only a
/// case difference (0x20 query-name randomization) is not flagged as a change.
fn diff_record_values(
    current_values: &[String],
    previous_keys: &HashSet<String>,
) -> (bool, Vec<String>, Vec<String>) {
    let current_keys: HashSet<String> = current_values
        .iter()
        .map(|v| v.to_ascii_lowercase())
        .collect();

    // Added: present now (by key) but absent from the previous key set.
    let mut added: Vec<String> = current_values
        .iter()
        .filter(|v| !previous_keys.contains(&v.to_ascii_lowercase()))
        .cloned()
        .collect();
    // Removed: previous keys no longer present. The original casing is not
    // retained across iterations, so emit the lowercased key.
    let mut removed: Vec<String> = previous_keys
        .iter()
        .filter(|k| !current_keys.contains(*k))
        .cloned()
        .collect();
    added.sort();
    added.dedup();
    removed.sort();
    removed.dedup();

    let changed = !added.is_empty() || !removed.is_empty();
    (changed, added, removed)
}

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

    /// Builds a case-folded key set the same way `follow()` stores
    /// `previous_values`, for use in the diff helper tests.
    fn lowercased_set<const N: usize>(values: [&str; N]) -> HashSet<String> {
        values.iter().map(|v| v.to_ascii_lowercase()).collect()
    }

    #[tokio::test]
    async fn test_follow_config_default() {
        let config = FollowConfig::default();
        assert_eq!(config.iterations, 10);
        assert_eq!(config.interval_secs, 60);
        assert!(!config.changes_only);
    }

    /// Records that differ only in case between iterations (e.g. a resolver that
    /// applies 0x20 query-name randomization to NS answers) must NOT be reported
    /// as a change. The comparison key is case-folded.
    #[test]
    fn diff_values_ignores_case_only_differences() {
        let previous = lowercased_set(["NS1.EXAMPLE.COM.", "ns2.example.com."]);
        let current = vec![
            "ns1.example.com.".to_string(),
            "NS2.EXAMPLE.COM.".to_string(),
        ];
        let (changed, added, removed) = diff_record_values(&current, &previous);
        assert!(!changed, "case-only differences must not count as a change");
        assert!(added.is_empty(), "no added values: {added:?}");
        assert!(removed.is_empty(), "no removed values: {removed:?}");
    }

    /// A genuine value change is still detected, and `added`/`removed` carry the
    /// original casing for display.
    #[test]
    fn diff_values_detects_real_change_preserving_case() {
        let previous = lowercased_set(["ns1.example.com."]);
        let current = vec!["NS2.Example.Com.".to_string()];
        let (changed, added, removed) = diff_record_values(&current, &previous);
        assert!(changed, "a different nameserver is a real change");
        assert_eq!(added, vec!["NS2.Example.Com.".to_string()]);
        assert_eq!(removed, vec!["ns1.example.com.".to_string()]);
    }

    #[test]
    fn follow_config_rejects_unbounded_iterations() {
        assert!(FollowConfig::new(MAX_FOLLOW_ITERATIONS, 1.0).is_ok());
        let err = FollowConfig::new(MAX_FOLLOW_ITERATIONS + 1, 1.0).unwrap_err();
        assert!(matches!(err, SeerError::InvalidInput(_)));
        assert!(FollowConfig::new(usize::MAX, 1.0).is_err());
    }

    #[tokio::test]
    async fn test_follow_config_new() {
        let config = FollowConfig::new(5, 0.5).unwrap();
        assert_eq!(config.iterations, 5);
        assert_eq!(config.interval_secs, 30);
    }

    #[tokio::test]
    #[ignore = "live network; run with --ignored or SEER_LIVE_TESTS=1"]
    async fn test_follow_single_iteration() {
        let follower = DnsFollower::new();
        let config = FollowConfig::new(1, 0.0).unwrap();

        let result = follower
            .follow_simple("example.com", RecordType::A, None, config)
            .await;

        assert!(result.is_ok());
        let result = result.unwrap();
        assert_eq!(result.completed_iterations(), 1);
        assert!(!result.interrupted);
    }

    #[test]
    fn follow_config_rejects_zero_iterations() {
        assert!(FollowConfig::new(0, 1.0).is_err());
    }

    #[test]
    fn follow_config_rejects_infinite_interval() {
        assert!(FollowConfig::new(10, f64::INFINITY).is_err());
        assert!(FollowConfig::new(10, f64::NEG_INFINITY).is_err());
    }

    #[test]
    fn follow_config_rejects_nan_interval() {
        assert!(FollowConfig::new(10, f64::NAN).is_err());
    }

    #[test]
    fn follow_config_rejects_negative_interval() {
        assert!(FollowConfig::new(10, -1.0).is_err());
    }

    #[test]
    fn follow_config_rejects_interval_above_cap() {
        assert!(FollowConfig::new(10, 60.1).is_err());
    }

    #[test]
    fn follow_config_accepts_valid() {
        assert!(FollowConfig::new(10, 1.5).is_ok());
        assert!(FollowConfig::new(1, 0.0).is_ok());
        assert!(FollowConfig::new(1, 60.0).is_ok());
    }

    #[test]
    fn follow_config_floors_subsecond_interval_for_multi_iteration() {
        // A sub-second interval truncates to 0s; with many iterations that is
        // a back-to-back live-DNS query flood. Multi-iteration follows must
        // be floored to at least 1s between queries.
        let config = FollowConfig::new(10_000, 0.001).unwrap();
        assert!(
            config.interval_secs >= 1,
            "multi-iteration interval must be floored to >= 1s, got {}",
            config.interval_secs
        );
    }

    #[test]
    fn follow_config_allows_zero_interval_for_single_iteration() {
        // A single-shot follow does no looping, so a 0s interval is harmless
        // and must not be forced to 1s.
        let config = FollowConfig::new(1, 0.0).unwrap();
        assert_eq!(config.interval_secs, 0);
    }

    #[tokio::test]
    #[ignore = "live network; run with --ignored or SEER_LIVE_TESTS=1"]
    async fn follow_honors_cancel() {
        use tokio::sync::watch;

        let (tx, rx) = watch::channel(false);
        // 100 iterations with 30s intervals would take ~50 minutes.
        let config = FollowConfig::new(100, 0.5).unwrap();
        let follower = DnsFollower::new();

        let handle = tokio::spawn(async move {
            follower
                .follow("example.com", RecordType::A, None, config, None, Some(rx))
                .await
        });

        // Give the follow a tick to start and get into its first sleep.
        tokio::time::sleep(Duration::from_millis(200)).await;
        tx.send(true).unwrap();

        let joined = tokio::time::timeout(Duration::from_secs(10), handle)
            .await
            .expect("follow should return promptly after cancel");
        let result = joined.expect("join").expect("follow result");
        assert!(result.interrupted, "follow should be interrupted");
        assert!(
            result.completed_iterations() < 100,
            "should not complete all iterations"
        );
    }
}