fix(network): preserve DNS answer order

Signed-off-by: John Myers <9696606+johntmyers@users.noreply.github.com>

Author: johntmyers

crates/openshell-supervisor-network/src/policy_dns/resolver.rs

@@ -11,7 +11,7 @@ use super::name::NormalizedName;
 use hickory_proto::op::{Message, MessageType, OpCode, Query, ResponseCode};
 use hickory_proto::rr::{Name, RData, RecordType};
 use openshell_core::net::connect_tcp_nodelay_best_effort;
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::{BTreeMap, BTreeSet, HashSet};
 use std::net::{IpAddr, SocketAddr};
 use std::time::Duration;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
@@ -204,8 +204,7 @@ impl TrustedResolver for SocketTrustedResolver {
                         .iter()
                         .map(|(address, _)| *address)
                         .collect::<Vec<_>>();
-                    addresses.sort_unstable();
-                    addresses.dedup();
+                    retain_first_addresses(&mut addresses);
                     addresses.truncate(MAX_RETAINED_ADDRESSES);
                     let address_ttl = records.iter().map(|(_, ttl)| *ttl).min().unwrap_or(1);
                     return Ok(TrustedAnswer {
@@ -239,6 +238,11 @@ impl TrustedResolver for SocketTrustedResolver {
     }
 }
 
+fn retain_first_addresses(addresses: &mut Vec<IpAddr>) {
+    let mut seen = HashSet::new();
+    addresses.retain(|address| seen.insert(*address));
+}
+
 fn parse_response(wire: &[u8], id: u16, query: &Query) -> Result<Message, ResolveError> {
     if wire.len() > MAX_DNS_MESSAGE_BYTES {
         return Err(ResolveError::Oversized);
@@ -311,6 +315,25 @@ mod tests {
     use openshell_core::net::set_tcp_nodelay_best_effort;
     use tokio::net::TcpListener;
 
+    #[test]
+    fn resolver_address_deduplication_preserves_answer_order() {
+        let mut addresses = vec![
+            "203.0.113.20".parse().unwrap(),
+            "203.0.113.10".parse().unwrap(),
+            "203.0.113.20".parse().unwrap(),
+        ];
+
+        retain_first_addresses(&mut addresses);
+
+        assert_eq!(
+            addresses,
+            vec![
+                "203.0.113.20".parse::<IpAddr>().unwrap(),
+                "203.0.113.10".parse::<IpAddr>().unwrap(),
+            ]
+        );
+    }
+
     #[test]
     fn answer_parser_keeps_only_requested_family_and_bounds_are_constants() {
         let owner = Name::from_ascii("db.example.").unwrap();

crates/openshell-supervisor-network/src/policy_dns/store.rs

@@ -9,7 +9,7 @@ use crate::proxy::destination::{
     DestinationRequest, DestinationValidationPlan, UpstreamConnector, build_pinned_validation_plan,
     validate_destination,
 };
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::{BTreeMap, BTreeSet, HashSet};
 use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
 use std::ops::RangeInclusive;
 use std::sync::RwLock;
@@ -89,14 +89,14 @@ impl MappingLookup {
         &self,
         endpoint_id: &PolicyEndpointId,
     ) -> Result<UpstreamConnector, MappingLookupError> {
+        let mut seen = HashSet::new();
         let addresses = self
             .record
             .contracts
             .iter()
             .filter(|contract| contract.port == self.port && &contract.endpoint_id == endpoint_id)
             .flat_map(|contract| contract.pinned_addresses.iter().copied())
-            .collect::<BTreeSet<_>>()
-            .into_iter()
+            .filter(|address| seen.insert(*address))
             .collect::<Vec<_>>();
         if addresses.is_empty() {
             return Err(MappingLookupError::EndpointMismatch);
@@ -655,4 +655,31 @@ mod tests {
         let connector = lookup.connector_for(&endpoint).await.unwrap();
         assert_eq!(connector.addrs(), &["203.0.113.8:5432".parse().unwrap()]);
     }
+
+    #[tokio::test]
+    async fn connector_preserves_resolver_address_order_while_deduplicating() {
+        let store = store(1);
+        let now = Instant::now();
+        let mut request = request("must-not-resolve.invalid", 1, Duration::from_secs(5));
+        request.contracts[0].pinned_addresses = vec![
+            "203.0.113.20".parse().unwrap(),
+            "203.0.113.10".parse().unwrap(),
+            "203.0.113.20".parse().unwrap(),
+        ];
+        let record = store.publish(request, 1, now).unwrap();
+        let lookup = store
+            .lookup(record.synthetic_address, 5432, 1, now)
+            .unwrap();
+        let endpoint = lookup.endpoint_ids().next().unwrap().clone();
+
+        let connector = lookup.connector_for(&endpoint).await.unwrap();
+
+        assert_eq!(
+            connector.addrs(),
+            &[
+                "203.0.113.20:5432".parse().unwrap(),
+                "203.0.113.10:5432".parse().unwrap(),
+            ]
+        );
+    }
 }