rust-dnsbox/dnsbox/src/bin/resolver/cache/mod.rs

use dnsbox_base::common_types::{DnsName, Type};
use dnsbox_base::ser::RRData;
use failure::Error;
use futures::unsync::oneshot;
use futures::{Async, Future, Poll};
use std::cell::RefCell;
use std::collections::HashMap;
use std::mem::replace;
use std::rc::Rc;
use std::time::Duration;

mod root_hints;

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
pub enum Source {
	Hint,
	Glue,
	Authority,
}

#[derive(Clone, Debug)]
pub struct Entry {
	rrset: Vec<Box<dyn RRData>>,
	source: Source,
	ttl: u32,
}

impl Entry {
	pub fn new(rrset: Vec<Box<dyn RRData>>, source: Source) -> Self {
		unimplemented!()
	}

	/// whether TTL not reached yet. might also check DNSSEC stamps
	pub fn alive(&self) -> bool {
		true
	}

	pub fn needs_refresh(&self) -> bool {
		// TTL > 60 and only 25% ttl left before getting invalid
		false
	}
}

enum InnerEntry {
	Missing,
	Hit(Rc<Entry>),
	Refreshing(Rc<Entry>),
	Pending(Vec<oneshot::Sender<Rc<Entry>>>),
}

impl InnerEntry {
	fn restart(&mut self) -> (bool, CacheResult) {
		// start new lookup
		let (tx, rx) = oneshot::channel();
		replace(self, InnerEntry::Pending(vec![tx]));
		(true, CacheResult(InnerResult::Waiting(rx)))
	}

	fn get(&mut self) -> (bool, CacheResult) {
		match replace(self, InnerEntry::Missing) {
			InnerEntry::Missing => self.restart(),
			InnerEntry::Hit(e) => {
				if !e.alive() {
					self.restart();
				}
				let res = CacheResult(InnerResult::Ready(e.clone()));
				let start_lookup = if e.needs_refresh() {
					replace(self, InnerEntry::Refreshing(e));
					true
				} else {
					replace(self, InnerEntry::Hit(e));
					false
				};
				(start_lookup, res)
			},
			InnerEntry::Refreshing(e) => {
				if !e.alive() {
					// new request already pending; but need new wait queue
					(false, self.restart().1)
				} else {
					let res = CacheResult(InnerResult::Ready(e.clone()));
					replace(self, InnerEntry::Refreshing(e));
					(false, res)
				}
			},
			InnerEntry::Pending(mut queue) => {
				let (tx, rx) = oneshot::channel();
				queue.push(tx);
				replace(self, InnerEntry::Pending(queue));
				(false, CacheResult(InnerResult::Waiting(rx)))
			},
		}
	}
}

type NameEntry = HashMap<Type, InnerEntry>;

pub type LookupFunction = Box<dyn FnMut(DnsName) -> Box<dyn Future<Item = Entry, Error = Error>>>;

pub struct Cache {
	names: Rc<RefCell<HashMap<String, NameEntry>>>,
	timeout: Duration,
	lookup: RefCell<LookupFunction>,
}

impl Cache {
	pub fn new(timeout: Duration, lookup: LookupFunction) -> Self {
		Cache {
			names: Rc::new(RefCell::new(HashMap::new())),
			timeout: timeout,
			lookup: RefCell::new(lookup),
		}
	}

	fn start_lookup(&self, name: DnsName, rrtype: Type) {
		unimplemented!()
	}

	pub fn lookup(&self, name: DnsName, rrtype: Type) -> CacheResult {
		let mut name_str = format!("{}", name);
		name_str.make_ascii_lowercase();

		let (start_lookup, res) = {
			let mut names = self.names.borrow_mut();
			let types = names.entry(name_str).or_insert_with(|| HashMap::new());
			let entry = types.entry(rrtype).or_insert(InnerEntry::Missing);
			entry.get()
		};

		if start_lookup {
			self.start_lookup(name, rrtype);
		}

		res
	}

	/// Insert hints
	///
	/// # Panics
	///
	/// Panics when the `rrset` entries don't match `rrtype`.
	pub fn insert_hint(&self, name: DnsName, rrtype: Type, rrset: Vec<Box<dyn RRData>>) {
		for e in &rrset {
			assert_eq!(rrtype, e.rr_type());
		}
	}
}

pub struct CacheResult(InnerResult);

enum InnerResult {
	Ready(Rc<Entry>),
	Waiting(oneshot::Receiver<Rc<Entry>>),
	Finished,
}

impl Future for CacheResult {
	type Item = Rc<Entry>;
	type Error = Error;

	fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
		match replace(&mut self.0, InnerResult::Finished) {
			InnerResult::Ready(res) => Ok(Async::Ready(res)),
			InnerResult::Waiting(mut rc) => {
				match rc.poll()? {
					Async::Ready(res) => Ok(Async::Ready(res)),
					Async::NotReady => {
						// keep waiting
						replace(&mut self.0, InnerResult::Waiting(rc));
						Ok(Async::NotReady)
					},
				}
			},
			InnerResult::Finished => Ok(Async::NotReady),
		}
	}
}