Fix crash on oversized nts cookies.

Author: davidv1992

ntp-proto/src/nts_record.rs

@@ -675,6 +675,7 @@ pub enum KeyExchangeError {
     NoValidAlgorithm,
     InvalidFixedKeyLength,
     NoCookies,
+    CookiesTooBig,
     Io(std::io::Error),
     Tls(tls_utils::Error),
     Certificate(tls_utils::Error),
@@ -704,6 +705,7 @@ impl Display for KeyExchangeError {
                 "The length of a fixed key does not match the algorithm used"
             ),
             Self::NoCookies => write!(f, "Missing cookies"),
+            Self::CookiesTooBig => write!(f, "Server returned cookies that are too large"),
             Self::Io(e) => write!(f, "{e}"),
             Self::Tls(e) => write!(f, "{e}"),
             Self::Certificate(e) => write!(f, "{e}"),
@@ -756,6 +758,7 @@ impl KeyExchangeError {
             | NoValidAlgorithm
             | InvalidFixedKeyLength
             | NoCookies
+            | CookiesTooBig
             | Tls(_)
             | Certificate(_)
             | DnsName(_)
@@ -964,6 +967,9 @@ pub struct KeyExchangeResultDecoder {
 }
 
 impl KeyExchangeResultDecoder {
+    // Chosen such that we can get new cookies if the need arises. In practice, the cookie should never be this big.
+    const MAX_COOKIE_SIZE: usize = 350;
+
     pub fn step_with_slice(
         mut self,
         bytes: &[u8],
@@ -1027,6 +1033,9 @@ impl KeyExchangeResultDecoder {
                 Continue(state)
             }
             NewCookie { cookie_data } => {
+                if cookie_data.len() > Self::MAX_COOKIE_SIZE {
+                    return ControlFlow::Break(Err(KeyExchangeError::CookiesTooBig));
+                }
                 state.cookies.store(cookie_data);
                 Continue(state)
             }
@@ -2507,6 +2516,22 @@ mod test {
         ]
     }
 
+    fn nts_oversized_cookie() -> [NtsRecord; 4] {
+        [
+            NtsRecord::NextProtocol {
+                protocol_ids: vec![0],
+            },
+            NtsRecord::AeadAlgorithm {
+                critical: false,
+                algorithm_ids: vec![15],
+            },
+            NtsRecord::NewCookie {
+                cookie_data: vec![0; 2048],
+            },
+            NtsRecord::EndOfMessage,
+        ]
+    }
+
     #[test]
     fn test_nts_time_nl_response() {
         let state = client_decode_records(nts_time_nl_records().as_slice()).unwrap();
@@ -2516,6 +2541,12 @@ mod test {
         assert_eq!(state.cookies.gap(), 0);
     }
 
+    #[test]
+    fn reject_oversized_cookie() {
+        let result = client_decode_records(nts_oversized_cookie().as_slice());
+        assert!(matches!(result, Err(KeyExchangeError::CookiesTooBig)));
+    }
+
     #[test]
     fn test_decode_nts_time_nl_response() {
         let mut decoder = NtsRecord::decoder();

ntp-proto/src/source.rs

@@ -571,6 +571,11 @@ impl<Controller: SourceController<MeasurementDelay = NtpDuration>> NtpSource<Con
                     .cookies
                     .gap()
                     .min(((self.buffer.len() - 300) / cookie.len()).min(u8::MAX as usize) as u8);
+                // Defence in depth, ensure we can get at least 1 new cookie.
+                if new_cookies == 0 {
+                    warn!("NTS Cookie too large, resetting source. This may be a problem with the source");
+                    return actions![NtpSourceAction::Reset];
+                }
                 match self.protocol_version {
                     ProtocolVersion::V4 => {
                         NtpPacket::nts_poll_message(&cookie, new_cookies, poll_interval)
@@ -861,7 +866,11 @@ impl<Controller: SourceController<MeasurementDelay = NtpDuration>> NtpSource<Con
 
 #[cfg(test)]
 mod test {
-    use crate::{packet::NoCipher, time_types::PollIntervalLimits, NtpClock};
+    use crate::{
+        packet::{AesSivCmac256, NoCipher},
+        time_types::PollIntervalLimits,
+        NtpClock,
+    };
 
     use super::*;
     #[cfg(feature = "ntpv5")]
@@ -1089,6 +1098,23 @@ mod test {
         assert!(source.current_poll_interval() >= source.controller.0);
     }
 
+    #[test]
+    fn test_oversize_cookie_doesnt_crash() {
+        let mut source = NtpSource::test_ntp_source(NoopController);
+        let mut ntsdata = SourceNtsData {
+            cookies: CookieStash::default(),
+            c2s: Box::new(AesSivCmac256::new([0; 32].into())),
+            s2c: Box::new(AesSivCmac256::new([0; 32].into())),
+        };
+        ntsdata.cookies.store(vec![0; 2048]);
+        ntsdata.cookies.store(vec![0; 2048]);
+        source.nts = Some(Box::new(ntsdata));
+        let actions = source.handle_timer();
+        for action in actions {
+            assert!(matches!(action, NtpSourceAction::Reset))
+        }
+    }
+
     #[test]
     fn test_handle_incoming() {
         let base = NtpInstant::now();