From 6fae7b778bf3ed7111c2302e403bbac7c1703d68 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Fri, 9 Feb 2024 23:06:04 +0000
Subject: [PATCH] Add a utility function to resolve CNAMEs in verified RRs

---
 src/query.rs      | 25 +++++++++++++++++++++++++
 src/validation.rs | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/src/query.rs b/src/query.rs
index b73b5b1..315c0a7 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -262,6 +262,31 @@ mod tests {
 		assert!(verified_rrs.expires > now);
 	}
 
+	#[test]
+	fn test_cname_query() {
+		for resolver in ["1.1.1.1:53", "8.8.8.8:53", "9.9.9.9:53"] {
+			let sockaddr = resolver.to_socket_addrs().unwrap().next().unwrap();
+			let query_name = "cname_test.matcorallo.com.".try_into().unwrap();
+			let (proof, _) = build_txt_proof(sockaddr, &query_name).unwrap();
+
+			let mut rrs = parse_rr_stream(&proof).unwrap();
+			rrs.shuffle(&mut rand::rngs::OsRng);
+			let verified_rrs = verify_rr_stream(&rrs).unwrap();
+			assert_eq!(verified_rrs.verified_rrs.len(), 2);
+
+			let now = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs();
+			assert!(verified_rrs.valid_from < now);
+			assert!(verified_rrs.expires > now);
+
+			let resolved_rrs = verified_rrs.resolve_name(&query_name);
+			assert_eq!(resolved_rrs.len(), 1);
+			if let RR::Txt(txt) = &resolved_rrs[0] {
+				assert_eq!(txt.name.as_str(), "txt_test.matcorallo.com.");
+				assert_eq!(txt.data, b"dnssec_prover_test");
+			} else { panic!(); }
+		}
+	}
+
 	#[cfg(feature = "tokio")]
 	use tokio_crate as tokio;
 
diff --git a/src/validation.rs b/src/validation.rs
index 615146e..7b2fab9 100644
--- a/src/validation.rs
+++ b/src/validation.rs
@@ -350,6 +350,29 @@ pub fn verify_rr_stream<'a>(inp: &'a [RR]) -> Result<VerifiedRRStream<'a>, Valid
 	}
 }
 
+impl<'a> VerifiedRRStream<'a> {
+	/// Given a name, resolve any [`CName`] records and return any verified records which were
+	/// pointed to by the original name.
+	///
+	/// Note that because of [`CName`]s, the [`RR::name`] in the returned records may or may not be
+	/// equal to `name`.
+	///
+	/// You MUST still check that the current UNIX time is between
+	/// [`VerifiedRRStream::valid_from`] and [`VerifiedRRStream::expires`] before
+	/// using any records returned here.
+	pub fn resolve_name<'b>(&self, mut name: &'b Name) -> Vec<&'a RR> where 'a: 'b {
+		loop {
+			let mut cname_search = self.verified_rrs.iter()
+				.filter(|rr| rr.name() == name)
+				.filter_map(|rr| if let RR::CName(cn) = rr { Some(cn) } else { None });
+			if let Some(cname) = cname_search.next() {
+				name = &cname.canonical_name;
+			}
+			return self.verified_rrs.iter().filter(|rr| rr.name() == name).map(|rr| *rr).collect();
+		}
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	#![allow(deprecated)]
@@ -664,6 +687,14 @@ mod tests {
 			assert_eq!(cname.name.as_str(), "cname_test.matcorallo.com.");
 			assert_eq!(cname.canonical_name.as_str(), "txt_test.matcorallo.com.");
 		} else { panic!(); }
+
+		let filtered_rrs =
+			verified_rrs.resolve_name(&"cname_test.matcorallo.com.".try_into().unwrap());
+		assert_eq!(filtered_rrs.len(), 1);
+		if let RR::Txt(txt) = &filtered_rrs[0] {
+			assert_eq!(txt.name.as_str(), "txt_test.matcorallo.com.");
+			assert_eq!(txt.data, b"dnssec_prover_test");
+		} else { panic!(); }
 	}
 
 	#[test]
@@ -697,6 +728,14 @@ mod tests {
 			assert_eq!(cname.name.as_str(), "test.cname_wildcard_test.matcorallo.com.");
 			assert_eq!(cname.canonical_name.as_str(), "cname.wildcard_test.matcorallo.com.");
 		} else { panic!(); }
+
+		let filtered_rrs =
+			verified_rrs.resolve_name(&"test.cname_wildcard_test.matcorallo.com.".try_into().unwrap());
+		assert_eq!(filtered_rrs.len(), 1);
+		if let RR::Txt(txt) = &filtered_rrs[0] {
+			assert_eq!(txt.name.as_str(), "cname.wildcard_test.matcorallo.com.");
+			assert_eq!(txt.data, b"wildcard_test");
+		} else { panic!(); }
 	}
 
 	#[test]
-- 
2.30.2