From 94ce9142b27dad831bbe987ffbc0e65e95f2fb27 Mon Sep 17 00:00:00 2001
From: Peter Cai <peter@typeblog.net>
Date: Tue, 6 Apr 2021 07:58:35 +0800
Subject: [PATCH] server: use parsed record data to build response

---
 src/client.rs | 18 +++++-------------
 src/server.rs | 10 +++++++---
 src/util.rs   | 18 +++++++++++++++++-
 3 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index a7ed97c..d3ceacf 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -1,13 +1,9 @@
 use crate::cache::DnsCache;
 use crate::r#override::OverrideResolver;
-use domain::{
-    base::{
-        iana::{Opcode, Rcode},
-        octets::Parser,
-        rdata::{ParseRecordData, UnknownRecordData},
-        Dname, Message, MessageBuilder, ParsedDname, Question, Record, ToDname,
-    },
-    rdata::AllRecordData,
+use domain::base::{
+    iana::{Opcode, Rcode},
+    rdata::UnknownRecordData,
+    Dname, Message, MessageBuilder, ParsedDname, Question, Record, ToDname,
 };
 use js_sys::{ArrayBuffer, Uint8Array};
 use wasm_bindgen_futures::JsFuture;
@@ -178,11 +174,7 @@ impl Client {
                 ),
             );
             // Try to parse the record; if failed, fail this entire query
-            let parsed_record_data: Result<Option<AllRecordData<&[u8], ParsedDname<&[u8]>>>, _> =
-                AllRecordData::parse_data(
-                    owned_record.rtype(),
-                    &mut Parser::from_ref(owned_record.data().data().as_slice()),
-                );
+            let parsed_record_data = crate::util::parse_record_data(record.data());
 
             if let Err(_) = parsed_record_data {
                 return Err("Failed to parse response from upstream".to_string());
diff --git a/src/server.rs b/src/server.rs
index ca18eee..3e53b35 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -222,9 +222,13 @@ impl Server {
         // Set up the answer section
         let mut answer_builder = question_builder.answer();
         for r in records {
-            answer_builder
-                .push(r)
-                .map_err(|_| "Max answer size exceeded".to_string())?;
+            if let Ok(Some(data)) = crate::util::parse_record_data(r.data()) {
+                answer_builder
+                    .push(Record::new(r.owner().clone(), r.class(), r.ttl(), data))
+                    .map_err(|_| "Max answer size exceeded".to_string())?;
+            } else {
+                return Err("Cannot parse record data responded by resolver client".to_string());
+            }
         }
         Ok(answer_builder.into_message())
     }
diff --git a/src/util.rs b/src/util.rs
index 697d6fa..3c78dfd 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -1,4 +1,11 @@
-use domain::base::Message;
+use domain::{
+    base::{
+        octets::{ParseError, Parser},
+        rdata::{ParseRecordData, UnknownRecordData},
+        Message, ParsedDname,
+    },
+    rdata::AllRecordData,
+};
 use js_sys::{Math, Promise};
 use std::ops::Add;
 use std::{collections::hash_map::DefaultHasher, hash::Hasher};
@@ -21,6 +28,15 @@ pub fn parse_dns_wireformat(msg: &[u8]) -> Result<Message<Vec<u8>>, String> {
         .map_err(|_| "Failed to parse DNS wireformat message".to_string())
 }
 
+pub fn parse_record_data<T: AsRef<[u8]>>(
+    record: &UnknownRecordData<T>,
+) -> Result<Option<AllRecordData<&[u8], ParsedDname<&[u8]>>>, ParseError> {
+    AllRecordData::parse_data(
+        record.rtype(),
+        &mut Parser::from_ref(record.data().as_ref()),
+    )
+}
+
 // Rust wrapper around JS functions
 // For convenience, and also to work around bugs in rust-analyzer
 // which thinks all JS functions are "unsafe"