From 74b117bc9f5de9a697ce62e760b3cf705afe7122 Mon Sep 17 00:00:00 2001 From: Peter Cai Date: Tue, 6 Apr 2021 11:11:42 +0800 Subject: [PATCH] refactor: actually parse upstream response by using AllRecordData * UnknownRecordData does not know how to handle all record types properly, which resulted in occasionally corrupted responses. * Let's change everything to AllRecordData and actually parse the record at our side. * There is no simple function to convert AllRecordData to owned, so we need a few helper functions in the util module. --- src/cache.rs | 18 ++++++++++-------- src/client.rs | 44 ++++++++++++++++++++++++-------------------- src/override.rs | 19 ++++++++----------- src/server.rs | 6 +++--- src/util.rs | 42 +++++++++++++++++++++++++++++++++++++++++- 5 files changed, 86 insertions(+), 43 deletions(-) diff --git a/src/cache.rs b/src/cache.rs index 1402dc9..ca7aae2 100644 --- a/src/cache.rs +++ b/src/cache.rs @@ -1,5 +1,6 @@ use crate::kv; -use domain::base::{rdata::UnknownRecordData, Dname, Question, Record}; +use crate::util::OwnedRecordData; +use domain::base::{Dname, Question, Record}; use js_sys::Date; use serde::{Deserialize, Serialize}; @@ -22,13 +23,14 @@ impl DnsCache { pub async fn put_cache( &self, - record: &Record>, UnknownRecordData>>, + record: &Record>, OwnedRecordData>, ) -> Result<(), String> { let ttl = record.ttl(); + let data = crate::util::owned_record_data_to_buffer(record.data())?; self.store .put_buf_ttl_metadata( - &Self::record_to_key(record), - record.data().data(), + &Self::record_to_key(record, &data), + &data, ttl as u64, DnsCacheMetadata { created_ts: (Date::now() / 1000f64) as u64, @@ -41,7 +43,7 @@ impl DnsCache { pub async fn get_cache( &self, question: &Question>>, - ) -> Option>, UnknownRecordData>>>> { + ) -> Option>, OwnedRecordData>>> { // One question can have multiple cached records; so we list by prefix // Note that list_prefix returns 1000 records at maximum by default // We don't expect one question to have that many answers, so it @@ -80,14 +82,14 @@ impl DnsCache { question.qname().to_owned(), question.qclass(), remaining_ttl as u32, - UnknownRecordData::from_octets(question.qtype(), value), + crate::util::octets_to_owned_record_data(question.qtype(), &value).ok()?, )); } Some(ret) } - fn record_to_key(record: &Record>, UnknownRecordData>>) -> String { + fn record_to_key(record: &Record>, OwnedRecordData>, buf: &[u8]) -> String { format!( "{};{};{};{}", record.owner(), @@ -96,7 +98,7 @@ impl DnsCache { // We need to append the hash of the record data to the key // because one question might have multiple answers // When reading, we need to list the keys first - crate::util::hash_buf(record.data().data()) + crate::util::hash_buf(buf) ) } diff --git a/src/client.rs b/src/client.rs index de420ac..dc24811 100644 --- a/src/client.rs +++ b/src/client.rs @@ -1,10 +1,11 @@ use crate::cache::DnsCache; use crate::r#override::OverrideResolver; +use crate::util::OwnedRecordData; use domain::base::{ iana::{Opcode, Rcode}, - rdata::UnknownRecordData, Dname, Message, MessageBuilder, ParsedDname, Question, Record, ToDname, }; +use domain::rdata::AllRecordData; use js_sys::{ArrayBuffer, Uint8Array}; use wasm_bindgen_futures::JsFuture; use web_sys::{Headers, Request, RequestInit, Response}; @@ -28,7 +29,7 @@ impl Client { pub async fn query( &self, questions: Vec>>>, - ) -> Result>, UnknownRecordData>>>, String> { + ) -> Result>, OwnedRecordData>>, String> { // Attempt to answer locally first let (mut local_answers, questions) = self.try_answer_from_local(questions).await; if questions.len() == 0 { @@ -59,7 +60,7 @@ impl Client { &self, questions: Vec>>>, retries: usize, - ) -> Result>, UnknownRecordData>>>, String> { + ) -> Result>, OwnedRecordData>>, String> { let mut last_res = Err("Dummy".to_string()); for _ in 0..retries { last_res = self.query(questions.clone()).await; @@ -140,7 +141,7 @@ impl Client { fn extract_answers( msg: Message>, - ) -> Result>, UnknownRecordData>>>, String> { + ) -> Result>, OwnedRecordData>>, String> { let answer_section = msg .answer() .map_err(|_| "Failed to parse DNS answer from upstream".to_string())?; @@ -149,18 +150,19 @@ impl Client { // this is different from the server impl let answers: Vec<_> = answer_section.collect(); - let mut ret: Vec>, UnknownRecordData>>> = Vec::new(); + let mut ret: Vec>, OwnedRecordData>> = Vec::new(); for a in answers { let parsed_record = a.map_err(|_| "Failed to parse DNS answer record".to_string())?; - // Use UnknownRecordData here because we don't really care about the actual type of the record - // It saves time and saves sanity (because of the type signature of AllRecordData) - let record: Record>, UnknownRecordData<&[u8]>> = parsed_record - .to_record() - .map_err(|_| "Cannot parse record".to_string())? - .ok_or("Cannot parse record".to_string())?; - // Convert everything to owned for sanity in type signature... - // We'll need to do a copy before returning outside of the main - // query function anyway + // Actually parse the record + // Note that we cannot just use UnknownRecordData here and not parse it; + // it does not know how to parse all types of records correctly, which + // could corrupt the actual record data + let record: Record>, AllRecordData<&[u8], ParsedDname<&Vec>>> = + parsed_record + .to_record() + .map_err(|_| "Cannot parse record".to_string())? + .ok_or("Cannot parse record".to_string())?; + // Convert the record to owned for sanity in type signature let owned_record = Record::new( record .owner() @@ -168,10 +170,12 @@ impl Client { .map_err(|_| "Failed to parse Dname".to_string())?, record.class(), record.ttl(), - UnknownRecordData::from_octets( - record.data().rtype(), - record.data().data().to_vec(), - ), + match crate::util::to_owned_record_data(record.data()) { + Ok(data) => data, + // If this fails, it means that our resolver doesn't support the type yet + // so just skip this record + Err(_) => continue, + }, ); ret.push(owned_record); } @@ -185,7 +189,7 @@ impl Client { &self, questions: Vec>>>, ) -> ( - Vec>, UnknownRecordData>>>, + Vec>, OwnedRecordData>>, Vec>>>, ) { let mut answers = Vec::new(); @@ -206,7 +210,7 @@ impl Client { } #[allow(unused_must_use)] - async fn cache_answers(&self, answers: &[Record>, UnknownRecordData>>]) { + async fn cache_answers(&self, answers: &[Record>, OwnedRecordData>]) { for a in answers { // Ignore error -- we don't really care self.cache.put_cache(a).await; diff --git a/src/override.rs b/src/override.rs index 0510f59..617e74e 100644 --- a/src/override.rs +++ b/src/override.rs @@ -1,5 +1,6 @@ use crate::trie_map::TrieMap; -use domain::base::{rdata::UnknownRecordData, Compose, Dname, Question, Record, Rtype}; +use crate::util::OwnedRecordData; +use domain::base::{Dname, Question, Record, Rtype}; use domain::rdata::{Aaaa, AllRecordData, A}; use lazy_static::lazy_static; use std::collections::{HashMap, HashSet}; @@ -74,7 +75,7 @@ impl OverrideResolver { pub fn try_resolve( &self, question: &Question>>, - ) -> Option>, UnknownRecordData>>> { + ) -> Option>, OwnedRecordData>> { match question.qtype() { // We only handle resolution of IP addresses Rtype::A | Rtype::A6 | Rtype::Aaaa | Rtype::Cname | Rtype::Any => (), @@ -101,21 +102,17 @@ impl OverrideResolver { &self, question: &Question>>, addr: &IpAddr, - ) -> Option>, UnknownRecordData>>> { - let (rtype, rdata): (_, AllRecordData, Dname>>) = match addr { - IpAddr::V4(addr) => (Rtype::A, AllRecordData::A(A::new(addr.clone()))), - IpAddr::V6(addr) => (Rtype::Aaaa, AllRecordData::Aaaa(Aaaa::new(addr.clone()))), + ) -> Option>, OwnedRecordData>> { + let rdata: OwnedRecordData = match addr { + IpAddr::V4(addr) => AllRecordData::A(A::new(addr.clone())), + IpAddr::V6(addr) => AllRecordData::Aaaa(Aaaa::new(addr.clone())), }; - // Convert AllRecordData to UnknownRecordData to match the type signature - // since our resolver client doesn't really care about the actual type - let mut rdata_buf: Vec = Vec::new(); - rdata.compose(&mut rdata_buf).ok()?; let record = Record::new( question.qname().clone(), question.qclass(), self.override_ttl, - UnknownRecordData::from_octets(rtype, rdata_buf), + rdata, ); return Some(record); } diff --git a/src/server.rs b/src/server.rs index 8286d16..3a47ae9 100644 --- a/src/server.rs +++ b/src/server.rs @@ -3,8 +3,8 @@ use crate::r#override::OverrideResolver; use async_static::async_static; use domain::base::{ iana::{Opcode, Rcode}, - rdata::UnknownRecordData, - Dname, Message, MessageBuilder, Question, Record, ToDname, + record::AsRecord, + Dname, Message, MessageBuilder, Question, ToDname, }; use js_sys::{ArrayBuffer, Uint8Array}; use serde::Deserialize; @@ -204,7 +204,7 @@ impl Server { fn build_answer_wireformat( id: u16, questions: Vec>>>, - records: Vec>, UnknownRecordData>>>, + records: Vec, ) -> Result>, String> { let mut message_builder = MessageBuilder::new_vec(); // Set up the response header diff --git a/src/util.rs b/src/util.rs index 697d6fa..bb62125 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,4 +1,7 @@ -use domain::base::Message; +use domain::base::{ + octets::Parser, rdata::ParseRecordData, Compose, Dname, Message, ParsedDname, Rtype, ToDname, +}; +use domain::rdata::{AllRecordData, Cname}; use js_sys::{Math, Promise}; use std::ops::Add; use std::{collections::hash_map::DefaultHasher, hash::Hasher}; @@ -58,3 +61,40 @@ pub fn hash_buf(buf: &[u8]) -> u64 { hasher.write(buf); hasher.finish() } + +// Shorthand for a fully-owned AllRecordData variant +pub type OwnedRecordData = AllRecordData, Dname>>; + +// Convert a parsed AllRecordData instance to owned +pub fn to_owned_record_data, U: AsRef<[u8]>>( + data: &AllRecordData>, +) -> Result { + match data { + AllRecordData::A(data) => Ok(AllRecordData::A(data.clone())), + AllRecordData::Aaaa(data) => Ok(AllRecordData::Aaaa(data.clone())), + AllRecordData::Cname(data) => Ok(AllRecordData::Cname(Cname::new( + data.cname() + .to_dname() + .map_err(|_| "Failed parsing CNAME".to_string())?, + ))), + // TODO: Fill all of these in + _ => Err("Unsupported record type".to_string()), + } +} + +// Convert owned record data to Vec buffer +pub fn owned_record_data_to_buffer(data: &OwnedRecordData) -> Result, String> { + let mut ret: Vec = Vec::new(); + data.compose(&mut ret) + .map_err(|_| "Cannot convert owned record data to buffer".to_string())?; + Ok(ret) +} + +// Parse record data buffer and convert to owned record data +pub fn octets_to_owned_record_data(rtype: Rtype, octets: &[u8]) -> Result { + let parsed: AllRecordData<&[u8], ParsedDname<&[u8]>> = + ParseRecordData::parse_data(rtype, &mut Parser::from_ref(octets)) + .map_err(|_| "Cannot parse given record data".to_string())? + .ok_or("Given record data parsed to nothing".to_string())?; + to_owned_record_data(&parsed) +}