Use cached dns results in data plane

data-plane/Cargo.toml

@@ -47,6 +47,7 @@ tower = { version = "0.4.13", features = ["util"] }
 tower-http = { version = "0.5.0", features = ["catch-panic"] }
 libc = "0.2.150"
 serial_test = "3.0.0"
+trust-dns-proto = "0.23.2"
 
 
 [dev-dependencies]

data-plane/src/dns/enclavedns.rs

@@ -1,7 +1,9 @@
 use super::error::DNSError;
 use bytes::Bytes;
+use shared::server::egress::cache_ip_for_allowlist;
 use shared::server::egress::check_dns_allowed_for_domain;
-use shared::server::egress::{cache_ip_for_allowlist, EgressDestinations};
+use shared::server::egress::get_cached_dns;
+use shared::server::egress::EgressDestinations;
 use shared::server::get_vsock_client;
 use shared::server::CID::Parent;
 use shared::DNS_PROXY_VSOCK_PORT;
@@ -12,6 +14,9 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::UdpSocket;
 use tokio::sync::{mpsc::Receiver, Semaphore};
 use tokio::time::timeout;
+use trust_dns_proto::op::{Message, ResponseCode};
+use trust_dns_proto::rr::Record;
+use trust_dns_proto::serialize::binary::BinEncodable;
 
 /// Empty struct for the DNS proxy that runs in the data plane
 pub struct EnclaveDnsProxy;
@@ -122,14 +127,42 @@ impl EnclaveDnsDriver {
         }
     }
 
+    fn get_dns_response(mut message: Message, records: Vec<Record>) -> Result<Vec<u8>, DNSError> {
+        message.set_response_code(ResponseCode::NoError);
+        message.add_answers(records);
+
+        let response_bytes = message.to_bytes()?;
+        Ok(response_bytes)
+    }
+
     /// Perform a DNS lookup using the proxy running on the Host
     async fn perform_dns_lookup(
         dns_packet: Bytes,
         request_upper_bound: Duration,
         allowed_destinations: EgressDestinations,
     ) -> Result<Bytes, DNSError> {
         // Check domain is allowed before proxying lookup
-        check_dns_allowed_for_domain(&dns_packet.clone(), &allowed_destinations)?;
+        let message = check_dns_allowed_for_domain(&dns_packet, &allowed_destinations)?;
+        let query = match message.query() {
+            Some(query) => query,
+            None => return Self::proxy_dns_request(request_upper_bound, dns_packet).await,
+        };
+        match get_cached_dns(query) {
+            Ok(record) => match record {
+                Some(record) => {
+                    let dns_response = Self::get_dns_response(message.clone(), record)?;
+                    Ok(Bytes::copy_from_slice(&dns_response))
+                }
+                None => Self::proxy_dns_request(request_upper_bound, dns_packet).await,
+            },
+            Err(_) => Self::proxy_dns_request(request_upper_bound, dns_packet).await,
+        }
+    }
+
+    async fn proxy_dns_request(
+        request_upper_bound: Duration,
+        dns_packet: Bytes,
+    ) -> Result<Bytes, DNSError> {
         // Attempt DNS lookup wth a timeout, flatten timeout errors into a DNS Error
         let dns_response =
             timeout(request_upper_bound, Self::forward_dns_lookup(dns_packet)).await??;

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),
 }

shared/Cargo.toml

@@ -32,7 +32,7 @@ base64 = "0.13.0"
 env_logger = "0.10.0"
 once_cell = { version = "1.19.0", optional = true }
 ttl_cache = { version ="0.5.1", optional = true }
-dns-parser = { version = "0.8.0", optional = true }
+trust-dns-proto = { version ="0.23.2", optional = true }
 
 [dev-dependencies]
 tokio-test = "0.4.2"
@@ -45,5 +45,5 @@ crate-type = ["lib"]
 
 [features]
 default = []
-network_egress = ["dep:once_cell", "dep:ttl_cache", "dep:dns-parser"]
+network_egress = ["dep:once_cell", "dep:ttl_cache", "dep:trust-dns-proto"]
 enclave = ["dep:tokio-vsock"]

shared/src/server/egress.rs

@@ -1,11 +1,16 @@
-use dns_parser::RData;
 use once_cell::sync::Lazy;
 use serde::Deserialize;
 use serde::Deserializer;
+use std::collections::HashMap;
+use std::net::IpAddr;
 use std::net::Ipv4Addr;
 use std::sync::Mutex;
 use std::time::Duration;
 use thiserror::Error;
+use trust_dns_proto::op::Message;
+use trust_dns_proto::op::Query;
+use trust_dns_proto::rr::Record;
+use trust_dns_proto::serialize::binary::BinDecodable;
 use ttl_cache::TtlCache;
 
 #[derive(Debug, Error)]
@@ -21,14 +26,19 @@ pub enum EgressError {
     #[error("TLS extension missing")]
     ExtensionMissing,
     #[error(transparent)]
-    DNSParseError(#[from] dns_parser::Error),
+    DNSParseError(#[from] trust_dns_proto::error::ProtoError),
     #[error("Could not obtain lock for IP cache")]
     CouldntObtainLock,
+    #[error("IP not found")]
+    EgressIpNotFound,
 }
 
 pub static ALLOWED_IPS_FROM_DNS: Lazy<Mutex<TtlCache<String, String>>> =
     Lazy::new(|| Mutex::new(TtlCache::new(1000)));
 
+pub static DOMAINS_CACHED_DNS: Lazy<Mutex<TtlCache<String, Vec<Record>>>> =
+    Lazy::new(|| Mutex::new(TtlCache::new(1000)));
+
 pub fn get_egress_allow_list_from_env() -> EgressDestinations {
     let domain_str = std::env::var("EV_EGRESS_ALLOW_LIST").unwrap_or("".to_string());
     get_egress_allow_list(domain_str)
@@ -76,35 +86,84 @@ pub fn check_domain_allow_list(
 pub fn check_dns_allowed_for_domain(
     packet: &[u8],
     destinations: &EgressDestinations,
-) -> Result<(), EgressError> {
-    let packet = dns_parser::Packet::parse(packet)?;
-    packet
-        .questions
-        .iter()
-        .try_for_each(|q| check_domain_allow_list(q.qname.to_string(), destinations))
+) -> Result<Message, EgressError> {
+    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];
+        check_domain_allow_list(domain.into(), destinations)
+    })?;
+    Ok(parsed_packet)
 }
 
 pub fn cache_ip_for_allowlist(packet: &[u8]) -> Result<(), EgressError> {
-    let packet = dns_parser::Packet::parse(packet)?;
-    packet.answers.iter().try_for_each(|ans| {
-        if let RData::A(ip) = ans.data {
-            cache_ip(ip.0.to_string(), ans)
-        } else {
-            Ok(())
+    let parsed_packet = Message::from_bytes(packet)?;
+    let mut dns_records: HashMap<String, (u32, Vec<Record>)> = HashMap::new();
+
+    for ans in parsed_packet.answers() {
+        let ip = ans
+            .data()
+            .ok_or(EgressError::EgressIpNotFound)?
+            .ip_addr()
+            .ok_or(EgressError::EgressIpNotFound)?;
+
+        if let IpAddr::V4(id_addr) = ip {
+            dns_records
+                .entry(ans.name().to_string())
+                .and_modify(|(ttl, records)| {
+                    records.push(ans.clone());
+                    if ans.ttl() < *ttl {
+                        *ttl = ans.ttl();
+                    }
+                })
+                .or_insert_with(|| (ans.ttl(), vec![ans.clone()]));
+
+            cache_ip(id_addr.to_string(), ans.name().to_string(), ans.ttl())?;
         }
-    })
+    }
+
+    for (name, (ttl, records)) in dns_records {
+        cache_dns_record(name, records, ttl as u64)?;
+    }
+
+    Ok(())
+}
+
+pub fn get_ip_from_cache(ip: String) -> Result<(), EgressError> {
+    let cache = match ALLOWED_IPS_FROM_DNS.lock() {
+        Ok(cache) => cache,
+        Err(_) => return Err(EgressError::CouldntObtainLock),
+    };
+    match cache.get(&ip) {
+        Some(_) => Ok(()),
+        None => Err(EgressError::EgressIpNotAllowed(ip)),
+    }
+}
+
+pub fn get_cached_dns(query: &Query) -> Result<Option<Vec<Record>>, EgressError> {
+    let domain = query.name().to_string();
+    let cache = match DOMAINS_CACHED_DNS.lock() {
+        Ok(cache) => cache,
+        Err(_) => return Err(EgressError::CouldntObtainLock),
+    };
+    Ok(cache.get(&domain).cloned())
 }
 
-fn cache_ip(ip: String, answer: &dns_parser::ResourceRecord<'_>) -> Result<(), EgressError> {
+fn cache_ip(ip: String, name: String, ttl: u32) -> Result<(), EgressError> {
     let mut cache = match ALLOWED_IPS_FROM_DNS.lock() {
         Ok(cache) => cache,
         Err(_) => return Err(EgressError::CouldntObtainLock),
     };
-    cache.insert(
-        ip,
-        answer.name.to_string(),
-        Duration::from_secs(answer.ttl.into()),
-    );
+    cache.insert(ip, name, Duration::from_secs(ttl as u64));
+    Ok(())
+}
+
+fn cache_dns_record(name: String, answers: Vec<Record>, ttl: u64) -> Result<(), EgressError> {
+    let mut cache = match DOMAINS_CACHED_DNS.lock() {
+        Ok(cache) => cache,
+        Err(_) => return Err(EgressError::CouldntObtainLock),
+    };
+    cache.insert(name, answers.clone(), Duration::from_secs(ttl));
     Ok(())
 }
 
@@ -114,22 +173,14 @@ pub fn check_ip_allow_list(
 ) -> Result<(), EgressError> {
     if allowed_destinations.allow_all
         || allowed_destinations.ips.contains(&ip)
-        || is_valid_ip_from_dns(ip.clone())?
+        || get_ip_from_cache(ip.clone()).is_ok()
     {
         Ok(())
     } else {
         Err(EgressError::EgressIpNotAllowed(ip))
     }
 }
 
-fn is_valid_ip_from_dns(ip: String) -> Result<bool, EgressError> {
-    let cache = match ALLOWED_IPS_FROM_DNS.lock() {
-        Ok(cache) => cache,
-        Err(_) => return Err(EgressError::CouldntObtainLock),
-    };
-    Ok(cache.get(&ip).is_some())
-}
-
 #[derive(Clone, PartialEq, Debug, Deserialize)]
 pub struct EgressDestinations {
     pub wildcard: Vec<String>,
@@ -154,11 +205,21 @@ pub struct EgressConfig {
 
 #[cfg(test)]
 mod tests {
+    use crate::server::egress::cache_ip_for_allowlist;
     use crate::server::egress::check_domain_allow_list;
     use crate::server::egress::check_ip_allow_list;
     use crate::server::egress::get_egress_allow_list_from_env;
     use crate::server::egress::EgressDestinations;
     use crate::server::egress::EgressError::{EgressDomainNotAllowed, EgressIpNotAllowed};
+    use crate::server::egress::ALLOWED_IPS_FROM_DNS;
+    use crate::server::egress::DOMAINS_CACHED_DNS;
+    use std::net::Ipv4Addr;
+    use std::str::FromStr;
+    use std::thread::sleep;
+    use std::time::Duration;
+    use trust_dns_proto::op::Message;
+    use trust_dns_proto::rr::{DNSClass, Name, RData, Record, RecordType};
+    use trust_dns_proto::serialize::binary::BinEncodable;
 
     #[test]
     fn test_sequentially() {
@@ -170,6 +231,7 @@ mod tests {
         test_allow_valid_ip();
         test_allow_valid_domain();
         test_allow_valid_ip_for_all_allowed();
+        test_cache_ip_for_allowlist();
     }
 
     fn test_valid_all_domains() {
@@ -271,4 +333,40 @@ mod tests {
         let result = check_domain_allow_list("a.domain.com".to_string(), &destinations);
         assert!(result.is_ok());
     }
+
+    fn test_cache_ip_for_allowlist() {
+        let mut message = Message::new();
+        let domain = "a.domain.com.".to_string();
+        let ip = "172.67.167.151".to_string();
+
+        let record_a = create_record(domain.clone(), 3);
+        let record_b = create_record(domain.clone(), 500);
+        let record_c = create_record(domain.clone(), 600);
+
+        let records = vec![record_a.clone(), record_b.clone(), record_c];
+
+        message.add_answers(records.clone());
+        let packet = message.to_bytes().unwrap();
+        cache_ip_for_allowlist(&packet).unwrap();
+        let cache = ALLOWED_IPS_FROM_DNS.lock().unwrap();
+        assert_eq!(cache.get(&ip), Some(domain.clone()).as_ref());
+        let cache = DOMAINS_CACHED_DNS.lock().unwrap();
+        assert_eq!(cache.get(&domain), Some(&records));
+
+        // Assert shortest TTL is used
+        sleep(Duration::from_secs(3));
+        assert_eq!(cache.get(&domain), None);
+    }
+
+    fn create_record(domain: String, ttl: u32) -> Record {
+        let mut record: Record = Record::new();
+        record.set_name(Name::from_str(&domain).unwrap());
+        record.set_record_type(RecordType::A);
+        record.set_dns_class(DNSClass::IN);
+        record.set_ttl(ttl);
+        record.set_data(Some(RData::A(trust_dns_proto::rr::rdata::A(
+            Ipv4Addr::new(172, 67, 167, 151),
+        ))));
+        record
+    }
 }