Implement Server Reference ID Bloom FIlter

ntp-proto/src/lib.rs

@@ -69,6 +69,11 @@ mod exports {
         KeyExchangeClient, KeyExchangeError, KeyExchangeResult, KeyExchangeServer, NtsRecord,
         NtsRecordDecoder, WriteError,
     };
+
+    #[cfg(feature = "ntpv5")]
+    pub mod v5 {
+        pub use crate::packet::v5::server_reference_id::{BloomFilter, ServerId};
+    }
 }
 
 #[cfg(feature = "__internal-api")]

ntp-proto/src/packet/extension_fields.rs

@@ -5,6 +5,9 @@ use std::{
 
 use crate::keyset::DecodedServerCookie;
 
+#[cfg(feature = "ntpv5")]
+use crate::packet::v5::extension_fields::{ReferenceIdRequest, ReferenceIdResponse};
+
 use super::{crypto::EncryptResult, error::ParsingError, Cipher, CipherProvider, Mac};
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -20,6 +23,10 @@ enum ExtensionFieldTypeId {
     DraftIdentification,
     #[cfg(feature = "ntpv5")]
     Padding,
+    #[cfg(feature = "ntpv5")]
+    ReferenceIdRequest,
+    #[cfg(feature = "ntpv5")]
+    ReferenceIdResponse,
 }
 
 impl ExtensionFieldTypeId {
@@ -33,6 +40,10 @@ impl ExtensionFieldTypeId {
             0xF5FF => Self::DraftIdentification,
             #[cfg(feature = "ntpv5")]
             0xF501 => Self::Padding,
+            #[cfg(feature = "ntpv5")]
+            0xF503 => Self::ReferenceIdRequest,
+            #[cfg(feature = "ntpv5")]
+            0xF504 => Self::ReferenceIdResponse,
             _ => Self::Unknown { type_id },
         }
     }
@@ -47,6 +58,10 @@ impl ExtensionFieldTypeId {
             ExtensionFieldTypeId::DraftIdentification => 0xF5FF,
             #[cfg(feature = "ntpv5")]
             ExtensionFieldTypeId::Padding => 0xF501,
+            #[cfg(feature = "ntpv5")]
+            ExtensionFieldTypeId::ReferenceIdRequest => 0xF503,
+            #[cfg(feature = "ntpv5")]
+            ExtensionFieldTypeId::ReferenceIdResponse => 0xF504,
             ExtensionFieldTypeId::Unknown { type_id } => type_id,
         }
     }
@@ -64,6 +79,10 @@ pub enum ExtensionField<'a> {
     DraftIdentification(Cow<'a, str>),
     #[cfg(feature = "ntpv5")]
     Padding(usize),
+    #[cfg(feature = "ntpv5")]
+    ReferenceIdRequest(super::v5::extension_fields::ReferenceIdRequest),
+    #[cfg(feature = "ntpv5")]
+    ReferenceIdResponse(super::v5::extension_fields::ReferenceIdResponse<'a>),
     Unknown {
         type_id: u16,
         data: Cow<'a, [u8]>,
@@ -88,6 +107,10 @@ impl<'a> std::fmt::Debug for ExtensionField<'a> {
             }
             #[cfg(feature = "ntpv5")]
             Self::Padding(len) => f.debug_struct("Padding").field("length", &len).finish(),
+            #[cfg(feature = "ntpv5")]
+            Self::ReferenceIdRequest(r) => f.debug_tuple("ReferenceIdRequest").field(r).finish(),
+            #[cfg(feature = "ntpv5")]
+            Self::ReferenceIdResponse(r) => f.debug_tuple("ReferenceIdResponse").field(r).finish(),
             Self::Unknown {
                 type_id: typeid,
                 data,
@@ -127,6 +150,10 @@ impl<'a> ExtensionField<'a> {
             DraftIdentification(data) => DraftIdentification(Cow::Owned(data.into_owned())),
             #[cfg(feature = "ntpv5")]
             Padding(len) => Padding(len),
+            #[cfg(feature = "ntpv5")]
+            ReferenceIdRequest(req) => ReferenceIdRequest(req),
+            #[cfg(feature = "ntpv5")]
+            ReferenceIdResponse(res) => ReferenceIdResponse(res.into_owned()),
         }
     }
 
@@ -156,6 +183,10 @@ impl<'a> ExtensionField<'a> {
             }
             #[cfg(feature = "ntpv5")]
             Padding(len) => Self::encode_padding_field(w, *len, minimum_size, version),
+            #[cfg(feature = "ntpv5")]
+            ReferenceIdRequest(req) => req.serialize(w),
+            #[cfg(feature = "ntpv5")]
+            ReferenceIdResponse(res) => res.serialize(w),
         }
     }
 
@@ -518,6 +549,10 @@ impl<'a> ExtensionField<'a> {
             TypeId::DraftIdentification => {
                 EF::decode_draft_identification(message, extension_header_version)
             }
+            #[cfg(feature = "ntpv5")]
+            TypeId::ReferenceIdRequest => Ok(ReferenceIdRequest::decode(message)?.into()),
+            #[cfg(feature = "ntpv5")]
+            TypeId::ReferenceIdResponse => Ok(ReferenceIdResponse::decode(message).into()),
             type_id => EF::decode_unknown(type_id.to_type_id(), message),
         }
     }

ntp-proto/src/packet/mod.rs

@@ -19,7 +19,7 @@ mod extension_fields;
 mod mac;
 
 #[cfg(feature = "ntpv5")]
-mod v5;
+pub mod v5;
 
 pub use crypto::{
     AesSivCmac256, AesSivCmac512, Cipher, CipherHolder, CipherProvider, DecryptError,
@@ -118,15 +118,15 @@ pub struct NtpPacket<'a> {
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
-enum NtpHeader {
+pub enum NtpHeader {
     V3(NtpHeaderV3V4),
     V4(NtpHeaderV3V4),
     #[cfg(feature = "ntpv5")]
     V5(v5::NtpHeaderV5),
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
-struct NtpHeaderV3V4 {
+pub struct NtpHeaderV3V4 {
     leap: NtpLeapIndicator,
     mode: NtpAssociationMode,
     stratum: u8,
@@ -652,7 +652,14 @@ impl<'a> NtpPacket<'a> {
                         .untrusted
                         .into_iter()
                         .chain(input.efdata.authenticated)
-                        .filter(|ef| matches!(ef, ExtensionField::UniqueIdentifier(_)))
+                        .filter_map(|ef| match ef {
+                            uid @ ExtensionField::UniqueIdentifier(_) => Some(uid),
+                            ExtensionField::ReferenceIdRequest(req) => {
+                                let response = req.to_response(&system.bloom_filter)?;
+                                Some(ExtensionField::ReferenceIdResponse(response).into_owned())
+                            }
+                            _ => None,
+                        })
                         .chain(std::iter::once(ExtensionField::DraftIdentification(
                             Cow::Borrowed(v5::DRAFT_VERSION),
                         )))
@@ -871,6 +878,10 @@ impl<'a> NtpPacket<'a> {
         }
     }
 
+    pub fn header(&self) -> NtpHeader {
+        self.header
+    }
+
     pub fn leap(&self) -> NtpLeapIndicator {
         match self.header {
             NtpHeader::V3(header) => header.leap,
@@ -1039,6 +1050,14 @@ impl<'a> NtpPacket<'a> {
             }
         }
     }
+
+    pub fn untrusted_extension_fields(&self) -> impl Iterator<Item = &ExtensionField> {
+        self.efdata.untrusted.iter()
+    }
+
+    pub fn push_untrusted(&mut self, ef: ExtensionField<'static>) {
+        self.efdata.untrusted.push(ef);
+    }
 }
 
 // Returns whether all uid extension fields found match the given uid, or

ntp-proto/src/packet/v5/extension_fields.rs

@@ -1,3 +1,10 @@
+use crate::packet::error::ParsingError;
+use crate::packet::v5::server_reference_id::BloomFilter;
+use crate::packet::ExtensionField;
+use std::borrow::Cow;
+use std::convert::Infallible;
+use std::io::Write;
+
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub enum Type {
     DraftIdentification,
@@ -64,6 +71,137 @@ impl Type {
     }
 }
 
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub struct ReferenceIdRequest {
+    payload_len: u16,
+    offset: u16,
+}
+
+impl ReferenceIdRequest {
+    pub const fn new(payload_len: u16, offset: u16) -> Option<Self> {
+        if payload_len % 4 != 0 {
+            return None;
+        }
+
+        if payload_len + offset > 512 {
+            return None;
+        }
+
+        Some(Self {
+            payload_len,
+            offset,
+        })
+    }
+
+    pub fn to_response(self, filter: &BloomFilter) -> Option<ReferenceIdResponse> {
+        let offset = usize::from(self.offset);
+        let payload_len = usize::from(self.payload_len);
+
+        let bytes = filter.as_bytes().get(offset..)?.get(..payload_len)?.into();
+
+        Some(ReferenceIdResponse { bytes })
+    }
+
+    pub fn serialize(&self, mut writer: impl Write) -> std::io::Result<()> {
+        let payload_len = self.payload_len;
+        let ef_len: u16 = payload_len + 4;
+
+        writer.write_all(&Type::ReferenceIdRequest.to_bits().to_be_bytes())?;
+        writer.write_all(&ef_len.to_be_bytes())?;
+        writer.write_all(&self.offset.to_be_bytes())?;
+        writer.write_all(&[0; 2])?;
+
+        let words = payload_len / 4;
+        assert_eq!(payload_len % 4, 0);
+
+        for _ in 1..words {
+            writer.write_all(&[0; 4])?;
+        }
+
+        Ok(())
+    }
+
+    pub fn decode(msg: &[u8]) -> Result<Self, ParsingError<Infallible>> {
+        let payload_len =
+            u16::try_from(msg.len()).expect("NTP fields can not be longer than u16::MAX");
+        let offset_bytes: [u8; 2] = msg[0..2].try_into().unwrap();
+
+        Ok(Self {
+            payload_len,
+            offset: u16::from_be_bytes(offset_bytes),
+        })
+    }
+
+    pub const fn offset(&self) -> u16 {
+        self.offset
+    }
+
+    pub const fn payload_len(&self) -> u16 {
+        self.payload_len
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct ReferenceIdResponse<'a> {
+    bytes: Cow<'a, [u8]>,
+}
+
+impl<'a> ReferenceIdResponse<'a> {
+    pub const fn new(bytes: &'a [u8]) -> Option<Self> {
+        if bytes.len() % 4 != 0 {
+            return None;
+        }
+
+        if bytes.len() > 512 {
+            return None;
+        }
+
+        Some(Self {
+            bytes: Cow::Borrowed(bytes),
+        })
+    }
+
+    pub fn into_owned(self) -> ReferenceIdResponse<'static> {
+        ReferenceIdResponse {
+            bytes: Cow::Owned(self.bytes.into_owned()),
+        }
+    }
+
+    pub fn serialize(&self, mut writer: impl Write) -> std::io::Result<()> {
+        let len: u16 = self.bytes.len().try_into().unwrap();
+        let len = len + 4; // Add room for type and length
+        assert_eq!(len % 4, 0);
+
+        writer.write_all(&Type::ReferenceIdResponse.to_bits().to_be_bytes())?;
+        writer.write_all(&len.to_be_bytes())?;
+        writer.write_all(self.bytes.as_ref())?;
+
+        Ok(())
+    }
+
+    pub const fn decode(bytes: &'a [u8]) -> Self {
+        Self {
+            bytes: Cow::Borrowed(bytes),
+        }
+    }
+
+    pub fn bytes(&self) -> &[u8] {
+        &self.bytes
+    }
+}
+
+impl From<ReferenceIdRequest> for ExtensionField<'static> {
+    fn from(value: ReferenceIdRequest) -> Self {
+        Self::ReferenceIdRequest(value)
+    }
+}
+
+impl<'a> From<ReferenceIdResponse<'a>> for ExtensionField<'a> {
+    fn from(value: ReferenceIdResponse<'a>) -> Self {
+        Self::ReferenceIdResponse(value)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

ntp-proto/src/packet/v5/mod.rs

@@ -5,6 +5,7 @@ use rand::random;
 mod error;
 #[allow(dead_code)]
 pub mod extension_fields;
+pub mod server_reference_id;
 
 use crate::packet::error::ParsingError;
 pub use error::V5Error;