Skip to main content

ddns/core/parser/record/
endpoint.rs

1use std::{
2    convert::TryFrom,
3    fmt::Display,
4    hash::{Hash, Hasher},
5    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6},
6    path::Path,
7};
8
9use base64::Engine;
10use bytes::BufMut;
11use dhttp_identity::certificate::{CertificateChainKey, CertificateChainKind, CertificateSequence};
12use dquic::qbase::net::addr::EndpointAddr as DquicEndpointAddr;
13use nom::{
14    IResult, Parser,
15    bytes::streaming::take,
16    combinator::{flat_map, map},
17    error::{ErrorKind, make_error},
18    number::streaming::{be_u8, be_u16, be_u32, be_u128},
19};
20use rustls::{SignatureScheme, pki_types::SubjectPublicKeyInfoDer};
21use snafu::{ResultExt, Snafu};
22
23use crate::core::parser::{
24    sigin,
25    varint::{VarInt, WriteVarInt, be_varint},
26};
27
28#[derive(Debug, Snafu)]
29#[snafu(module)]
30pub enum SignEndpointError {
31    #[snafu(display("failed to determine endpoint signature scheme"))]
32    SignatureScheme { source: sigin::SignatureSchemeError },
33    #[snafu(display("failed to sign endpoint address"))]
34    Sign {
35        source: dhttp_identity::identity::SignError,
36    },
37}
38
39#[derive(Debug, Snafu)]
40#[snafu(module)]
41pub enum EndpointSelectorError {
42    #[snafu(display("endpoint record sequence does not fit certificate sequence"))]
43    SequenceTooLarge { sequence: u64 },
44}
45
46/// EndpointAddress record (Type E = 266)
47///
48/// Unified endpoint format that encodes address family, routing, clustering and NAT information
49/// in a single flags byte, followed by optional fields.
50///
51/// ## Wire format
52///
53/// ```text
54/// +-------+-----------------+--------------------+----------------+----------------------------+
55/// | flags | sequence(varint)| addr               | load(optional) | signature (optional)       |
56/// +-------+-----------------+--------------------+----------------+----------------------------+
57/// | u8    | QUIC varint     | see addr layout    | f32            | scheme(u16)+len(varint)+N |
58/// +-------+-----------------+--------------------+----------------+----------------------------+
59///
60/// addr layout:
61/// +---------+-----------------------------------------+
62/// | kind    | addr fields                             |
63/// +---------+-----------------------------------------+
64/// | direct  | port(u16) + IP(u32/u128)                |
65/// +---------+-----------------------------------------+
66/// | nat     | outer_port(u16) + outer_IP(u32/u128)    |
67/// |         | + agent_port(u16) + agent_IP(u32/u128)  |
68/// +---------+-----------------------------------------+
69/// ```
70///
71/// ## flags (u8) bit layout (bit 1 = MSB)
72///
73/// - bit 1 (0x80): FAMILY    — address family: 0=IPv4, 1=IPv6
74/// - bit 2 (0x40): MAIN      — primary/backup: 1=primary, 0=backup
75/// - bit 3 (0x20): CLUSTERED — multiple hosts share this name; record includes a device sequence number
76/// - bit 4 (0x10): NAT       — endpoint is behind NAT; record includes the agent endpoint address
77/// - bit 5 (0x08): LOAD      — record carries 1-minute load average (f32)
78/// - bit 6 (0x04): reserved
79/// - bit 7 (0x02): reserved
80/// - bit 8 (0x01): SIGNED    — record carries a publisher key signature to prevent DNS poisoning
81///
82/// ## Field order
83///
84/// `flags` → `sequence` (if CLUSTERED) → `addr` → `agent addr` (if NAT) → `load` (if LOAD) → `signature` (if SIGNED)
85///
86#[derive(Debug, Clone, PartialEq, Eq, Hash)]
87pub struct EndpointSignature {
88    scheme: u16,
89    signature: Vec<u8>,
90}
91
92#[derive(Debug, Clone)]
93pub struct EndpointAddr {
94    flags: u8,
95    /// Device sequence number used when multiple hosts share a domain (CLUSTERED).
96    /// None means no sequence number.
97    sequence: Option<VarInt>,
98    /// 1-minute load average (present when LOAD flag is set)
99    load: Option<f32>,
100    signature: Option<EndpointSignature>,
101    /// Primary address (the direct address, or the outer/public address for NAT)
102    pub primary: SocketAddr,
103    /// Agent address used for NAT traversal (present when NAT flag is set)
104    pub agent: Option<SocketAddr>,
105}
106
107impl PartialEq for EndpointAddr {
108    fn eq(&self, other: &Self) -> bool {
109        self.flags == other.flags
110            && self.sequence == other.sequence
111            && self.load.map(f32::to_bits) == other.load.map(f32::to_bits)
112            && self.signature == other.signature
113            && self.primary == other.primary
114            && self.agent == other.agent
115    }
116}
117
118impl Eq for EndpointAddr {}
119
120impl Hash for EndpointAddr {
121    fn hash<H: Hasher>(&self, state: &mut H) {
122        self.flags.hash(state);
123        self.sequence.hash(state);
124        self.load.map(f32::to_bits).hash(state);
125        self.signature.hash(state);
126        self.primary.hash(state);
127        self.agent.hash(state);
128    }
129}
130
131impl EndpointAddr {
132    const FLAG_FAMILY: u8 = 0b1000_0000; // bit 1: 0=IPv4, 1=IPv6
133    const FLAG_MAIN: u8 = 0b0100_0000; // bit 2: 1=primary, 0=backup
134    const FLAG_CLUSTERED: u8 = 0b0010_0000; // bit 3: multiple hosts share this domain
135    const FLAG_NAT: u8 = 0b0001_0000; // bit 4: endpoint is behind NAT, agent address present
136    const FLAG_LOAD: u8 = 0b0000_1000; // bit 5: record carries load average
137    // bit 6 (0x04): reserved
138    // bit 7 (0x02): reserved
139    const FLAG_SIGNED: u8 = 0b0000_0001; // bit 8: record carries publisher signature
140
141    pub fn direct_v4(addr: SocketAddrV4) -> Self {
142        Self {
143            flags: 0, // IPv4 direct: FAMILY=0, NAT=0
144            sequence: None,
145            load: None,
146            signature: None,
147            primary: addr.into(),
148            agent: None,
149        }
150    }
151
152    pub fn direct_v6(addr: SocketAddrV6) -> Self {
153        Self {
154            flags: Self::FLAG_FAMILY, // IPv6 direct: FAMILY=1, NAT=0
155            sequence: None,
156            load: None,
157            signature: None,
158            primary: addr.into(),
159            agent: None,
160        }
161    }
162
163    /// Create an IPv4 endpoint that requires NAT traversal.
164    /// `outer` is the public address; `agent` is the NAT helper address.
165    pub fn nat_v4(outer: SocketAddrV4, agent: SocketAddrV4) -> Self {
166        Self {
167            flags: Self::FLAG_NAT, // IPv4 NAT: FAMILY=0, NAT=1
168            sequence: None,
169            load: None,
170            signature: None,
171            primary: outer.into(),
172            agent: Some(agent.into()),
173        }
174    }
175
176    /// Create an IPv6 endpoint that requires NAT traversal.
177    /// `outer` is the public address; `agent` is the NAT helper address.
178    pub fn nat_v6(outer: SocketAddrV6, agent: SocketAddrV6) -> Self {
179        Self {
180            flags: Self::FLAG_FAMILY | Self::FLAG_NAT, // IPv6 NAT: FAMILY=1, NAT=1
181            sequence: None,
182            load: None,
183            signature: None,
184            primary: outer.into(),
185            agent: Some(agent.into()),
186        }
187    }
188
189    /// Returns true if the address family is IPv6.
190    pub fn is_ipv6(&self) -> bool {
191        self.flags & Self::FLAG_FAMILY != 0
192    }
193
194    /// Returns true if NAT traversal is required (agent address is present).
195    pub fn is_nat(&self) -> bool {
196        self.flags & Self::FLAG_NAT != 0
197    }
198
199    /// Returns true if this domain maps to a cluster of hosts (sequence number is present).
200    pub fn is_clustered(&self) -> bool {
201        self.flags & Self::FLAG_CLUSTERED != 0
202    }
203
204    /// Returns true if the record carries a load average value.
205    pub fn is_load(&self) -> bool {
206        self.flags & Self::FLAG_LOAD != 0
207    }
208
209    pub fn set_clustered(&mut self, clustered: bool) {
210        if clustered {
211            self.flags |= Self::FLAG_CLUSTERED;
212        } else {
213            self.flags &= !Self::FLAG_CLUSTERED;
214            self.sequence = None; // clear sequence number
215        }
216    }
217
218    pub fn set_load(&mut self, load: Option<f32>) {
219        self.load = load;
220        if self.load.is_some() {
221            self.flags |= Self::FLAG_LOAD;
222        } else {
223            self.flags &= !Self::FLAG_LOAD;
224        }
225    }
226
227    pub async fn sign_with_authority(
228        &mut self,
229        authority: &(impl dhttp_identity::identity::LocalAuthority + ?Sized),
230    ) -> Result<(), SignEndpointError> {
231        self.set_signed(true);
232        let data = self.signed_data();
233        let scheme = sigin::signature_scheme(authority.public_key())
234            .context(sign_endpoint_error::SignatureSchemeSnafu)?;
235        let signature = authority
236            .sign(&data)
237            .await
238            .context(sign_endpoint_error::SignSnafu)?;
239        self.signature = Some(EndpointSignature {
240            scheme: u16::from(scheme),
241            signature,
242        });
243        Ok(())
244    }
245
246    pub fn verify_signature(
247        &self,
248        spki: SubjectPublicKeyInfoDer<'_>,
249    ) -> Result<bool, sigin::VerifyError> {
250        let Some(sig) = &self.signature else {
251            return Ok(false);
252        };
253        let data = self.signed_data();
254        sigin::verify(
255            spki,
256            SignatureScheme::from(sig.scheme),
257            &data,
258            &sig.signature,
259        )
260    }
261
262    pub fn verify_signature_from_der(&self, cert_der: &[u8]) -> Result<bool, sigin::VerifyError> {
263        let (_, cert) = x509_parser::parse_x509_certificate(cert_der).map_err(|e| {
264            sigin::VerifyError::InvalidCertificate {
265                details: e.to_string(),
266            }
267        })?;
268
269        let spki = SubjectPublicKeyInfoDer::from(cert.tbs_certificate.subject_pki.raw);
270        self.verify_signature(spki)
271    }
272
273    pub fn verify_signature_from_pem(&self, cert_pem: &[u8]) -> Result<bool, sigin::VerifyError> {
274        let mut reader = std::io::Cursor::new(cert_pem);
275        if let Some(item) = rustls_pemfile::certs(&mut reader).next() {
276            let cert_der = item.map_err(|e| sigin::VerifyError::InvalidPem { source: e })?;
277            return self.verify_signature_from_der(&cert_der);
278        }
279        Err(sigin::VerifyError::InvalidCertificate {
280            details: "No certificate found in PEM".to_string(),
281        })
282    }
283
284    pub fn verify_signature_from_base64(
285        &self,
286        cert_base64: &str,
287    ) -> Result<bool, sigin::VerifyError> {
288        let cert_base64 = cert_base64.trim();
289        let cert_der = base64::engine::general_purpose::STANDARD
290            .decode(cert_base64)
291            .map_err(|e| sigin::VerifyError::InvalidBase64 { source: e })?;
292        self.verify_signature_from_der(&cert_der)
293    }
294
295    pub fn verify_signature_from_file(
296        &self,
297        path: impl AsRef<Path>,
298    ) -> Result<bool, sigin::VerifyError> {
299        let contents = std::fs::read(path).map_err(|e| sigin::VerifyError::Io { source: e })?;
300        // Try PEM first
301        if let Ok(res) = self.verify_signature_from_pem(&contents) {
302            return Ok(res);
303        }
304        // If PEM failed, try DER
305        self.verify_signature_from_der(&contents)
306    }
307
308    pub fn is_main(&self) -> bool {
309        self.flags() & Self::FLAG_MAIN == Self::FLAG_MAIN
310    }
311
312    pub fn set_main(&mut self, is_main: bool) {
313        let flags = self.flags_mut();
314        if is_main {
315            *flags |= Self::FLAG_MAIN;
316        } else {
317            *flags &= !Self::FLAG_MAIN;
318        }
319    }
320
321    pub fn is_signed(&self) -> bool {
322        self.flags() & Self::FLAG_SIGNED == Self::FLAG_SIGNED
323    }
324
325    pub fn set_signed(&mut self, is_signed: bool) {
326        let flags = self.flags_mut();
327        if is_signed {
328            *flags |= Self::FLAG_SIGNED;
329        } else {
330            *flags &= !Self::FLAG_SIGNED;
331        }
332    }
333
334    pub fn encpding_size(&self) -> usize {
335        let mut meta_len = 1; // flags
336
337        // sequence is only encoded when CLUSTERED flag is set
338        if let Some(seq) = &self.sequence {
339            meta_len += seq.encoding_size();
340        }
341
342        if self.load.is_some() {
343            meta_len += 4; // f32
344        }
345
346        if self.is_signed()
347            && let Some(sig) = &self.signature
348        {
349            let sig_len =
350                VarInt::try_from(sig.signature.len() as u64).unwrap_or(VarInt::from_u32(0));
351            meta_len += 2 + sig_len.encoding_size() + sig.signature.len();
352        }
353
354        let addr_len = match (self.is_ipv6(), self.is_nat()) {
355            (false, false) => 2 + 4,      // IPv4 direct: port + ipv4
356            (false, true) => (2 + 4) * 2, // IPv4 NAT: (port + ipv4) * 2
357            (true, false) => 2 + 16,      // IPv6 direct: port + ipv6
358            (true, true) => (2 + 16) * 2, // IPv6 NAT: (port + ipv6) * 2
359        };
360
361        meta_len + addr_len
362    }
363
364    pub fn addr(&self) -> SocketAddr {
365        self.primary
366    }
367
368    pub fn agent_addr(&self) -> Option<SocketAddr> {
369        self.agent
370    }
371
372    pub fn set_sequence(&mut self, sequence: u64) {
373        if sequence > 0 {
374            self.sequence = Some(VarInt::from_u64(sequence).expect("Sequence too large"));
375            self.set_clustered(true);
376        } else {
377            self.sequence = None;
378            self.set_clustered(false);
379        }
380    }
381
382    pub fn certificate_chain_key(&self) -> Result<CertificateChainKey, EndpointSelectorError> {
383        let kind = if self.is_main() {
384            CertificateChainKind::Primary
385        } else {
386            CertificateChainKind::Secondary
387        };
388        let sequence = self.sequence.map(VarInt::into_inner).unwrap_or(0);
389        if sequence > u64::from(u32::MAX) {
390            return endpoint_selector_error::SequenceTooLargeSnafu { sequence }.fail();
391        }
392        let sequence = sequence as u32;
393        Ok(CertificateChainKey::new(
394            CertificateSequence::from(sequence),
395            kind,
396        ))
397    }
398
399    pub fn set_certificate_chain_key(&mut self, chain: &CertificateChainKey) {
400        self.set_main(chain.kind() == CertificateChainKind::Primary);
401        self.set_sequence(u64::from(chain.sequence().get()));
402    }
403
404    pub fn load(&self) -> Option<f32> {
405        self.load
406    }
407
408    fn flags(&self) -> u8 {
409        self.flags
410    }
411
412    fn flags_mut(&mut self) -> &mut u8 {
413        &mut self.flags
414    }
415
416    pub fn signature(&self) -> Option<&EndpointSignature> {
417        self.signature.as_ref()
418    }
419
420    pub fn signature_base64(&self) -> Option<String> {
421        self.signature
422            .as_ref()
423            .map(|sig| base64::engine::general_purpose::STANDARD.encode(&sig.signature))
424    }
425
426    fn write_base<B: BufMut>(&self, buf: &mut B) {
427        buf.put_u8(self.flags);
428
429        // Sequence is only written when CLUSTERED is set
430        if let Some(seq) = &self.sequence {
431            buf.put_varint(*seq);
432        }
433
434        // Write primary address
435        match self.primary {
436            SocketAddr::V4(addr) => buf.put_socket_addr_v4(&addr),
437            SocketAddr::V6(addr) => buf.put_socket_addr_v6(&addr),
438        }
439
440        // Write agent address when NAT traversal is required
441        if let Some(agent_addr) = &self.agent {
442            match agent_addr {
443                SocketAddr::V4(addr) => buf.put_socket_addr_v4(addr),
444                SocketAddr::V6(addr) => buf.put_socket_addr_v6(addr),
445            }
446        }
447
448        if let Some(load) = self.load {
449            buf.put_u32(load.to_bits());
450        }
451    }
452
453    fn signed_data(&self) -> Vec<u8> {
454        let mut unsigned = self.clone();
455        unsigned.set_signed(true);
456        unsigned.signature = None;
457        let mut buf = bytes::BytesMut::with_capacity(unsigned.encpding_size());
458        unsigned.write_base(&mut buf);
459        buf.to_vec()
460    }
461}
462
463pub(crate) trait WriteEndpointAddr {
464    fn put_endpoint_addr(&mut self, endpoint: &EndpointAddr);
465}
466
467impl<B: BufMut> WriteEndpointAddr for B {
468    fn put_endpoint_addr(&mut self, endpoint: &EndpointAddr) {
469        endpoint.write_base(self);
470        if endpoint.is_signed()
471            && let Some(sig) = endpoint.signature()
472        {
473            self.put_u16(sig.scheme);
474            let len = VarInt::try_from(sig.signature.len() as u64).unwrap_or(VarInt::from_u32(0));
475            self.put_varint(len);
476            self.put_slice(&sig.signature);
477        }
478    }
479}
480
481pub fn be_endpoint_addr(input: &[u8]) -> nom::IResult<&[u8], EndpointAddr> {
482    let (remain, flags) = be_u8(input)?;
483
484    let is_clustered = flags & EndpointAddr::FLAG_CLUSTERED != 0;
485    let is_ipv6 = flags & EndpointAddr::FLAG_FAMILY != 0;
486    let is_nat = flags & EndpointAddr::FLAG_NAT != 0;
487    let has_load = flags & EndpointAddr::FLAG_LOAD != 0;
488
489    // Sequence number is only present when CLUSTERED is set
490    let (remain, sequence) = if is_clustered {
491        let (remain, seq) = be_varint(remain)?;
492        (remain, Some(seq))
493    } else {
494        (remain, None)
495    };
496
497    let (remain, primary) = if is_ipv6 {
498        let (remain, addr) = be_socket_addr_v6(remain)?;
499        (remain, SocketAddr::V6(addr))
500    } else {
501        let (remain, addr) = be_socket_addr_v4(remain)?;
502        (remain, SocketAddr::V4(addr))
503    };
504
505    let (remain, agent) = if is_nat {
506        let agent_addr = if is_ipv6 {
507            let (remain, addr) = be_socket_addr_v6(remain)?;
508            (remain, SocketAddr::V6(addr))
509        } else {
510            let (remain, addr) = be_socket_addr_v4(remain)?;
511            (remain, SocketAddr::V4(addr))
512        };
513        let (remain, addr) = agent_addr;
514        (remain, Some(addr))
515    } else {
516        (remain, None)
517    };
518
519    let (remain, load) = if has_load {
520        let (remain, load) = be_u32(remain)?;
521        (remain, Some(f32::from_bits(load)))
522    } else {
523        (remain, None)
524    };
525
526    let (remain, signature) = be_endpoint_signature(remain, flags)?;
527
528    Ok((
529        remain,
530        EndpointAddr {
531            flags,
532            sequence,
533            load,
534            signature,
535            primary,
536            agent,
537        },
538    ))
539}
540
541/// Parse an EndpointAddr with backward compatibility.
542///
543/// - If `rdlen` matches a legacy fixed length, parse as address-only and fill in default flags.
544/// - Otherwise parse as the modern format (with `flags` + optional fields).
545///
546/// Note: Legacy and modern records differ in their leading bytes.
547/// Disambiguation is done via `RDLENGTH` to avoid misinterpreting the port's high byte as `flags`.
548pub(crate) fn be_endpoint_addr_compat(
549    input: &[u8],
550    rdlen: u16,
551) -> nom::IResult<&[u8], EndpointAddr> {
552    // Check for legacy fixed lengths
553    let legacy_lengths = [
554        6,  // IPv4 direct: port(2) + ip(4)
555        12, // IPv4 NAT:    (port(2) + ip(4)) * 2
556        18, // IPv6 direct: port(2) + ip(16)
557        36, // IPv6 NAT:    (port(2) + ip(16)) * 2
558    ];
559
560    if legacy_lengths.contains(&(rdlen as usize)) {
561        return be_legacy_endpoint_addr_by_length(input, rdlen);
562    }
563
564    // Modern format
565    be_endpoint_addr(input)
566}
567
568/// Parse a legacy EndpointAddr record, identified by its fixed RDLENGTH.
569fn be_legacy_endpoint_addr_by_length(
570    input: &[u8],
571    rdlen: u16,
572) -> nom::IResult<&[u8], EndpointAddr> {
573    match rdlen {
574        6 => {
575            // IPv4 direct
576            let (remain, addr) = be_socket_addr_v4(input)?;
577            Ok((
578                remain,
579                EndpointAddr {
580                    flags: 0,
581                    sequence: None,
582                    load: None,
583                    signature: None,
584                    primary: addr.into(),
585                    agent: None,
586                },
587            ))
588        }
589        12 => {
590            // IPv4 NAT
591            let (remain, primary) = be_socket_addr_v4(input)?;
592            let (remain, agent) = be_socket_addr_v4(remain)?;
593            Ok((
594                remain,
595                EndpointAddr {
596                    flags: EndpointAddr::FLAG_NAT,
597                    sequence: None,
598                    load: None,
599                    signature: None,
600                    primary: primary.into(),
601                    agent: Some(agent.into()),
602                },
603            ))
604        }
605        18 => {
606            // IPv6 direct
607            let (remain, addr) = be_socket_addr_v6(input)?;
608            Ok((
609                remain,
610                EndpointAddr {
611                    flags: EndpointAddr::FLAG_FAMILY,
612                    sequence: None,
613                    load: None,
614                    signature: None,
615                    primary: addr.into(),
616                    agent: None,
617                },
618            ))
619        }
620        36 => {
621            // IPv6 NAT
622            let (remain, primary) = be_socket_addr_v6(input)?;
623            let (remain, agent) = be_socket_addr_v6(remain)?;
624            Ok((
625                remain,
626                EndpointAddr {
627                    flags: EndpointAddr::FLAG_FAMILY | EndpointAddr::FLAG_NAT,
628                    sequence: None,
629                    load: None,
630                    signature: None,
631                    primary: primary.into(),
632                    agent: Some(agent.into()),
633                },
634            ))
635        }
636        _ => Err(nom::Err::Error(nom::error::make_error(
637            input,
638            nom::error::ErrorKind::LengthValue,
639        ))),
640    }
641}
642
643fn be_endpoint_signature(input: &[u8], flags: u8) -> IResult<&[u8], Option<EndpointSignature>> {
644    if (flags & EndpointAddr::FLAG_SIGNED) != EndpointAddr::FLAG_SIGNED {
645        if !input.is_empty() {
646            return Err(nom::Err::Error(make_error(input, ErrorKind::Eof)));
647        }
648        return Ok((input, None));
649    }
650
651    let (remain, scheme_u16) = be_u16(input)?;
652    let (remain, sig_len) = be_varint(remain)?;
653    let sig_len = usize::try_from(sig_len.into_inner())
654        .map_err(|_| nom::Err::Error(make_error(remain, ErrorKind::TooLarge)))?;
655    let (remain, sig) = take(sig_len)(remain)?;
656    Ok((
657        remain,
658        Some(EndpointSignature {
659            scheme: scheme_u16,
660            signature: sig.to_vec(),
661        }),
662    ))
663}
664
665pub trait WriteSocketAddr {
666    fn put_socket_addr_v4(&mut self, addr: &SocketAddrV4);
667
668    fn put_socket_addr_v6(&mut self, addr: &SocketAddrV6);
669
670    fn put_socket_addr(&mut self, addr: &SocketAddr) {
671        match addr {
672            SocketAddr::V4(v4) => self.put_socket_addr_v4(v4),
673            SocketAddr::V6(v6) => self.put_socket_addr_v6(v6),
674        }
675    }
676}
677
678impl<T: BufMut> WriteSocketAddr for T {
679    fn put_socket_addr_v4(&mut self, addr: &SocketAddrV4) {
680        self.put_u16(addr.port());
681        self.put_u32(u32::from(*addr.ip()));
682    }
683
684    fn put_socket_addr_v6(&mut self, addr: &SocketAddrV6) {
685        self.put_u16(addr.port());
686        self.put_u128(u128::from(*addr.ip()));
687    }
688}
689
690pub fn be_socket_addr_v4(input: &[u8]) -> IResult<&[u8], SocketAddrV4> {
691    flat_map(be_u16, |port| {
692        map(be_ipv4_addr, move |ip| SocketAddrV4::new(ip, port))
693    })
694    .parse(input)
695}
696
697pub fn be_socket_addr_v6(input: &[u8]) -> IResult<&[u8], SocketAddrV6> {
698    flat_map(be_u16, |port| {
699        map(be_ipv6_addr, move |ip| SocketAddrV6::new(ip, port, 0, 0))
700    })
701    .parse(input)
702}
703
704pub fn be_ipv4_addr(input: &[u8]) -> IResult<&[u8], Ipv4Addr> {
705    map(be_u32, Ipv4Addr::from).parse(input)
706}
707
708pub fn be_ipv6_addr(input: &[u8]) -> IResult<&[u8], Ipv6Addr> {
709    map(be_u128, Ipv6Addr::from).parse(input)
710}
711
712pub fn be_ip_addr(is_v6: bool) -> impl Fn(&[u8]) -> IResult<&[u8], IpAddr> {
713    move |input| match is_v6 {
714        true => map(be_u128, |ip| IpAddr::V6(Ipv6Addr::from(ip))).parse(input),
715        false => map(be_u32, |ip| IpAddr::V4(Ipv4Addr::from(ip))).parse(input),
716    }
717}
718
719impl Display for EndpointAddr {
720    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
721        if let Some(agent_addr) = &self.agent {
722            write!(f, "{}-{agent_addr}", self.primary)
723        } else {
724            write!(f, "{}", self.primary)
725        }
726    }
727}
728
729impl TryFrom<DquicEndpointAddr> for EndpointAddr {
730    type Error = ();
731
732    fn try_from(value: DquicEndpointAddr) -> Result<Self, Self::Error> {
733        match value {
734            DquicEndpointAddr::Direct {
735                addr: SocketAddr::V4(addr),
736            } => Ok(Self::direct_v4(addr)),
737            DquicEndpointAddr::Direct {
738                addr: SocketAddr::V6(addr),
739            } => Ok(Self::direct_v6(addr)),
740            DquicEndpointAddr::Agent {
741                agent: SocketAddr::V4(agent),
742                outer: SocketAddr::V4(outer),
743            } => Ok(Self::nat_v4(outer, agent)),
744            DquicEndpointAddr::Agent {
745                agent: SocketAddr::V6(agent),
746                outer: SocketAddr::V6(outer),
747            } => Ok(Self::nat_v6(outer, agent)),
748            _ => Err(()),
749        }
750    }
751}
752
753impl TryFrom<EndpointAddr> for DquicEndpointAddr {
754    type Error = ();
755
756    fn try_from(value: EndpointAddr) -> Result<Self, Self::Error> {
757        if let Some(agent_addr) = value.agent {
758            match (value.primary, agent_addr) {
759                (SocketAddr::V4(outer), SocketAddr::V4(agent)) => Ok(DquicEndpointAddr::Agent {
760                    outer: SocketAddr::V4(outer),
761                    agent: SocketAddr::V4(agent),
762                }),
763                (SocketAddr::V6(outer), SocketAddr::V6(agent)) => Ok(DquicEndpointAddr::Agent {
764                    outer: SocketAddr::V6(outer),
765                    agent: SocketAddr::V6(agent),
766                }),
767                _ => Err(()),
768            }
769        } else {
770            match value.primary {
771                SocketAddr::V4(addr) => Ok(DquicEndpointAddr::Direct {
772                    addr: SocketAddr::V4(addr),
773                }),
774                SocketAddr::V6(addr) => Ok(DquicEndpointAddr::Direct {
775                    addr: SocketAddr::V6(addr),
776                }),
777            }
778        }
779    }
780}
781
782#[cfg(test)]
783mod tests {
784    use std::{
785        net::{Ipv4Addr, Ipv6Addr},
786        sync::Arc,
787    };
788
789    use bytes::BytesMut;
790    use dhttp_identity::certificate::{
791        CertificateChainKey, CertificateChainKind, CertificateSequence,
792    };
793    use futures::future::BoxFuture;
794    use ring::signature::KeyPair;
795    use rustls::{
796        SignatureScheme,
797        sign::{Signer, SigningKey},
798    };
799
800    use super::*;
801
802    fn chain(sequence: u32, kind: CertificateChainKind) -> CertificateChainKey {
803        CertificateChainKey::new(CertificateSequence::from(sequence), kind)
804    }
805
806    fn ed25519_spki(public_key: &[u8]) -> Vec<u8> {
807        let mut spki = Vec::with_capacity(44);
808        spki.extend_from_slice(&[
809            0x30, 0x2a, 0x30, 0x05, 0x06, 0x03, 0x2b, 0x65, 0x70, 0x03, 0x21, 0x00,
810        ]);
811        spki.extend_from_slice(public_key);
812        spki
813    }
814
815    #[test]
816    fn endpoint_selector_normalizes_missing_sequence_to_primary_zero() {
817        let addr = SocketAddrV4::new(Ipv4Addr::new(10, 0, 0, 1), 5353);
818        let mut endpoint = EndpointAddr::direct_v4(addr);
819        endpoint.set_main(true);
820
821        let selector = endpoint
822            .certificate_chain_key()
823            .expect("missing sequence normalizes to selector");
824
825        assert_eq!(selector, chain(0, CertificateChainKind::Primary));
826    }
827
828    #[test]
829    fn endpoint_selector_normalizes_missing_sequence_to_secondary_zero() {
830        let addr = SocketAddrV4::new(Ipv4Addr::new(10, 0, 0, 2), 5353);
831        let endpoint = EndpointAddr::direct_v4(addr);
832
833        let selector = endpoint
834            .certificate_chain_key()
835            .expect("missing sequence normalizes to selector");
836
837        assert_eq!(selector, chain(0, CertificateChainKind::Secondary));
838    }
839
840    #[test]
841    fn endpoint_selector_sets_primary_and_secondary_chains() {
842        let addr = SocketAddrV4::new(Ipv4Addr::new(10, 0, 0, 3), 5353);
843        let mut endpoint = EndpointAddr::direct_v4(addr);
844
845        endpoint.set_certificate_chain_key(&chain(7, CertificateChainKind::Primary));
846        assert!(endpoint.is_main());
847        assert!(endpoint.is_clustered());
848        assert_eq!(
849            endpoint.certificate_chain_key().unwrap(),
850            chain(7, CertificateChainKind::Primary)
851        );
852
853        endpoint.set_certificate_chain_key(&chain(0, CertificateChainKind::Secondary));
854        assert!(!endpoint.is_main());
855        assert!(!endpoint.is_clustered());
856        assert_eq!(
857            endpoint.certificate_chain_key().unwrap(),
858            chain(0, CertificateChainKind::Secondary)
859        );
860    }
861
862    #[test]
863    fn legacy_endpoint_v4_direct_without_meta() {
864        let port = 5353u16;
865        let ip = Ipv4Addr::new(10, 0, 0, 1);
866        let mut buf = BytesMut::new();
867        buf.extend_from_slice(&port.to_be_bytes());
868        buf.extend_from_slice(&u32::from(ip).to_be_bytes());
869        let (remain, decoded) = be_endpoint_addr_compat(&buf, 6).unwrap();
870        assert!(remain.is_empty());
871        assert_eq!(
872            decoded,
873            EndpointAddr::direct_v4(SocketAddrV4::new(ip, port))
874        );
875    }
876
877    #[test]
878    fn legacy_endpoint_v4_nat_without_meta() {
879        let outer = SocketAddrV4::new(Ipv4Addr::new(10, 0, 0, 1), 1000);
880        let agent = SocketAddrV4::new(Ipv4Addr::new(10, 0, 0, 2), 2000);
881        let mut buf = BytesMut::new();
882        buf.extend_from_slice(&outer.port().to_be_bytes());
883        buf.extend_from_slice(&u32::from(*outer.ip()).to_be_bytes());
884        buf.extend_from_slice(&agent.port().to_be_bytes());
885        buf.extend_from_slice(&u32::from(*agent.ip()).to_be_bytes());
886        let (remain, decoded) = be_endpoint_addr_compat(&buf, 12).unwrap();
887        assert!(remain.is_empty());
888        assert_eq!(decoded, EndpointAddr::nat_v4(outer, agent));
889    }
890
891    #[test]
892    fn flag_bit_ops_work() {
893        let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 5353);
894        let mut ep = EndpointAddr {
895            flags: 0b0011_0000,
896            sequence: None,
897            load: None,
898            signature: None,
899            primary: addr.into(),
900            agent: None,
901        };
902
903        assert!(!ep.is_main());
904        assert!(!ep.is_signed());
905
906        ep.set_main(true);
907        assert!(ep.is_main());
908        assert_eq!(ep.flags, 0b0111_0000);
909
910        ep.set_signed(true);
911        assert!(ep.is_signed());
912        assert_eq!(ep.flags, 0b0111_0001);
913
914        ep.set_main(false);
915        assert!(!ep.is_main());
916        assert!(ep.is_signed());
917        assert_eq!(ep.flags, 0b0011_0001);
918
919        ep.set_signed(false);
920        assert!(!ep.is_signed());
921        assert_eq!(ep.flags, 0b0011_0000);
922    }
923
924    #[test]
925    fn varint_roundtrip_and_len() {
926        fn roundtrip(v: u64) {
927            let v = VarInt::from_u64(v).unwrap();
928            let mut buf = BytesMut::new();
929            buf.put_varint(v);
930            assert_eq!(buf.len(), v.encoding_size());
931            let (remain, decoded) = be_varint(&buf).unwrap();
932            assert!(remain.is_empty());
933            assert_eq!(decoded, v);
934        }
935
936        for v in [
937            0u64,
938            1,
939            63,
940            64,
941            16383,
942            16384,
943            (1 << 30) - 1,
944            1 << 30,
945            (1 << 62) - 1,
946        ] {
947            roundtrip(v);
948        }
949    }
950
951    #[test]
952    fn varint_rejects_overflow_and_incomplete() {
953        assert!(VarInt::from_u64((1 << 62) + 1).is_err());
954
955        let incomplete = [0b01_000000u8];
956        match be_varint(&incomplete) {
957            Err(nom::Err::Incomplete(_)) => {}
958            other => panic!("expected Incomplete, got {other:?}"),
959        }
960    }
961
962    #[test]
963    fn endpoint_encode_decode_roundtrip() {
964        let v4_outer = SocketAddrV4::new(Ipv4Addr::new(10, 0, 0, 1), 1000);
965        let v4_agent = SocketAddrV4::new(Ipv4Addr::new(10, 0, 0, 2), 2000);
966        let v6_outer = SocketAddrV6::new(Ipv6Addr::LOCALHOST, 3000, 0, 0);
967        let v6_agent = SocketAddrV6::new(Ipv6Addr::LOCALHOST, 4000, 0, 0);
968
969        let mut with_load = EndpointAddr::direct_v4(v4_outer);
970        with_load.set_load(Some(0.42_f32));
971
972        let cases = vec![
973            // IPv4 direct, MAIN + CLUSTERED flags
974            EndpointAddr {
975                flags: EndpointAddr::FLAG_MAIN | EndpointAddr::FLAG_CLUSTERED,
976                sequence: Some(VarInt::from_u32(0)),
977                load: None,
978                signature: None,
979                primary: v4_outer.into(),
980                agent: None,
981            },
982            // IPv4 NAT, CLUSTERED flag
983            EndpointAddr {
984                flags: EndpointAddr::FLAG_NAT | EndpointAddr::FLAG_CLUSTERED,
985                sequence: Some(VarInt::from_u32(127)),
986                load: None,
987                signature: None,
988                primary: v4_outer.into(),
989                agent: Some(v4_agent.into()),
990            },
991            // IPv6 direct, MAIN + CLUSTERED flags
992            EndpointAddr {
993                flags: EndpointAddr::FLAG_FAMILY
994                    | EndpointAddr::FLAG_MAIN
995                    | EndpointAddr::FLAG_CLUSTERED,
996                sequence: Some(VarInt::from_u32(128)),
997                load: None,
998                signature: None,
999                primary: v6_outer.into(),
1000                agent: None,
1001            },
1002            // IPv6 NAT, CLUSTERED flag
1003            EndpointAddr {
1004                flags: EndpointAddr::FLAG_FAMILY
1005                    | EndpointAddr::FLAG_NAT
1006                    | EndpointAddr::FLAG_CLUSTERED,
1007                sequence: Some(VarInt::from_u64((1 << 62) - 1).unwrap()),
1008                load: None,
1009                signature: None,
1010                primary: v6_outer.into(),
1011                agent: Some(v6_agent.into()),
1012            },
1013            // IPv4 direct with LOAD
1014            with_load,
1015        ];
1016
1017        for ep in cases {
1018            let mut buf = BytesMut::new();
1019            buf.put_endpoint_addr(&ep);
1020            assert_eq!(buf.len(), ep.encpding_size());
1021
1022            let (remain, decoded) = be_endpoint_addr(&buf).unwrap();
1023            assert!(remain.is_empty());
1024            assert_eq!(decoded, ep);
1025        }
1026    }
1027
1028    #[test]
1029    fn signed_endpoint_accepts_scheme_inclusive_signature() {
1030        let addr = SocketAddrV4::new(Ipv4Addr::new(10, 10, 0, 7), 20004);
1031        let scheme = u16::from(SignatureScheme::ED25519);
1032        let signature = vec![0xaa; 64];
1033        let sig_len = VarInt::try_from(signature.len() as u64).unwrap();
1034
1035        let mut buf = BytesMut::new();
1036        buf.put_u8(EndpointAddr::FLAG_SIGNED);
1037        buf.put_socket_addr_v4(&addr);
1038        buf.put_u16(scheme);
1039        buf.put_varint(sig_len);
1040        buf.extend_from_slice(&signature);
1041
1042        let (remain, decoded) = be_endpoint_addr(&buf).unwrap();
1043
1044        assert!(remain.is_empty());
1045        assert!(decoded.is_signed());
1046        assert_eq!(decoded.addr(), SocketAddr::V4(addr));
1047        assert_eq!(decoded.signature().unwrap().signature, signature);
1048    }
1049
1050    #[test]
1051    fn signed_endpoint_rejects_signature_without_scheme() {
1052        let addr = SocketAddrV4::new(Ipv4Addr::new(10, 10, 0, 7), 20004);
1053        let signature = vec![0xaa; 64];
1054        let sig_len = VarInt::try_from(signature.len() as u64).unwrap();
1055
1056        let mut buf = BytesMut::new();
1057        buf.put_u8(EndpointAddr::FLAG_SIGNED);
1058        buf.put_socket_addr_v4(&addr);
1059        buf.put_varint(sig_len);
1060        buf.extend_from_slice(&signature);
1061
1062        assert!(be_endpoint_addr(&buf).is_err());
1063    }
1064
1065    #[test]
1066    fn signed_endpoint_writes_actual_scheme_before_signature_length() {
1067        #[derive(Debug)]
1068        struct Ed25519Key {
1069            keypair: Arc<ring::signature::Ed25519KeyPair>,
1070            spki: Vec<u8>,
1071        }
1072
1073        #[derive(Debug)]
1074        struct Ed25519Signer(Arc<ring::signature::Ed25519KeyPair>);
1075
1076        impl Signer for Ed25519Signer {
1077            fn sign(&self, message: &[u8]) -> Result<Vec<u8>, rustls::Error> {
1078                Ok(self.0.sign(message).as_ref().to_vec())
1079            }
1080
1081            fn scheme(&self) -> SignatureScheme {
1082                SignatureScheme::ED25519
1083            }
1084        }
1085
1086        impl SigningKey for Ed25519Key {
1087            fn choose_scheme(&self, offered: &[SignatureScheme]) -> Option<Box<dyn Signer>> {
1088                offered
1089                    .contains(&SignatureScheme::ED25519)
1090                    .then(|| Box::new(Ed25519Signer(self.keypair.clone())) as Box<dyn Signer>)
1091            }
1092
1093            fn algorithm(&self) -> rustls::SignatureAlgorithm {
1094                rustls::SignatureAlgorithm::ED25519
1095            }
1096        }
1097
1098        impl dhttp_identity::identity::LocalAuthority for Ed25519Key {
1099            fn name(&self) -> &str {
1100                "authority.example"
1101            }
1102
1103            fn cert_chain(&self) -> &[rustls::pki_types::CertificateDer<'static>] {
1104                &[]
1105            }
1106
1107            fn public_key(&self) -> SubjectPublicKeyInfoDer<'_> {
1108                SubjectPublicKeyInfoDer::from(self.spki.as_slice())
1109            }
1110
1111            fn sign(
1112                &self,
1113                data: &[u8],
1114            ) -> BoxFuture<'_, Result<Vec<u8>, dhttp_identity::identity::SignError>> {
1115                let result = dhttp_identity::identity::sign_with_key(self, data);
1116                Box::pin(std::future::ready(result))
1117            }
1118        }
1119
1120        let rng = ring::rand::SystemRandom::new();
1121        let pkcs8 = ring::signature::Ed25519KeyPair::generate_pkcs8(&rng).unwrap();
1122        let keypair =
1123            Arc::new(ring::signature::Ed25519KeyPair::from_pkcs8(pkcs8.as_ref()).unwrap());
1124        let spki = ed25519_spki(keypair.public_key().as_ref());
1125        let key = Ed25519Key { keypair, spki };
1126
1127        let mut ep = EndpointAddr::direct_v4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 5353));
1128        futures::executor::block_on(ep.sign_with_authority(&key)).unwrap();
1129
1130        let mut buf = BytesMut::new();
1131        buf.put_endpoint_addr(&ep);
1132
1133        let scheme_offset = 1 + 2 + 4;
1134        let encoded_scheme = u16::from_be_bytes([buf[scheme_offset], buf[scheme_offset + 1]]);
1135        assert_eq!(encoded_scheme, u16::from(SignatureScheme::ED25519));
1136    }
1137
1138    #[test]
1139    fn endpoint_signature_roundtrip_and_verify() {
1140        #[derive(Debug)]
1141        struct Ed25519Key {
1142            keypair: Arc<ring::signature::Ed25519KeyPair>,
1143            spki: Vec<u8>,
1144        }
1145
1146        #[derive(Debug)]
1147        struct Ed25519Signer(Arc<ring::signature::Ed25519KeyPair>);
1148
1149        impl Signer for Ed25519Signer {
1150            fn sign(&self, message: &[u8]) -> Result<Vec<u8>, rustls::Error> {
1151                Ok(self.0.sign(message).as_ref().to_vec())
1152            }
1153
1154            fn scheme(&self) -> SignatureScheme {
1155                SignatureScheme::ED25519
1156            }
1157        }
1158
1159        impl SigningKey for Ed25519Key {
1160            fn choose_scheme(&self, offered: &[SignatureScheme]) -> Option<Box<dyn Signer>> {
1161                offered
1162                    .contains(&SignatureScheme::ED25519)
1163                    .then(|| Box::new(Ed25519Signer(self.keypair.clone())) as Box<dyn Signer>)
1164            }
1165
1166            fn algorithm(&self) -> rustls::SignatureAlgorithm {
1167                rustls::SignatureAlgorithm::ED25519
1168            }
1169        }
1170
1171        impl dhttp_identity::identity::LocalAuthority for Ed25519Key {
1172            fn name(&self) -> &str {
1173                "authority.example"
1174            }
1175
1176            fn cert_chain(&self) -> &[rustls::pki_types::CertificateDer<'static>] {
1177                &[]
1178            }
1179
1180            fn public_key(&self) -> SubjectPublicKeyInfoDer<'_> {
1181                SubjectPublicKeyInfoDer::from(self.spki.as_slice())
1182            }
1183
1184            fn sign(
1185                &self,
1186                data: &[u8],
1187            ) -> BoxFuture<'_, Result<Vec<u8>, dhttp_identity::identity::SignError>> {
1188                let result = dhttp_identity::identity::sign_with_key(self, data);
1189                Box::pin(std::future::ready(result))
1190            }
1191        }
1192
1193        let rng = ring::rand::SystemRandom::new();
1194        let pkcs8 = ring::signature::Ed25519KeyPair::generate_pkcs8(&rng).unwrap();
1195        let keypair =
1196            Arc::new(ring::signature::Ed25519KeyPair::from_pkcs8(pkcs8.as_ref()).unwrap());
1197        let spki = ed25519_spki(keypair.public_key().as_ref());
1198        let key = Ed25519Key {
1199            keypair: keypair.clone(),
1200            spki: spki.clone(),
1201        };
1202
1203        let addr = SocketAddrV4::new(Ipv4Addr::new(10, 0, 0, 1), 5353);
1204        let mut ep = EndpointAddr::direct_v4(addr);
1205        ep.set_main(true);
1206        futures::executor::block_on(ep.sign_with_authority(&key)).unwrap();
1207
1208        let mut buf = BytesMut::new();
1209        buf.put_endpoint_addr(&ep);
1210        assert_eq!(buf.len(), ep.encpding_size());
1211
1212        let (remain, decoded) = be_endpoint_addr(&buf).unwrap();
1213        assert!(remain.is_empty());
1214        assert!(decoded.is_signed());
1215        assert!(decoded.signature().is_some());
1216        assert!(
1217            decoded
1218                .verify_signature(SubjectPublicKeyInfoDer::from(spki.as_slice()))
1219                .unwrap()
1220        );
1221
1222        let mut tampered = decoded.clone();
1223        tampered.set_main(false);
1224        assert!(
1225            !tampered
1226                .verify_signature(SubjectPublicKeyInfoDer::from(spki.as_slice()))
1227                .unwrap()
1228        );
1229    }
1230
1231    #[test]
1232    fn sign_with_authority_stores_canonical_signature() {
1233        #[derive(Debug)]
1234        struct StaticAuthority {
1235            spki: Vec<u8>,
1236        }
1237
1238        impl dhttp_identity::identity::LocalAuthority for StaticAuthority {
1239            fn name(&self) -> &str {
1240                "authority.example"
1241            }
1242
1243            fn cert_chain(&self) -> &[rustls::pki_types::CertificateDer<'static>] {
1244                &[]
1245            }
1246
1247            fn public_key(&self) -> SubjectPublicKeyInfoDer<'_> {
1248                SubjectPublicKeyInfoDer::from(self.spki.as_slice())
1249            }
1250
1251            fn sign(
1252                &self,
1253                _data: &[u8],
1254            ) -> BoxFuture<'_, Result<Vec<u8>, dhttp_identity::identity::SignError>> {
1255                Box::pin(std::future::ready(Ok(vec![1, 2, 3])))
1256            }
1257        }
1258
1259        let mut ep = EndpointAddr::direct_v4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 5353));
1260        let authority = StaticAuthority {
1261            spki: ed25519_spki(&[0; 32]),
1262        };
1263        futures::executor::block_on(ep.sign_with_authority(&authority)).unwrap();
1264
1265        let signature = ep.signature().unwrap();
1266        assert_eq!(signature.scheme, u16::from(SignatureScheme::ED25519));
1267        assert_eq!(signature.signature, vec![1, 2, 3]);
1268    }
1269
1270    #[test]
1271    fn optional_fields_flags_follow_values() {
1272        let addr = SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 5353);
1273        let mut ep = EndpointAddr::direct_v4(addr);
1274
1275        assert!(!ep.is_load());
1276
1277        ep.set_load(Some(0.5_f32));
1278
1279        assert!(ep.is_load());
1280        assert_eq!(ep.load(), Some(0.5_f32));
1281
1282        ep.set_load(None);
1283
1284        assert!(!ep.is_load());
1285        assert_eq!(ep.load(), None);
1286    }
1287}