1use std::{io, num::NonZeroUsize};
2
3use dhttp_identity::certificate::{CertificateChainKey, CertificateSequence};
4use dquic::{
5 qbase::net::{AddrFamily, Family, addr::EndpointAddr as DquicEndpointAddr},
6 qresolve::{Resolve, Source},
7};
8use futures::future::BoxFuture;
9
10use crate::core::parser::record::endpoint::EndpointAddr as DnsEndpointAddr;
11
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct EndpointCandidateGroup {
14 pub chain: CertificateChainKey,
15 pub endpoints: Vec<DquicEndpointAddr>,
16 pub sources: Vec<Source>,
17}
18
19#[derive(Debug, Clone, Default, PartialEq, Eq)]
20pub struct EndpointCandidates {
21 pub groups: Vec<EndpointCandidateGroup>,
22}
23
24#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
25pub enum SequenceQuery {
26 #[default]
27 Default,
28 Exact(CertificateSequence),
29 Limit(NonZeroUsize),
30 All,
31}
32
33#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
34pub struct EndpointLookup {
35 pub sequences: SequenceQuery,
36 pub record_limit: Option<NonZeroUsize>,
37 pub family: Option<Family>,
38}
39
40impl EndpointLookup {
41 #[must_use]
42 pub fn exact(sequence: CertificateSequence) -> Self {
43 Self {
44 sequences: SequenceQuery::Exact(sequence),
45 record_limit: None,
46 family: None,
47 }
48 }
49
50 #[must_use]
51 pub fn limit(count: NonZeroUsize) -> Self {
52 Self {
53 sequences: SequenceQuery::Limit(count),
54 record_limit: None,
55 family: None,
56 }
57 }
58
59 #[must_use]
60 pub fn all() -> Self {
61 Self {
62 sequences: SequenceQuery::All,
63 record_limit: None,
64 family: None,
65 }
66 }
67
68 #[must_use]
69 pub fn with_record_limit(mut self, count: NonZeroUsize) -> Self {
70 self.record_limit = Some(count);
71 self
72 }
73
74 #[must_use]
75 pub fn with_family(mut self, family: Option<Family>) -> Self {
76 self.family = family;
77 self
78 }
79}
80
81#[cfg_attr(
82 not(any(feature = "h3", feature = "http", feature = "mdns", test)),
83 allow(dead_code)
84)]
85pub(crate) fn endpoint_matches_family(
86 endpoint: &DquicEndpointAddr,
87 family: Option<Family>,
88) -> bool {
89 family.is_none_or(|family| endpoint.addr().family() == family)
90}
91
92#[cfg_attr(
93 not(any(feature = "h3", feature = "http", feature = "mdns", test)),
94 allow(dead_code)
95)]
96pub(crate) fn filter_endpoint_candidate_groups<T>(
97 mut groups: EndpointCandidateGroups<T>,
98 family: Option<Family>,
99) -> EndpointCandidateGroups<T> {
100 if family.is_some() {
101 for (_, endpoints) in &mut groups {
102 endpoints.retain(|(_, endpoint)| endpoint_matches_family(endpoint, family));
103 }
104 groups.retain(|(_, endpoints)| !endpoints.is_empty());
105 }
106 groups
107}
108
109#[cfg(any(feature = "h3", feature = "http"))]
110pub(crate) fn append_endpoint_lookup_query(url: &mut url::Url, lookup: EndpointLookup) {
111 let mut pairs = url.query_pairs_mut();
112 match lookup.sequences {
113 SequenceQuery::Default => {}
114 SequenceQuery::Exact(sequence) => {
115 pairs.append_pair("sequence", &sequence.get().to_string());
116 }
117 SequenceQuery::Limit(limit) => {
118 pairs.append_pair("sequence_limit", &limit.get().to_string());
119 }
120 SequenceQuery::All => {
121 pairs.append_pair("sequence_limit", "all");
122 }
123 }
124 if let Some(limit) = lookup.record_limit {
125 pairs.append_pair("record_limit", &limit.get().to_string());
126 }
127}
128
129#[cfg(any(feature = "h3", feature = "http", test))]
130pub(crate) fn select_group_pairs<T>(
131 groups: Vec<(CertificateChainKey, T)>,
132 query: SequenceQuery,
133) -> Vec<(CertificateChainKey, T)> {
134 match query {
135 SequenceQuery::Default => groups.into_iter().take(3).collect(),
136 SequenceQuery::Exact(sequence) => groups
137 .into_iter()
138 .filter(|(chain, _)| chain.sequence() == sequence)
139 .collect(),
140 SequenceQuery::Limit(limit) => groups.into_iter().take(limit.get()).collect(),
141 SequenceQuery::All => groups,
142 }
143}
144
145pub type EndpointCandidateFuture<'a> = BoxFuture<'a, io::Result<EndpointCandidates>>;
146
147pub trait ResolveEndpointCandidates: Resolve {
148 fn lookup_endpoint_candidates<'a>(
149 &'a self,
150 name: &'a str,
151 lookup: EndpointLookup,
152 ) -> EndpointCandidateFuture<'a>;
153}
154
155pub type ArcEndpointCandidateResolver =
156 std::sync::Arc<dyn ResolveEndpointCandidates + Send + Sync + 'static>;
157
158#[cfg_attr(
159 not(any(feature = "h3", feature = "http", feature = "mdns", test)),
160 allow(dead_code)
161)]
162pub(crate) type EndpointCandidateGroups<T> =
163 Vec<(CertificateChainKey, Vec<(T, DquicEndpointAddr)>)>;
164
165#[cfg_attr(
166 not(any(feature = "h3", feature = "http", feature = "mdns", test)),
167 allow(dead_code)
168)]
169#[derive(Debug, Clone)]
170pub(crate) struct TaggedEndpointCandidate<T> {
171 pub(crate) tag: T,
172 pub(crate) record: DnsEndpointAddr,
173 pub(crate) fallback_chain_key: Option<CertificateChainKey>,
174}
175
176#[cfg_attr(
177 not(any(feature = "h3", feature = "http", feature = "mdns", test)),
178 allow(dead_code)
179)]
180pub(crate) fn grouped_endpoint_candidates<T>(
181 records: impl IntoIterator<Item = TaggedEndpointCandidate<T>>,
182) -> EndpointCandidateGroups<T> {
183 let mut groups: Vec<(CertificateChainKey, Vec<(T, DquicEndpointAddr)>)> = Vec::new();
184
185 for TaggedEndpointCandidate {
186 tag,
187 record,
188 fallback_chain_key,
189 } in records
190 {
191 let chain_key = effective_chain_key(&record, fallback_chain_key);
192 if !crate::core::certificate::is_primary_chain_key(&chain_key) {
193 continue;
194 }
195 let Ok(endpoint) = DquicEndpointAddr::try_from(record) else {
196 continue;
197 };
198
199 if let Some((_key, endpoints)) = groups.iter_mut().find(|(key, _)| *key == chain_key) {
200 endpoints.push((tag, endpoint));
201 } else {
202 groups.push((chain_key, vec![(tag, endpoint)]));
203 }
204 }
205
206 groups
207}
208
209#[cfg_attr(
210 not(any(feature = "h3", feature = "http", feature = "mdns", test)),
211 allow(dead_code)
212)]
213fn effective_chain_key(
214 record: &DnsEndpointAddr,
215 fallback_chain_key: Option<CertificateChainKey>,
216) -> CertificateChainKey {
217 if record.is_main() || record.sequence().is_some() {
218 return record.certificate_chain_key();
219 }
220
221 fallback_chain_key.unwrap_or_else(|| record.certificate_chain_key())
222}
223
224#[cfg(test)]
225mod tests {
226 use std::{
227 net::{SocketAddrV4, SocketAddrV6},
228 num::NonZeroUsize,
229 };
230
231 use dhttp_identity::certificate::CertificateSequence;
232
233 use super::*;
234
235 #[test]
236 fn endpoint_lookup_constructors_encode_valid_states() {
237 let one = NonZeroUsize::new(1).unwrap();
238 let exact = CertificateSequence::from(2u8);
239
240 assert_eq!(EndpointLookup::default().sequences, SequenceQuery::Default);
241 assert_eq!(
242 EndpointLookup::exact(exact).sequences,
243 SequenceQuery::Exact(exact)
244 );
245 assert_eq!(
246 EndpointLookup::limit(one).sequences,
247 SequenceQuery::Limit(one)
248 );
249 assert_eq!(EndpointLookup::all().sequences, SequenceQuery::All);
250 assert_eq!(
251 EndpointLookup::all().with_record_limit(one).record_limit,
252 Some(one)
253 );
254 assert_eq!(
255 EndpointLookup::default()
256 .with_family(Some(Family::V6))
257 .family,
258 Some(Family::V6)
259 );
260 }
261
262 fn direct(addr: &str, main: bool, sequence: u32) -> DnsEndpointAddr {
263 let socket: SocketAddrV4 = addr.parse().expect("socket addr");
264 let mut endpoint = DnsEndpointAddr::direct_v4(socket);
265 endpoint.set_main(main);
266 endpoint.set_sequence(CertificateSequence::try_from(sequence).unwrap());
267 endpoint
268 }
269
270 fn direct_v6(addr: &str, main: bool, sequence: u32) -> DnsEndpointAddr {
271 let socket: SocketAddrV6 = addr.parse().expect("socket addr");
272 let mut endpoint = DnsEndpointAddr::direct_v6(socket);
273 endpoint.set_main(main);
274 endpoint.set_sequence(CertificateSequence::try_from(sequence).unwrap());
275 endpoint
276 }
277
278 #[test]
279 fn family_filter_removes_mismatched_endpoints_and_empty_groups() {
280 let groups = grouped_endpoint_candidates([
281 TaggedEndpointCandidate {
282 tag: "v4",
283 record: direct("192.0.2.10:4433", true, 1),
284 fallback_chain_key: None,
285 },
286 TaggedEndpointCandidate {
287 tag: "v6",
288 record: direct_v6("[2001:db8::10]:4433", true, 1),
289 fallback_chain_key: None,
290 },
291 TaggedEndpointCandidate {
292 tag: "v6-only-group",
293 record: direct_v6("[2001:db8::20]:4433", true, 2),
294 fallback_chain_key: None,
295 },
296 ]);
297
298 let groups = filter_endpoint_candidate_groups(groups, Some(Family::V4));
299
300 assert_eq!(groups.len(), 1);
301 assert_eq!(groups[0].0.sequence().get(), 1);
302 assert_eq!(groups[0].1.len(), 1);
303 assert_eq!(groups[0].1[0].0, "v4");
304 }
305
306 #[test]
307 fn grouping_preserves_input_order_between_primary_sequences() {
308 let groups = grouped_endpoint_candidates([
309 TaggedEndpointCandidate {
310 tag: "wifi",
311 record: direct("192.0.2.10:4433", true, 2),
312 fallback_chain_key: None,
313 },
314 TaggedEndpointCandidate {
315 tag: "ethernet",
316 record: direct("192.0.2.20:4433", true, 1),
317 fallback_chain_key: None,
318 },
319 TaggedEndpointCandidate {
320 tag: "wifi-backup",
321 record: direct("192.0.2.11:4433", true, 2),
322 fallback_chain_key: None,
323 },
324 ]);
325
326 assert_eq!(groups.len(), 2);
327 assert_eq!(groups[0].0.usage().kind_flag(), "0");
328 assert_eq!(groups[0].0.sequence().get(), 2);
329 assert_eq!(groups[0].1.len(), 2);
330 assert_eq!(groups[1].0.usage().kind_flag(), "0");
331 assert_eq!(groups[1].0.sequence().get(), 1);
332 assert_eq!(groups[1].1.len(), 1);
333 }
334
335 #[test]
336 fn grouping_ignores_secondary_records() {
337 let groups = grouped_endpoint_candidates([
338 TaggedEndpointCandidate {
339 tag: "secondary",
340 record: direct("192.0.2.20:4433", false, 1),
341 fallback_chain_key: None,
342 },
343 TaggedEndpointCandidate {
344 tag: "primary",
345 record: direct("192.0.2.10:4433", true, 2),
346 fallback_chain_key: None,
347 },
348 ]);
349
350 assert_eq!(groups.len(), 1);
351 assert_eq!(groups[0].0.usage().kind_flag(), "0");
352 assert_eq!(groups[0].0.sequence().get(), 2);
353 assert_eq!(groups[0].1[0].0, "primary");
354 }
355
356 #[test]
357 fn sequence_query_selects_ordered_group_pairs() {
358 let pairs = || {
359 vec![
360 (
361 crate::core::certificate::primary_chain_key(CertificateSequence::from(2u8)),
362 "two",
363 ),
364 (
365 crate::core::certificate::primary_chain_key(CertificateSequence::from(1u8)),
366 "one",
367 ),
368 (
369 crate::core::certificate::primary_chain_key(CertificateSequence::from(3u8)),
370 "three",
371 ),
372 ]
373 };
374
375 assert_eq!(
376 select_group_pairs(
377 pairs(),
378 SequenceQuery::Exact(CertificateSequence::from(1u8)),
379 ),
380 vec![(
381 crate::core::certificate::primary_chain_key(CertificateSequence::from(1u8)),
382 "one",
383 )]
384 );
385 assert_eq!(
386 select_group_pairs(pairs(), SequenceQuery::Limit(NonZeroUsize::new(2).unwrap()),)
387 .into_iter()
388 .map(|(_, value)| value)
389 .collect::<Vec<_>>(),
390 vec!["two", "one"]
391 );
392 }
393
394 #[test]
395 fn grouping_uses_fallback_chain_key_for_unmarked_endpoint() {
396 let endpoint = DnsEndpointAddr::direct_v4("192.0.2.60:4433".parse().unwrap());
397 let groups = grouped_endpoint_candidates([TaggedEndpointCandidate {
398 tag: "h3",
399 record: endpoint,
400 fallback_chain_key: Some(crate::core::certificate::primary_chain_key(
401 CertificateSequence::from(3u8),
402 )),
403 }]);
404
405 assert_eq!(groups.len(), 1);
406 assert_eq!(groups[0].0.usage().kind_flag(), "0");
407 assert_eq!(groups[0].0.sequence().get(), 3);
408 assert_eq!(groups[0].1[0].0, "h3");
409 }
410}