iroh-dns-server 1.0.0

use core::fmt;
use std::{
    collections::{BTreeMap, btree_map},
    str::FromStr,
    sync::Arc,
};

use hickory_server::proto::{
    ProtoError,
    op::Message,
    rr::{
        Name, Record, RecordSet, RecordType, RrKey,
        domain::{IntoLabel, Label},
    },
    serialize::binary::BinDecodable,
};
use iroh_base::PublicKey;
use iroh_dns::pkarr::SignedPacket;
use n0_error::{e, stack_error};

/// A lightweight `[u8; 32]` wrapper used as a key for caches and database lookups.
///
/// Does not validate that the bytes are a valid public key on construction.
/// If constructed from invalid bytes, methods like `Display` and `to_z32`
/// will panic. In practice, bytes always originate from a validated `PublicKey`
/// or a database that was written from one.
#[derive(derive_more::Into, Eq, PartialEq, Ord, PartialOrd, Hash, Clone, Copy)]
pub(crate) struct PublicKeyBytes([u8; 32]);

#[stack_error(derive, add_meta, from_sources)]
pub(crate) enum InvalidPublicKeyBytes {
    #[error("invalid z-base-32 encoding")]
    InvalidEncoding,
    #[error("invalid length, must be 32 bytes")]
    InvalidLength,
}

impl PublicKeyBytes {
    /// Wraps raw bytes without validating they are a valid public key.
    ///
    /// # Safety (logical)
    ///
    /// The caller must ensure the bytes represent a valid Ed25519 public key.
    /// Passing invalid bytes will cause panics in `Display`, `Debug`, and `to_z32`.
    pub(crate) fn new_unchecked(bytes: [u8; 32]) -> Self {
        Self(bytes)
    }

    pub(crate) fn from_z32(s: &str) -> Result<Self, InvalidPublicKeyBytes> {
        let pk = PublicKey::from_z32(s).map_err(|_| e!(InvalidPublicKeyBytes::InvalidEncoding))?;
        Ok(Self(*pk.as_bytes()))
    }

    pub(crate) fn to_z32(self) -> String {
        PublicKey::from_bytes(&self.0).expect("valid key").to_z32()
    }

    pub(crate) fn as_bytes(&self) -> &[u8; 32] {
        &self.0
    }

    pub(crate) fn from_signed_packet(packet: &SignedPacket) -> Self {
        Self(*packet.public_key().as_bytes())
    }
}

impl fmt::Display for PublicKeyBytes {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.to_z32())
    }
}

impl fmt::Debug for PublicKeyBytes {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "PublicKeyBytes({})", self.to_z32())
    }
}

impl FromStr for PublicKeyBytes {
    type Err = InvalidPublicKeyBytes;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Self::from_z32(s)
    }
}

impl AsRef<[u8; 32]> for PublicKeyBytes {
    fn as_ref(&self) -> &[u8; 32] {
        &self.0
    }
}

pub(crate) fn signed_packet_to_hickory_message(
    signed_packet: &SignedPacket,
) -> Result<Message, ProtoError> {
    let encoded = signed_packet.encoded_packet();
    let message = Message::from_bytes(encoded)?;
    Ok(message)
}

pub(crate) fn signed_packet_to_hickory_records_without_origin(
    signed_packet: &SignedPacket,
    filter: impl Fn(&Record) -> bool,
) -> Result<(Label, BTreeMap<RrKey, Arc<RecordSet>>), ProtoError> {
    let common_zone = Label::from_utf8(&signed_packet.public_key().to_z32())?;
    let mut message = signed_packet_to_hickory_message(signed_packet)?;
    let answers = std::mem::take(&mut message.answers);
    let mut output: BTreeMap<RrKey, Arc<RecordSet>> = BTreeMap::new();
    for mut record in answers.into_iter() {
        // disallow SOA and NS records
        if matches!(record.record_type(), RecordType::SOA | RecordType::NS) {
            continue;
        }
        // expect the z32 encoded pubkey as root name
        let name = &record.name;
        if name.num_labels() < 1 {
            continue;
        }
        let zone = name.iter().next_back().unwrap().into_label()?;
        if zone != common_zone {
            continue;
        }
        if !filter(&record) {
            continue;
        }

        let name_without_zone =
            Name::from_labels(name.iter().take(name.num_labels() as usize - 1))?;
        record.name = name_without_zone;

        let rrkey = RrKey::new(record.name.clone().into(), record.record_type());
        match output.entry(rrkey) {
            btree_map::Entry::Vacant(e) => {
                let set: RecordSet = record.into();
                e.insert(Arc::new(set));
            }
            btree_map::Entry::Occupied(mut e) => {
                let set = e.get_mut();
                let serial = set.serial();
                // safe because we just created the arc and are sync iterating
                Arc::get_mut(set).unwrap().insert(record, serial);
            }
        }
    }
    Ok((common_zone, output))
}

pub(crate) fn record_set_append_origin(
    input: &RecordSet,
    origin: &Name,
    serial: u32,
) -> Result<RecordSet, ProtoError> {
    let new_name = input.name().clone().append_name(origin)?;
    let mut output = RecordSet::new(new_name.clone(), input.record_type(), serial);
    // TODO: less clones
    for record in input.records_without_rrsigs() {
        let mut record = record.clone();
        record.name = new_name.clone();
        output.insert(record, serial);
    }
    Ok(output)
}