From 6a609514871e0643ad281c5dea3db5fbfc32be50 Mon Sep 17 00:00:00 2001
From: mice_on_drugs <tobi@ottenweller.net>
Date: Sun, 8 May 2022 16:15:02 +0200
Subject: [PATCH] DNSMessageType tests added (#21)

Co-authored-by: Tobias Ottenweller <tobi@ottenweller.net>
Reviewed-on: https://git.comff.net/rustics/dns/pulls/21
Co-authored-by: mice_on_drugs <tobi@ottenweller.net>
Co-committed-by: mice_on_drugs <tobi@ottenweller.net>
---
 src/proto.rs  | 41 +++++++++++++++++++++++++++++------------
 udp_sender.py |  6 +++---
 2 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/src/proto.rs b/src/proto.rs
index 6818126..7680a7e 100644
--- a/src/proto.rs
+++ b/src/proto.rs
@@ -110,20 +110,17 @@ impl TryFrom<u8> for DNSRCode {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub enum DNSMessageType {
     Query,
     Response,
 }
 
-impl TryFrom<u8> for DNSMessageType {
-    type Error = DNSParseError;
-
-    fn try_from(value: u8) -> Result<Self, Self::Error> {
+impl From<bool> for DNSMessageType {
+    fn from(value: bool) -> DNSMessageType {
         match value {
-            0 => Ok(Self::Query),
-            1 => Ok(Self::Response),
-            _ => Err(DNSParseError::DNSRCodeInvalid),
+            false => Self::Query,
+            true => Self::Response,
         }
     }
 }
@@ -187,15 +184,15 @@ impl DNSHeader {
 
         let id = u16::from_be_bytes((&datagram[..2]).try_into().unwrap());
 
-        let message_type = DNSMessageType::try_from(datagram[2] & Self::QR_MASK)?;
+        let message_type = DNSMessageType::from(datagram[2] & Self::QR_MASK != 0);
         let opcode = DNSOpCode::try_from((datagram[2] & Self::OPCODE_MASK) >> Self::OPCODE_OFFSET)?;
         let authorative_answer = (datagram[2] & Self::AA_MASK) != 0;
         let truncated = (datagram[2] & Self::TC_MASK) != 0;
         let recursion_desired = (datagram[2] & Self::RD_MASK) != 0;
 
-        let recursion_available = (datagram[3] & Self::RA_MASK) != 0;
-        let authentic_data = (datagram[3] & Self::AD_MASK) != 0;
-        let checking_disabled = (datagram[3] & Self::CD_MASK) != 0;
+        let recursion_available = (datagram[2] & Self::RA_MASK) != 0;
+        let authentic_data = (datagram[2] & Self::AD_MASK) != 0;
+        let checking_disabled = (datagram[2] & Self::CD_MASK) != 0;
         let response_code = DNSRCode::try_from(datagram[3] & Self::RCODE_MASK)?;
 
         let query_count = u16::from_be_bytes((datagram[4..6]).try_into().unwrap());
@@ -287,4 +284,24 @@ mod tests {
         }
         Ok(())
     }
+
+    #[test]
+    fn parse_dns_header_message_type_query() -> Result<(), DNSParseError> {
+        let dns_query = [
+            0xff, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        ];
+        let parse_result = DNSHeader::from_udp_datagram(&dns_query);
+        assert_eq!(parse_result?.message_type, DNSMessageType::Query);
+        Ok(())
+    }
+
+    #[test]
+    fn parse_dns_header_message_type_response() -> Result<(), DNSParseError> {
+        let dns_query = [
+            0xff, 0x00, 0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        ];
+        let parse_result = DNSHeader::from_udp_datagram(&dns_query);
+        assert_eq!(parse_result?.message_type, DNSMessageType::Response);
+        Ok(())
+    }
 }
diff --git a/udp_sender.py b/udp_sender.py
index e5f68a5..9246d2f 100755
--- a/udp_sender.py
+++ b/udp_sender.py
@@ -3,11 +3,11 @@ import socket
 import sys
 
 
-UDP_SIZE = int(sys.argv[1])
+# UDP_SIZE = int(sys.argv[1])
 
 UDP_IP = "127.0.0.1"
 UDP_PORT = 13337
-MESSAGE = "A"
+MESSAGE = b"\xff\x00\x80\x00\x00\x00\x00\x00\x00\x00\x00\x00"
 
 sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-sock.sendto(bytes(MESSAGE*UDP_SIZE, "utf-8"), (UDP_IP, UDP_PORT))
\ No newline at end of file
+sock.sendto(MESSAGE, (UDP_IP, UDP_PORT))
\ No newline at end of file