From cccdf5b5e9baa7007756b8de52f5511450a73593 Mon Sep 17 00:00:00 2001
From: mice_on_drugs <tobi@ottenweller.net>
Date: Sun, 3 Apr 2022 19:03:46 +0200
Subject: [PATCH] Implement TryFrom for DNSOpCode and DNSRCode. (#15)

Co-authored-by: Tobias Ottenweller <tobi@ottenweller.net>
Reviewed-on: https://git.comff.net/rustics/dns/pulls/15
---
 src/main.rs  |  8 +++--
 src/proto.rs | 95 ++++++++++++++++++++++++++--------------------------
 2 files changed, 53 insertions(+), 50 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 874a1a9..20c1e42 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,11 +1,15 @@
-use std::sync::mpsc;
 use std::net::UdpSocket;
+use std::sync::mpsc;
 use std::thread::{self, JoinHandle};
 
 mod proto;
 use proto::DNSHeader;
 
-fn listen() -> (JoinHandle<()>, mpsc::Sender<Vec<u8>>, mpsc::Receiver<Vec<u8>>) {
+fn listen() -> (
+    JoinHandle<()>,
+    mpsc::Sender<Vec<u8>>,
+    mpsc::Receiver<Vec<u8>>,
+) {
     let (tx, rx) = mpsc::channel();
 
     let tx_clone = tx.clone();
diff --git a/src/proto.rs b/src/proto.rs
index 8eb962a..290b9a1 100644
--- a/src/proto.rs
+++ b/src/proto.rs
@@ -20,15 +20,17 @@ pub enum DNSOpCode {
     Update = 5,
 }
 
-impl From<u16> for DNSOpCode {
-    fn from(val: u16) -> Self {
-        match val {
-            0 => DNSOpCode::Query,
-            1 => DNSOpCode::IQuery,
-            2 => DNSOpCode::Status,
-            4 => DNSOpCode::Notify,
-            5 => DNSOpCode::Update,
-            _ => panic!("KACKE"),
+impl TryFrom<u16> for DNSOpCode {
+    type Error = DNSParseError;
+
+    fn try_from(value: u16) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(DNSOpCode::Query),
+            1 => Ok(DNSOpCode::IQuery),
+            2 => Ok(DNSOpCode::Status),
+            4 => Ok(DNSOpCode::Notify),
+            5 => Ok(DNSOpCode::Update),
+            _ => Err(DNSParseError::DNSOpCodeInvalid),
         }
     }
 }
@@ -87,21 +89,23 @@ pub enum DNSRCode {
     NotZone = 10,
 }
 
-impl From<u16> for DNSRCode {
-    fn from(val: u16) -> Self {
-        match val {
-            0 => DNSRCode::NoError,
-            1 => DNSRCode::FormErr,
-            2 => DNSRCode::ServFail,
-            3 => DNSRCode::NXDomain,
-            4 => DNSRCode::NotImp,
-            5 => DNSRCode::Refused,
-            6 => DNSRCode::YXDomain,
-            7 => DNSRCode::YXRRSet,
-            8 => DNSRCode::NXRRSet,
-            9 => DNSRCode::NotAuth,
-            10 => DNSRCode::NotZone,
-            _ => panic!("KACKE"),
+impl TryFrom<u16> for DNSRCode {
+    type Error = DNSParseError;
+
+    fn try_from(value: u16) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(DNSRCode::NoError),
+            1 => Ok(DNSRCode::FormErr),
+            2 => Ok(DNSRCode::ServFail),
+            3 => Ok(DNSRCode::NXDomain),
+            4 => Ok(DNSRCode::NotImp),
+            5 => Ok(DNSRCode::Refused),
+            6 => Ok(DNSRCode::YXDomain),
+            7 => Ok(DNSRCode::YXRRSet),
+            8 => Ok(DNSRCode::NXRRSet),
+            9 => Ok(DNSRCode::NotAuth),
+            10 => Ok(DNSRCode::NotZone),
+            _ => Err(DNSParseError::DNSRCodeInvalid),
         }
     }
 }
@@ -160,26 +164,26 @@ impl DNSHeader {
 
     pub fn from_udp_datagram(datagram: &[u8]) -> Result<Self, DNSParseError> {
         if datagram.len() < 11 {
-            return Err(DNSParseError::DatagramLengthError);
+            return Err(DNSParseError::DatagramTooShort);
         }
 
-        let id = u16::from_be_bytes((&datagram[..2]).try_into()?);
+        let id = u16::from_be_bytes((&datagram[..2]).try_into().unwrap());
 
-        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;
-        let tc = (flags & DNSHeader::TC_MASK) != 0;
-        let rd = (flags & DNSHeader::RD_MASK) != 0;
-        let ra = (flags & DNSHeader::RA_MASK) != 0;
-        let ad = (flags & DNSHeader::AD_MASK) != 0;
-        let cd = (flags & DNSHeader::CD_MASK) != 0;
-        let rcode = DNSRCode::from((flags & DNSHeader::RCODE_MASK) >> DNSHeader::RCODE_OFFSET);
+        let flags = u16::from_be_bytes((&datagram[2..4]).try_into().unwrap());
+        let qr = (flags & Self::QR_MASK) != 0;
+        let opcode = DNSOpCode::try_from((flags & Self::OPCODE_MASK) >> Self::OPCODE_OFFSET)?;
+        let aa = (flags & Self::AA_MASK) != 0;
+        let tc = (flags & Self::TC_MASK) != 0;
+        let rd = (flags & Self::RD_MASK) != 0;
+        let ra = (flags & Self::RA_MASK) != 0;
+        let ad = (flags & Self::AD_MASK) != 0;
+        let cd = (flags & Self::CD_MASK) != 0;
+        let rcode = DNSRCode::try_from((flags & Self::RCODE_MASK) >> Self::RCODE_OFFSET)?;
 
-        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()?);
+        let qd_zo_count = u16::from_be_bytes((&datagram[4..6]).try_into().unwrap());
+        let an_pr_count = u16::from_be_bytes((&datagram[6..8]).try_into().unwrap());
+        let ns_up_count = u16::from_be_bytes((&datagram[8..10]).try_into().unwrap());
+        let ar_count = u16::from_be_bytes((&datagram[10..12]).try_into().unwrap());
 
         Ok(DNSHeader {
             id,
@@ -202,12 +206,7 @@ impl DNSHeader {
 
 #[derive(Debug)]
 pub enum DNSParseError {
-    DatagramLengthError,
-    SliceError(std::array::TryFromSliceError)
-}
-
-impl From<std::array::TryFromSliceError> for DNSParseError {
-    fn from(err: std::array::TryFromSliceError) -> Self {
-        DNSParseError::SliceError(err)
-    }
+    DatagramTooShort,
+    DNSOpCodeInvalid,
+    DNSRCodeInvalid,
 }