From b5f1d8d438b809e1c57f2983d1a4000e0ca88b15 Mon Sep 17 00:00:00 2001
From: Nils Andreas Svee <me@lochnair.net>
Date: Thu, 7 May 2026 21:02:45 +0700
Subject: [PATCH] fix(client): check probe scheduling arithmetic

---
 crates/irtt-client/src/client.rs              | 37 ++++++++++++---
 .../src/client/tests/probes_replies.rs        | 45 +++++++++++++++++++
 2 files changed, 77 insertions(+), 5 deletions(-)

diff --git a/crates/irtt-client/src/client.rs b/crates/irtt-client/src/client.rs
index 187f13f..356721f 100644
--- a/crates/irtt-client/src/client.rs
+++ b/crates/irtt-client/src/client.rs
@@ -231,13 +231,22 @@ impl Client {
             logical_seq,
             wire_seq,
             sent_at,
-            timeout_at: sent_at.mono + self.config.probe_timeout,
+            timeout_at: sent_at
+                .mono
+                .checked_add(self.config.probe_timeout)
+                .ok_or(ClientError::DurationOverflow)?,
         };
         session.pending.insert(pending)?;
 
         session.next_wire_seq = session.next_wire_seq.wrapping_add(1);
-        session.next_logical_seq += 1;
-        session.packets_sent += 1;
+        session.next_logical_seq = session
+            .next_logical_seq
+            .checked_add(1)
+            .ok_or(ClientError::DurationOverflow)?;
+        session.packets_sent = session
+            .packets_sent
+            .checked_add(1)
+            .ok_or(ClientError::DurationOverflow)?;
 
         let negotiated = self
             .negotiated
@@ -249,7 +258,7 @@ impl Client {
             .as_mut()
             .expect("session must exist when phase is Open");
         session.next_send_at =
-            session.start_mono + Duration::from_nanos(interval_ns * session.packets_sent);
+            next_probe_deadline(session.start_mono, interval_ns, session.packets_sent)?;
 
         if let Some(end) = session.end_mono {
             if session.next_send_at >= end {
@@ -564,7 +573,12 @@ impl Client {
 
         let duration_ns = negotiated.params.duration_ns;
         let end_mono = if duration_ns > 0 {
-            Some(now.mono + Duration::from_nanos(duration_ns as u64))
+            let duration = Duration::from_nanos(duration_ns as u64);
+            Some(
+                now.mono
+                    .checked_add(duration)
+                    .ok_or(ClientError::DurationOverflow)?,
+            )
         } else {
             None
         };
@@ -647,6 +661,19 @@ fn update_highest_received(session: &mut ActiveSession, wire_seq: u32) {
     }));
 }
 
+fn next_probe_deadline(
+    start: Instant,
+    interval_ns: u64,
+    packets_sent: u64,
+) -> Result<Instant, ClientError> {
+    let offset_ns = interval_ns
+        .checked_mul(packets_sent)
+        .ok_or(ClientError::DurationOverflow)?;
+    start
+        .checked_add(Duration::from_nanos(offset_ns))
+        .ok_or(ClientError::DurationOverflow)
+}
+
 fn sequence_is_after(candidate: u32, current: u32) -> bool {
     candidate != current && candidate.wrapping_sub(current) < (1 << 31)
 }
diff --git a/crates/irtt-client/src/client/tests/probes_replies.rs b/crates/irtt-client/src/client/tests/probes_replies.rs
index c478ff9..ebb6b6a 100644
--- a/crates/irtt-client/src/client/tests/probes_replies.rs
+++ b/crates/irtt-client/src/client/tests/probes_replies.rs
@@ -266,6 +266,51 @@ fn next_send_deadline_advances_by_interval() {
     server.join();
 }
 
+#[test]
+fn next_probe_deadline_reports_overflow() {
+    let start = Instant::now();
+    assert_eq!(
+        next_probe_deadline(start, 1_000_000_000, 2).unwrap(),
+        start + Duration::from_secs(2)
+    );
+    assert!(matches!(
+        next_probe_deadline(start, u64::MAX, 2),
+        Err(ClientError::DurationOverflow)
+    ));
+}
+
+#[test]
+fn send_probe_reports_schedule_overflow() {
+    let params = Params {
+        protocol_version: 1,
+        duration_ns: 0,
+        interval_ns: 1_000_000_000,
+        received_stats: ReceivedStats::Both,
+        stamp_at: StampAt::Both,
+        clock: Clock::Both,
+        ..Params::default()
+    };
+    let server = silent_open_server(params);
+    let config = ClientConfig {
+        duration: None,
+        interval: Duration::from_secs(1),
+        socket_config: crate::SocketConfig {
+            recv_timeout: Some(Duration::from_millis(200)),
+            ..Default::default()
+        },
+        ..default_test_config(server.addr)
+    };
+    let mut client = Client::connect(config).unwrap();
+    assert_open_started(client.open(ClientTimestamp::now()).unwrap());
+    client.session.as_mut().unwrap().packets_sent = u64::MAX;
+
+    assert!(matches!(
+        client.send_probe(),
+        Err(ClientError::DurationOverflow)
+    ));
+    server.join();
+}
+
 #[test]
 fn recv_once_returns_empty_on_timeout() {
     let params = default_params();