Kez/rust/crates/kez-channels/src/dns.rs

//! DNS channel: looks up `_kez.<domain>` TXT records and verifies the first
//! one whose value parses as a KEZ proof (compact or legacy form).

use std::sync::Arc;

use async_trait::async_trait;
use hickory_resolver::{Resolver, proto::rr::RData};
use kez_core::{COMPACT_PROOF_PREFIX, Identity, dns_txt_name};

use crate::{Channel, ChannelError, ChannelHit, ChannelResult, parse_and_verify_for};

/// Resolver abstraction so tests can substitute a fake. The real
/// implementation uses `hickory-resolver` against the system config.
#[async_trait]
pub trait TxtResolver: Send + Sync {
    async fn lookup_txt(&self, name: &str) -> ChannelResult<Vec<String>>;
}

/// Production resolver: builds a tokio-backed hickory resolver per call.
pub struct SystemResolver;

#[async_trait]
impl TxtResolver for SystemResolver {
    async fn lookup_txt(&self, name: &str) -> ChannelResult<Vec<String>> {
        let resolver = Resolver::builder_tokio()
            .map_err(|e| ChannelError::Unreachable(format!("resolver config: {e}")))?
            .build()
            .map_err(|e| ChannelError::Unreachable(format!("resolver build: {e}")))?;
        let lookup = resolver
            .txt_lookup(name)
            .await
            .map_err(|e| ChannelError::Unreachable(format!("TXT lookup {name}: {e}")))?;

        let mut out = Vec::new();
        for record in lookup.answers() {
            let RData::TXT(txt) = &record.data else {
                continue;
            };
            // TXT RDATA is a sequence of <=255-byte segments; concatenate them
            // back into the original payload.
            let value: String = txt
                .txt_data
                .iter()
                .map(|bytes| String::from_utf8_lossy(bytes))
                .collect();
            out.push(value);
        }
        Ok(out)
    }
}

#[derive(Clone)]
pub struct DnsChannel {
    resolver: Arc<dyn TxtResolver>,
}

impl DnsChannel {
    pub fn new() -> Self {
        Self {
            resolver: Arc::new(SystemResolver),
        }
    }

    /// Inject a custom resolver (used by tests and any non-system DNS path).
    pub fn with_resolver(resolver: Arc<dyn TxtResolver>) -> Self {
        Self { resolver }
    }
}

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

#[async_trait]
impl Channel for DnsChannel {
    fn system(&self) -> &'static str {
        "dns"
    }

    async fn fetch_and_verify(&self, identity: &Identity) -> ChannelResult<ChannelHit> {
        let name = dns_txt_name(identity).map_err(|e| ChannelError::Other(e.into()))?;
        let records = self.resolver.lookup_txt(&name).await?;

        let mut last_error: Option<ChannelError> = None;
        for value in records {
            if !looks_like_kez_txt(&value) {
                continue;
            }
            match parse_and_verify_for(&value, identity) {
                Ok(hit) => return Ok(hit),
                Err(err) => last_error = Some(err),
            }
        }

        Err(last_error.unwrap_or_else(|| ChannelError::NotFound(identity.clone())))
    }
}

/// Pure: a TXT value looks like a KEZ proof if it starts with the compact
/// prefix (spec form) or the legacy `kez1:` prefix.
pub fn looks_like_kez_txt(value: &str) -> bool {
    value.starts_with(COMPACT_PROOF_PREFIX) || value.starts_with("kez1:")
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;
    use kez_core::{ClaimPayload, NostrSecret, SignedClaim};

    struct FakeResolver(Vec<String>);

    #[async_trait]
    impl TxtResolver for FakeResolver {
        async fn lookup_txt(&self, _name: &str) -> ChannelResult<Vec<String>> {
            Ok(self.0.clone())
        }
    }

    fn sign_dns(subject: &str) -> SignedClaim {
        let secret = NostrSecret::generate();
        let primary = Identity::parse(format!("nostr:{}", secret.npub())).unwrap();
        let subject = Identity::parse(subject).unwrap();
        SignedClaim::sign(ClaimPayload::new(subject, primary, Utc::now()), &secret).unwrap()
    }

    #[test]
    fn looks_like_kez_txt_accepts_both_prefixes() {
        assert!(looks_like_kez_txt("kez:z1:foo"));
        assert!(looks_like_kez_txt("kez1:{...}"));
        assert!(!looks_like_kez_txt("v=spf1 -all"));
        assert!(!looks_like_kez_txt(""));
    }

    #[tokio::test]
    async fn no_records_yields_not_found() {
        let channel = DnsChannel::with_resolver(Arc::new(FakeResolver(vec![])));
        let identity = Identity::parse("dns:jason.example.com").unwrap();
        let err = channel.fetch_and_verify(&identity).await.unwrap_err();
        assert!(matches!(err, ChannelError::NotFound(_)));
    }

    #[tokio::test]
    async fn ignores_non_kez_txt_then_falls_through() {
        let channel = DnsChannel::with_resolver(Arc::new(FakeResolver(vec![
            "v=spf1 -all".into(),
            "google-site-verification=abc".into(),
        ])));
        let identity = Identity::parse("dns:jason.example.com").unwrap();
        let err = channel.fetch_and_verify(&identity).await.unwrap_err();
        assert!(matches!(err, ChannelError::NotFound(_)));
    }

    #[tokio::test]
    async fn verifies_compact_proof() {
        let signed = sign_dns("dns:jason.example.com");
        let compact = signed.to_compact().unwrap();
        let channel = DnsChannel::with_resolver(Arc::new(FakeResolver(vec![compact])));
        let identity = Identity::parse("dns:jason.example.com").unwrap();
        let hit = channel.fetch_and_verify(&identity).await.unwrap();
        assert_eq!(hit.proof, signed);
    }

    #[tokio::test]
    async fn rejects_proof_for_wrong_subject() {
        let signed = sign_dns("dns:mallory.example.com");
        let compact = signed.to_compact().unwrap();
        let channel = DnsChannel::with_resolver(Arc::new(FakeResolver(vec![compact])));
        let identity = Identity::parse("dns:jason.example.com").unwrap();
        let err = channel.fetch_and_verify(&identity).await.unwrap_err();
        assert!(matches!(err, ChannelError::SubjectMismatch { .. }));
    }
}