From bc278b7d399d2192f327f29e4b44f46cc3e73723 Mon Sep 17 00:00:00 2001
From: Tobias Ottenweller <tobi@ottenweller.net>
Date: Sun, 3 Apr 2022 17:01:07 +0200
Subject: [PATCH] Custom error type for parsing

---
 src/proto.rs | 48 ++++++++++++++++++++----------------------------
 1 file changed, 20 insertions(+), 28 deletions(-)

diff --git a/src/proto.rs b/src/proto.rs
index 6565a40..a29828a 100644
--- a/src/proto.rs
+++ b/src/proto.rs
@@ -38,6 +38,12 @@ pub enum DNSRCode {
     NOTINZONE = 10,
 }
 
+#[derive(Debug)]
+pub enum DNSParseErr {
+    DatagramLengthError,
+    SliceError(std::array::TryFromSliceError)
+}
+
 impl From<u16> for DNSRCode {
     fn from(val: u16) -> Self {
         match val {
@@ -109,18 +115,14 @@ impl DNSHeader {
     const OPCODE_OFFSET: u16 = 1;
     const RCODE_OFFSET: u16 = 11;
 
-    pub fn from_udp_datagram(datagram: &[u8]) -> Result<Self, &'static str> {
+    pub fn from_udp_datagram(datagram: &[u8]) -> Result<Self, DNSParseErr> {
         if datagram.len() < 11 {
-            return Err("Not enough data");
+            return Err(DNSParseErr::DatagramLengthError);
         }
 
-        let id = u16::from_be_bytes((&datagram[..2]).try_into().map_err(|_err| "Invalid id")?);
+        let id = u16::from_be_bytes((&datagram[..2]).try_into()?);
 
-        let flags = u16::from_be_bytes(
-            (&datagram[2..4])
-                .try_into()
-                .map_err(|_err| "Invalid flags")?,
-        );
+        let flags = u16::from_be_bytes((&datagram[2..4]).try_into()?);
         let qr = (flags & DNSHeader::QR_MASK) != 0;
         let opcode = DNSOpCode::from((flags & DNSHeader::OPCODE_MASK) >> DNSHeader::OPCODE_OFFSET);
         let aa = (flags & DNSHeader::AA_MASK) != 0;
@@ -131,26 +133,10 @@ impl DNSHeader {
         let cd = (flags & DNSHeader::CD_MASK) != 0;
         let rcode = DNSRCode::from((flags & DNSHeader::RCODE_MASK) >> DNSHeader::RCODE_OFFSET);
 
-        let qd_zo_count = u16::from_be_bytes(
-            (&datagram[4..6])
-                .try_into()
-                .map_err(|_err| "Invalid qd_zo_count")?,
-        );
-        let an_pr_count = u16::from_be_bytes(
-            (&datagram[6..8])
-                .try_into()
-                .map_err(|_err| "Invalid an_pr_count")?,
-        );
-        let ns_up_count = u16::from_be_bytes(
-            (&datagram[8..10])
-                .try_into()
-                .map_err(|_err| "Invalid ns_up_count")?,
-        );
-        let ar_count = u16::from_be_bytes(
-            (&datagram[10..12])
-                .try_into()
-                .map_err(|_err| "Invalid ar_count")?,
-        );
+        let qd_zo_count = u16::from_be_bytes((&datagram[4..6]).try_into()?);
+        let an_pr_count = u16::from_be_bytes((&datagram[6..8]).try_into()?);
+        let ns_up_count = u16::from_be_bytes((&datagram[8..10]).try_into()?);
+        let ar_count = u16::from_be_bytes((&datagram[10..12]).try_into()?);
 
         Ok(DNSHeader {
             id,
@@ -170,3 +156,9 @@ impl DNSHeader {
         })
     }
 }
+
+impl From<std::array::TryFromSliceError> for DNSParseErr {
+    fn from(err: std::array::TryFromSliceError) -> Self {
+        DNSParseErr::SliceError(err)
+    }
+}