remove unwraps

evervault · Jul 30, 2024 · da78712 · da78712
1 parent 9120453
commit da78712
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 9 deletions.
diff --git a/data-plane/src/dns/enclavedns.rs b/data-plane/src/dns/enclavedns.rs
@@ -126,7 +126,7 @@ impl EnclaveDnsDriver {
         }
     }
 
-    fn get_dns_answer(id: u16, record: Record) -> Vec<u8> {
+    fn get_dns_answer(id: u16, record: Record) -> Result<Vec<u8>, DNSError> {
         let mut message = Message::new();
         message.set_id(id);
         message.set_message_type(MessageType::Response);
@@ -137,10 +137,8 @@ impl EnclaveDnsDriver {
 
         message.add_answer(record);
 
-        let response_bytes = message.to_bytes().unwrap();
-        println!("TESTTTTT::::::: {:?}", response_bytes.len());
-
-        response_bytes
+        let response_bytes = message.to_bytes()?;
+        Ok(response_bytes)
     }
 
     /// Perform a DNS lookup using the proxy running on the Host
@@ -154,7 +152,7 @@ impl EnclaveDnsDriver {
         match get_cached_dns(packet.clone()) {
             Ok(record) => match record {
                 Some(record) => {
-                    let dns_response = Self::get_dns_answer(packet.header().id(), record);
+                    let dns_response = Self::get_dns_answer(packet.header().id(), record)?;
                     Ok(Bytes::copy_from_slice(&dns_response))
                 }
                 None => {

diff --git a/data-plane/src/dns/error.rs b/data-plane/src/dns/error.rs
@@ -19,4 +19,6 @@ pub enum DNSError {
     EgressError(#[from] EgressError),
     #[error("DNS lookup failed due to a timeout after: {0}")]
     DNSTimeout(#[from] tokio::time::error::Elapsed),
+    #[error("DNS parse error: {0}")]
+    ProtoError(#[from] trust_dns_proto::error::ProtoError),
 }
diff --git a/shared/src/server/egress.rs b/shared/src/server/egress.rs
@@ -64,7 +64,6 @@ pub fn check_domain_allow_list(
     domain: String,
     allowed_destinations: &EgressDestinations,
 ) -> Result<(), EgressError> {
-    println!("Checking domain: {}", domain);
     let valid_wildcard = allowed_destinations
         .wildcard
         .iter()
@@ -83,7 +82,7 @@ pub fn check_dns_allowed_for_domain(
     packet: &[u8],
     destinations: &EgressDestinations,
 ) -> Result<Message, EgressError> {
-    let parsed_packet = Message::from_bytes(packet).unwrap();
+    let parsed_packet = Message::from_bytes(packet)?;
     parsed_packet.queries().iter().try_for_each(|q| {
         let domain = q.name().to_string();
         let domain = &domain[..domain.len() - 1];
@@ -93,7 +92,7 @@ pub fn check_dns_allowed_for_domain(
 }
 
 pub fn cache_ip_for_allowlist(packet: &[u8]) -> Result<(), EgressError> {
-    let parsed_packet = Message::from_bytes(packet).unwrap();
+    let parsed_packet = Message::from_bytes(packet)?;
     parsed_packet.answers().iter().try_for_each(|ans| {
         let ip = ans.data().unwrap().ip_addr().unwrap().to_string();
         cache_ip(ip, ans.name().to_string(), ans.ttl())?;