pkarr 7.0.0

Public-Key Addressable Resource Records (Pkarr); publish and resolve DNS records over Mainline DHT
Documentation
use crate::{PublicKey, SignedPacket, StoredNodeCount};

use super::dht::DhtBackend;
use super::publish_result_accumulator::PublishResultAccumulator;
use super::relay::RelayBackend;
use super::resolve_result_accumulator::ResolveResultAccumulator;
use super::{BackendResolvePolicy, CacheContext};
use crate::client::{PublishError, ResolveError};

#[derive(Debug)]
pub(in crate::client) struct CombinedBackend {
    relay: RelayBackend,
    dht: DhtBackend,
}

impl CombinedBackend {
    pub(super) fn new(relay: RelayBackend, dht: DhtBackend) -> Self {
        Self { relay, dht }
    }

    pub(super) async fn publish(
        &self,
        signed_packet: &SignedPacket,
    ) -> Result<StoredNodeCount, PublishError> {
        let relay_publish = self.relay.publish(signed_packet);
        let dht_publish = self.dht.publish(signed_packet);
        let (relay_result, dht_result) = tokio::join!(relay_publish, dht_publish);

        let mut accumulator = PublishResultAccumulator::default();
        accumulator.record_result(relay_result);
        accumulator.record_result(dht_result);
        accumulator.into_result()
    }

    pub(super) async fn resolve(
        &self,
        public_key: &PublicKey,
        policy: BackendResolvePolicy<'_>,
    ) -> Result<SignedPacket, ResolveError> {
        let first_resolve = self.relay.resolve(public_key, policy);
        let second_resolve = self.dht.resolve(public_key, policy);

        match policy {
            BackendResolvePolicy::CacheOnly => first_resolve.await,
            BackendResolvePolicy::CacheFirst(context) => {
                tokio::pin!(first_resolve);
                tokio::pin!(second_resolve);
                tokio::select!(
                    result = &mut first_resolve => {
                        first_acceptable_or_wait_second(result, second_resolve, context).await
                    }
                    result = &mut second_resolve => {
                        first_acceptable_or_wait_second(result, first_resolve, context).await
                    }
                )
            }
            BackendResolvePolicy::NetworkOnly => {
                let (first_result, second_result) = tokio::join!(first_resolve, second_resolve);
                merge_resolve_results(first_result, second_result)
            }
        }
    }
}

async fn first_acceptable_or_wait_second(
    first: Result<SignedPacket, ResolveError>,
    second: impl std::future::Future<Output = Result<SignedPacket, ResolveError>>,
    context: CacheContext<'_>,
) -> Result<SignedPacket, ResolveError> {
    let mut accumulator = ResolveResultAccumulator::new(Some(context));
    if accumulator.record_result(first) {
        return accumulator.into_result();
    }
    accumulator.record_result(second.await);
    accumulator.into_result()
}

fn merge_resolve_results(
    first_result: Result<SignedPacket, ResolveError>,
    second_result: Result<SignedPacket, ResolveError>,
) -> Result<SignedPacket, ResolveError> {
    let mut accumulator = ResolveResultAccumulator::default();
    accumulator.record_result(first_result);
    accumulator.record_result(second_result);
    accumulator.into_result()
}

#[cfg(test)]
mod tests {
    use ntimestamp::Timestamp;

    use super::*;
    use crate::Keypair;

    fn signed_packet(timestamp: u64) -> SignedPacket {
        SignedPacket::builder()
            .timestamp(Timestamp::from(timestamp))
            .sign(&Keypair::random())
            .unwrap()
    }

    #[test]
    fn merge_resolve_results_preserves_packet_when_invalid_seq_is_not_newer() {
        let packet = signed_packet(10);

        assert_eq!(
            merge_resolve_results(
                Ok(packet.clone()),
                Err(ResolveError::InvalidSignedPacket { seq: 10 }),
            ),
            Ok(packet)
        );
    }

    #[test]
    fn merge_resolve_results_returns_invalid_signed_packet_when_seq_is_newer() {
        assert_eq!(
            merge_resolve_results(
                Ok(signed_packet(10)),
                Err(ResolveError::InvalidSignedPacket { seq: 11 }),
            ),
            Err(ResolveError::InvalidSignedPacket { seq: 11 })
        );
    }

    #[test]
    fn merge_resolve_results_keeps_highest_invalid_signed_packet_seq() {
        assert_eq!(
            merge_resolve_results(
                Err(ResolveError::InvalidSignedPacket { seq: 10 }),
                Err(ResolveError::InvalidSignedPacket { seq: 11 }),
            ),
            Err(ResolveError::InvalidSignedPacket { seq: 11 })
        );
    }

    #[tokio::test]
    async fn cache_first_waits_for_fresh_second_backend() {
        let keypair = Keypair::random();
        let mut expired = SignedPacket::builder()
            .timestamp(Timestamp::from(10))
            .sign(&keypair)
            .unwrap();
        expired.set_last_seen(&(Timestamp::now() - 60 * 1_000_000_u64));
        let fresh = SignedPacket::builder()
            .timestamp(Timestamp::from(11))
            .sign(&keypair)
            .unwrap();

        let resolved = first_acceptable_or_wait_second(
            Ok(expired),
            std::future::ready(Ok(fresh.clone())),
            CacheContext::new(None, 30, 30),
        )
        .await
        .unwrap();

        assert_eq!(resolved, fresh);
    }
}