From a0751f7bd4c37ec9315bed745da1c430f7abaaab Mon Sep 17 00:00:00 2001 From: relaxing-more <2068804516@qq.com> Date: Tue, 31 Mar 2026 10:42:21 +0800 Subject: [PATCH] feat: implement 1-RTT key update management mechanism Signed-off-by: relaxing-more <2068804516@qq.com> --- .github/workflows/codecov.yml | 2 +- codecov.yml | 2 +- gm-quic/tests/common/mod.rs | 2 +- qbase/src/packet/keys.rs | 926 ++++++++++++++++++++++- qbase/src/packet/signal.rs | 8 - qconnection/src/space.rs | 16 +- qconnection/src/space/data.rs | 119 +-- qconnection/src/tls.rs | 2 +- qconnection/src/tx.rs | 85 ++- qinterface/src/component/route/packet.rs | 101 ++- 10 files changed, 1132 insertions(+), 131 deletions(-) diff --git a/.github/workflows/codecov.yml b/.github/workflows/codecov.yml index 48f115f80..563a0f720 100644 --- a/.github/workflows/codecov.yml +++ b/.github/workflows/codecov.yml @@ -21,4 +21,4 @@ jobs: with: token: ${{ secrets.CODECOV_TOKEN }} files: lcov.info - fail_ci_if_error: true + fail_ci_if_error: true \ No newline at end of file diff --git a/codecov.yml b/codecov.yml index 323e65f41..8367c2cf1 100644 --- a/codecov.yml +++ b/codecov.yml @@ -2,4 +2,4 @@ coverage: status: patch: off project: off - range: "70..100" + range: "70..100" \ No newline at end of file diff --git a/gm-quic/tests/common/mod.rs b/gm-quic/tests/common/mod.rs index f31522d8c..b64bc2602 100644 --- a/gm-quic/tests/common/mod.rs +++ b/gm-quic/tests/common/mod.rs @@ -57,7 +57,7 @@ pub fn run(future: F) -> F::Output { RT.block_on(async move { LazyLock::force(&TRACING); - match time::timeout(Duration::from_secs(60), future).await { + match time::timeout(Duration::from_secs(120), future).await { Ok(output) => output, Err(_timedout) => panic!("test timed out"), } diff --git a/qbase/src/packet/keys.rs b/qbase/src/packet/keys.rs index a3deea785..9e01d18fc 100644 --- a/qbase/src/packet/keys.rs +++ b/qbase/src/packet/keys.rs @@ -1,6 +1,7 @@ use std::{ + collections::{HashMap, VecDeque}, future::Future, - ops::DerefMut, + ops::{DerefMut, Range}, pin::Pin, sync::{Arc, Mutex, MutexGuard}, task::{Context, Poll, Waker}, @@ -49,7 +50,10 @@ impl From for Keys { } use super::KeyPhaseBit; -use crate::role::Role; +use crate::{ + error::{ErrorKind, QuicError}, + role::Role, +}; #[derive(Clone)] enum KeysState { @@ -155,7 +159,7 @@ impl ArcKeys { /// Asynchronously obtain the remote keys for removing header protection and packet decryption. /// - /// Rreturn [`GetRemoteKeys`], which implemented Future trait. + /// Returns [`GetRemoteKeys`], which implements the Future trait. /// /// ## Example /// @@ -224,7 +228,7 @@ impl ArcKeys { /// Especially in the closing state, the return keys are used to generate the final packet /// containing the ConnectionClose frame, and decrypt the data packets received from the /// peer for a while. - pub fn invalid(&self) -> Option { + pub fn invalidate(&self) -> Option { self.lock_guard().invalid() } } @@ -277,7 +281,7 @@ impl ArcZeroRttKeys { } } - pub fn invalid(&self) -> Option { + pub fn invalidate(&self) -> Option { self.lock_guard().invalid() } } @@ -287,64 +291,369 @@ impl ArcZeroRttKeys { /// /// See [key update](https://www.rfc-editor.org/rfc/rfc9001#name-key-update) /// of [RFC 9001](https://www.rfc-editor.org/rfc/rfc9001) for more details. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GetKeyError { + Retired, + Unknown, +} + +/// Entry in the 1-RTT key update window. +/// Each entry represents a key generation with both send and receive keys. +#[derive(Clone)] +pub struct KeyEntry { + /// Generation counter (0, 1, 2, ...) + pub generation: u64, + /// Send key for outgoing packets + pub local_key: Arc, + /// Receive key for incoming packets + pub remote_key: Arc, + /// Record received packet number range for this key generation + pub rcvd_pn_range: Option>, +} + +impl KeyEntry { + fn new(generation: u64, local_key: Arc, remote_key: Arc) -> Self { + Self { + generation, + local_key, + remote_key, + rcvd_pn_range: None, + } + } + + /// Update the received packet number range when a packet is successfully decrypted. + fn update_rcvd_pn(&mut self, pn: u64) { + match &mut self.rcvd_pn_range { + None => self.rcvd_pn_range = Some(pn..pn + 1), + Some(range) => { + range.start = range.start.min(pn); + range.end = range.end.max(pn + 1); + } + } + } +} + pub struct OneRttPacketKeys { + /// Key update counter, incremented on each update + counter: u64, + /// Ordered window of key entries (max 3 generations) + keys: VecDeque, + /// Secrets for derive next key pair + secrets: Option, + /// Current key phase bit cur_phase: KeyPhaseBit, - secrets: Secrets, - // TODO: 保存三个 - remote: [Option>; 2], - local: Arc, + /// Sent packet number ranges for each generation (for ACK tracking) + sent_pn_ranges: HashMap>, + /// Largest acknowledged packet number in 1-RTT space + largest_acked_pn: Option, + /// Map packet number to its generation + sent_pn_stage: HashMap, + /// Count of unacked packets in each generation + outstanding_pn_count: HashMap, + /// Consecutive decryption failures + contiguous_decrypt_failures: u32, +} + +/// Result of trying to decrypt with one key. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DecryptAttemptResult { + /// Decryption succeeded and returned plaintext body length. + Success(usize), + /// Decryption failed, but caller should continue trying other keys. + ContinueTrying, + /// Decryption failed and failure threshold is exceeded. + GiveUp, } impl OneRttPacketKeys { + /// Maximum number of consecutive decryption failures before giving up + const MAX_CONTIGUOUS_DECRYPT_FAILURES: u32 = 3; + /// Create new [`OneRttPacketKeys`]. /// /// The TLS handshake session must exchange enough information to generate the 1-RTT keys. fn new(remote: Box, local: Box, secrets: Secrets) -> Self { + let sk: Arc = Arc::from(local); + let rk: Arc = Arc::from(remote); + + let mut keys = VecDeque::new(); + let entry = KeyEntry::new(0, sk.clone(), rk); + keys.push_back(entry); + Self { + counter: 0, + keys, + secrets: Some(secrets), cur_phase: KeyPhaseBit::default(), - secrets, - remote: [Some(Arc::from(remote)), None], - local: Arc::from(local), + sent_pn_ranges: HashMap::new(), + largest_acked_pn: None, + sent_pn_stage: HashMap::new(), + outstanding_pn_count: HashMap::new(), + contiguous_decrypt_failures: 0, + } + } + + /// Get the local key for encrypting outgoing packets. + /// + /// Returns (generation, key_phase, packet_key). + fn get_local_key(&self) -> (u64, KeyPhaseBit, Arc) { + let entry = self.keys.back().expect("keys should not be empty"); + (entry.generation, self.cur_phase, entry.local_key.clone()) + } + + /// Get candidate generations based on key phase bit. + fn candidate_generations(&self, key_phase: KeyPhaseBit) -> [Option; 2] { + if let Some(latest) = self.keys.back().map(|e| e.generation) { + if key_phase == self.cur_phase { + [Some(latest), latest.checked_sub(2)] + } else { + [latest.checked_sub(1), latest.checked_add(1)] + } + } else { + [None, None] } } - /// Proactively update the 1-RTT packet key locally. - /// Or be informed by the peer to update the key. + /// Returns true if the local endpoint is ready to initiate a key update on send path. /// - /// The key phase bit will be toggled and sent to the peer, - /// informing the peer to update the key to next 1-RTT packet key too. - pub fn update(&mut self) { + /// Readiness means at least one packet from the current generation was acknowledged, + /// which avoids rotating too early before the peer has confirmed current keys. + fn is_send_key_update_ready(&self) -> bool { + let latest_gen = self.keys.back().map(|e| e.generation).unwrap_or(0); + let Some(min_sent_pn) = self.sent_pn_ranges.get(&latest_gen).map(|r| r.start) else { + return false; + }; + let Some(largest_acked) = self.largest_acked_pn else { + return false; + }; + largest_acked >= min_sent_pn + } + + /// Derive and install the next key generation. + /// + /// This advances key phase, appends a new generation, and keeps only recent generations. + fn rotate_to_next_key_generation(&mut self) { + let key_set = self + .secrets + .as_mut() + .expect("1-RTT secrets must exist when updating keys") + .next_packet_keys(); + self.counter += 1; + + let entry = KeyEntry::new( + self.counter, + Arc::from(key_set.local), + Arc::from(key_set.remote), + ); + + self.keys.push_back(entry); + while self.keys.len() > 3 { + self.keys.pop_front(); + } + self.cur_phase.toggle(); - let key_set = self.secrets.next_packet_keys(); - self.remote[self.cur_phase.as_index()] = Some(Arc::from(key_set.remote)); - self.local = Arc::from(key_set.local); } - /// Old key must be phased out within a certain period of time. + /// When a packet in generation i is acked, we know the peer has received + /// a packet encrypted with key i, confirming the key update. + fn apply_packet_acked( + &mut self, + pn: u64, + largest_ack: u64, + rcvd_generation: Option, + ) -> Result<(), QuicError> { + self.largest_acked_pn = Some( + self.largest_acked_pn + .map_or(largest_ack, |cur| cur.max(largest_ack)), + ); + + let Some(stage) = self.sent_pn_stage.remove(&pn) else { + return Ok(()); + }; + + if let Some(rcvd_generation) = rcvd_generation + && stage > rcvd_generation + { + return Err(QuicError::with_default_fty( + ErrorKind::KeyUpdate, + "peer acknowledged new-key packet with older key phase", + )); + } + + let Some(count) = self.outstanding_pn_count.get_mut(&stage) else { + return Ok(()); + }; + + *count = count.saturating_sub(1); + if *count == 0 { + self.outstanding_pn_count.remove(&stage); + } + Ok(()) + } + + /// Record sent packet number for key update tracking + fn track_packet_sent(&mut self, generation: u64, pn: u64) { + let entry = self.sent_pn_ranges.entry(generation).or_insert(pn..pn); + entry.start = entry.start.min(pn); + entry.end = entry.end.max(pn + 1); + + if let Some(old_stage) = self.sent_pn_stage.insert(pn, generation) + && let Some(count) = self.outstanding_pn_count.get_mut(&old_stage) + { + *count = count.saturating_sub(1); + if *count == 0 { + self.outstanding_pn_count.remove(&old_stage); + } + } + + *self.outstanding_pn_count.entry(generation).or_default() += 1; + } + + /// Get the remote packet key for a specific generation. /// - /// If the old one don't go, the new ones won't come. - /// If it is not phased out, it will be considered as new keys and - /// fail to decrypt the packet in future. - pub fn phase_out(&mut self) { - self.remote[(!self.cur_phase).as_index()].take(); + /// Returns `Some(key)` if the remote key for the given generation exists. + /// Returns `None` if the generation is the next expected one (key not yet derived). + /// Returns `Err` if the generation is invalid or has been discarded. + fn get_remote_key_for_generation( + &self, + generation: u64, + ) -> Result>, GetKeyError> { + for entry in self.keys.iter() { + if entry.generation == generation { + return Ok(Some(entry.remote_key.clone())); + } + } + + if let Some(latest) = self.keys.back().map(|e| e.generation) { + if generation == latest + 1 { + return Ok(None); + } + } + + Err(GetKeyError::Unknown) } - /// Get the remote key to decrypt the incoming 1-RTT packet. - /// If the key phase is not the current key phase, update the key, see [`Self::update`]. + /// Passive update in response to peer-initiated key phase change + fn update_by_peer(&mut self) -> Result<(), QuicError> { + if let Some(latest) = self.keys.back() + && latest.rcvd_pn_range.is_none() + && latest.generation > 0 + { + return Err(QuicError::with_default_fty( + ErrorKind::KeyUpdate, + "received consecutive peer key updates before confirming prior update", + )); + } + + self.rotate_to_next_key_generation(); + Ok(()) + } + + /// Called when a packet is successfully decrypted with key of generation i. + /// + /// This marks the generation as confirmed by received data and updates the pn range. + fn on_packet_decrypted(&mut self, rcvd_pn: u64, generation: u64) -> Result<(), QuicError> { + let max_newer_pn = self + .keys + .iter() + .filter(|entry| entry.generation > generation) + .filter_map(|entry| entry.rcvd_pn_range.as_ref().map(|r| r.end - 1)) + .max(); + + if let Some(max_newer_pn) = max_newer_pn + && rcvd_pn > max_newer_pn + { + return Err(QuicError::with_default_fty( + ErrorKind::KeyUpdate, + "key downgrade detected: higher packet number decrypted with old key", + )); + } + + for entry in self.keys.iter_mut() { + if entry.generation == generation { + entry.update_rcvd_pn(rcvd_pn); + break; + } + } + + self.contiguous_decrypt_failures = 0; + Ok(()) + } + + /// Called when packet decryption fails with key of generation i. /// - /// Return `Arc` to decrypt the incoming 1-RTT packet. - pub fn get_remote(&mut self, key_phase: KeyPhaseBit, _pn: u64) -> Arc { - if key_phase != self.cur_phase && self.remote[key_phase.as_index()].is_none() { - self.update(); + /// Returns true if we should give up (threshold exceeded), false otherwise. + fn on_packet_decrypt_failed(&mut self) -> bool { + self.contiguous_decrypt_failures += 1; + self.contiguous_decrypt_failures > Self::MAX_CONTIGUOUS_DECRYPT_FAILURES + } + + /// Handle the result of one decryption attempt and update key state. + fn apply_decrypt_attempt( + &mut self, + generation: u64, + decoded_pn: u64, + decrypt_result: Result, + ) -> Result { + match decrypt_result { + Ok(body_length) => { + self.on_packet_decrypted(decoded_pn, generation)?; + Ok(DecryptAttemptResult::Success(body_length)) + } + Err(_) => { + let reached_threshold = self.on_packet_decrypt_failed(); + if reached_threshold { + Ok(DecryptAttemptResult::GiveUp) + } else { + Ok(DecryptAttemptResult::ContinueTrying) + } + } } - self.remote[key_phase.as_index()].clone().unwrap() } - /// Get the local current key to encrypt the outgoing packet. + /// Get all candidate keys to try for decryption, in order of preference. /// - /// Return `Arc` to encrypt the outgoing 1-RTT packet. - pub fn get_local(&self) -> (KeyPhaseBit, Arc) { - (self.cur_phase, self.local.clone()) + /// Returns an iterator of (generation, key) tuples, and an optional next_generation + /// that may need key update. + fn collect_decrypt_key_candidates( + &self, + decoded_pn: u64, + key_phase: KeyPhaseBit, + ) -> (Vec<(u64, Arc)>, Option) { + let mut candidates = Vec::new(); + let mut next_generation = None; + + // First, try primary key by packet number range + for entry in self.keys.iter() { + if let Some(range) = &entry.rcvd_pn_range { + if decoded_pn >= range.start && decoded_pn < range.end { + let expected_phase = if entry.generation % 2 == 0 { + KeyPhaseBit::Zero + } else { + KeyPhaseBit::One + }; + if key_phase == expected_phase { + candidates.push((entry.generation, entry.remote_key.clone())); + return (candidates, None); + } + } + } + } + + // Try candidate generations + for generation in self.candidate_generations(key_phase).into_iter().flatten() { + match self.get_remote_key_for_generation(generation) { + Ok(Some(key)) => candidates.push((generation, key)), + Ok(None) => { + if next_generation.is_none() { + next_generation = Some(generation); + } + } + Err(_) => {} + } + } + + (candidates, next_generation) } } @@ -361,10 +670,75 @@ impl ArcOneRttPacketKeys { /// Obtain exclusive access to the 1-RTT packet keys. /// During the exclusive period of encrypting or decrypting packets, /// the keys must not be updated elsewhere. - pub fn lock_guard(&self) -> MutexGuard<'_, OneRttPacketKeys> { + fn lock_guard(&self) -> MutexGuard<'_, OneRttPacketKeys> { self.0.0.lock().unwrap() } + /// Update key phase if possible and return current local packet key material for sending. + pub fn update_and_get_local_key(&self) -> (u64, KeyPhaseBit, Arc) { + let mut keys = self.lock_guard(); + if keys.is_send_key_update_ready() { + keys.rotate_to_next_key_generation(); + } + keys.get_local_key() + } + + /// Record a sent packet for key update tracking. + pub fn on_packet_sent(&self, generation: u64, pn: u64) { + self.lock_guard().track_packet_sent(generation, pn); + } + + /// Handle ACK feedback for key update tracking. + pub fn on_packet_acked( + &self, + pn: u64, + largest_ack: u64, + rcvd_generation: Option, + ) -> Result<(), QuicError> { + self.lock_guard() + .apply_packet_acked(pn, largest_ack, rcvd_generation) + } + + /// Get decryption candidate keys in preference order. + pub fn get_candidate_keys_for_decryption( + &self, + decoded_pn: u64, + key_phase: KeyPhaseBit, + ) -> (Vec<(u64, Arc)>, Option) { + self.lock_guard() + .collect_decrypt_key_candidates(decoded_pn, key_phase) + } + + /// Get remote key for generation; if generation is the next one, passively update once. + pub fn resolve_remote_key_for_generation( + &self, + generation: u64, + ) -> Result>, QuicError> { + let mut keys = self.lock_guard(); + match keys.get_remote_key_for_generation(generation) { + Ok(Some(key)) => Ok(Some(key)), + Ok(None) => { + keys.update_by_peer()?; + Ok(keys + .get_remote_key_for_generation(generation) + .ok() + .flatten()) + } + Err(_) => Ok(None), + } + } + + /// Update decryption-related key state by one decryption attempt result. + pub fn record_decrypt_attempt( + &self, + generation: u64, + decoded_pn: u64, + decrypt_result: Result, + ) -> Result { + self.lock_guard() + .apply_decrypt_attempt(generation, decoded_pn, decrypt_result) + } + /// Get the length of the tag of the packet key's underlying AEAD algorithm. /// /// For example, when collecting data to send, buffer needs to reserve @@ -451,7 +825,7 @@ impl ArcOneRttKeys { } } - pub fn invalid(&self) -> Option<(HeaderProtectionKeys, ArcOneRttPacketKeys)> { + pub fn invalidate(&self) -> Option<(HeaderProtectionKeys, ArcOneRttPacketKeys)> { let mut state = self.lock_guard(); match std::mem::replace(state.deref_mut(), OneRttKeysState::Invalid) { OneRttKeysState::Pending(rx_waker) => { @@ -487,7 +861,7 @@ impl ArcOneRttKeys { /// Asynchronously obtain the remote keys for removing header protection and packet decryption. /// - /// Rreturn [`GetRemoteKeys`], which implemented the Future trait. + /// Returns [`GetRemoteKeys`], which implements the Future trait. pub fn get_remote_keys(&self) -> GetRemoteOneRttKeys<'_> { GetRemoteOneRttKeys(self) } @@ -522,3 +896,475 @@ impl Future for GetRemoteOneRttKeys<'_> { } } } + +#[cfg(test)] +mod tests { + use rustls::{Error as RustlsError, quic::Tag}; + + use super::*; + + struct DummyPacketKey; + + impl HeaderProtectionKey for DummyPacketKey { + fn encrypt_in_place( + &self, + _sample: &[u8], + _first_byte: &mut u8, + _packet_number: &mut [u8], + ) -> Result<(), RustlsError> { + Err(RustlsError::General("dummy key".into())) + } + + fn decrypt_in_place( + &self, + _sample: &[u8], + _first_byte: &mut u8, + _packet_number: &mut [u8], + ) -> Result<(), RustlsError> { + Err(RustlsError::General("dummy key".into())) + } + + fn sample_len(&self) -> usize { + 16 + } + } + + impl PacketKey for DummyPacketKey { + fn encrypt_in_place( + &self, + _packet_number: u64, + _header: &[u8], + _payload: &mut [u8], + ) -> Result { + Err(RustlsError::General("dummy key".into())) + } + + fn decrypt_in_place<'a>( + &self, + _packet_number: u64, + _header: &[u8], + _payload: &'a mut [u8], + ) -> Result<&'a [u8], RustlsError> { + Err(RustlsError::General("dummy key".into())) + } + + fn tag_len(&self) -> usize { + 16 + } + + fn confidentiality_limit(&self) -> u64 { + u64::MAX + } + + fn integrity_limit(&self) -> u64 { + u64::MAX + } + } + + fn dummy_pk() -> Arc { + Arc::new(DummyPacketKey) + } + + fn test_keys_with_entries(entries: VecDeque) -> OneRttPacketKeys { + OneRttPacketKeys { + counter: entries.back().map(|e| e.generation).unwrap_or(0), + keys: entries, + secrets: None, + cur_phase: KeyPhaseBit::Zero, + sent_pn_ranges: HashMap::new(), + largest_acked_pn: None, + sent_pn_stage: HashMap::new(), + outstanding_pn_count: HashMap::new(), + contiguous_decrypt_failures: 0, + } + } + + #[test] + fn key_update_error_on_consecutive_peer_updates_before_confirming_first() { + let mut entries = VecDeque::new(); + let mut e0 = KeyEntry::new(0, dummy_pk(), dummy_pk()); + e0.rcvd_pn_range = Some(1..11); + entries.push_back(e0); + + let mut e1 = KeyEntry::new(1, dummy_pk(), dummy_pk()); + e1.rcvd_pn_range = None; + entries.push_back(e1); + + let mut keys = test_keys_with_entries(entries); + let err = keys + .update_by_peer() + .expect_err("must reject detect consecutive peer key update"); + assert_eq!(err.kind(), ErrorKind::KeyUpdate); + } + + #[test] + fn key_update_error_on_peer_unsynced_ack_using_old_key_phase() { + let mut entries = VecDeque::new(); + entries.push_back(KeyEntry::new(0, dummy_pk(), dummy_pk())); + entries.push_back(KeyEntry::new(1, dummy_pk(), dummy_pk())); + let mut keys = test_keys_with_entries(entries); + + keys.sent_pn_stage.insert(42, 1); + + let err = keys + .apply_packet_acked(42, 42, Some(0)) + .expect_err("must reject old-key ACKing new-key packet"); + assert_eq!(err.kind(), ErrorKind::KeyUpdate); + } + + #[test] + fn key_update_error_on_key_downgrade_high_pn_decrypted_with_old_key() { + let mut entries = VecDeque::new(); + let mut e0 = KeyEntry::new(0, dummy_pk(), dummy_pk()); + e0.rcvd_pn_range = Some(1..51); + entries.push_back(e0); + + let mut e1 = KeyEntry::new(1, dummy_pk(), dummy_pk()); + e1.rcvd_pn_range = Some(100..121); + entries.push_back(e1); + + let mut keys = test_keys_with_entries(entries); + let err = keys + .on_packet_decrypted(130, 0) + .expect_err("must reject decrypting higher PN with old key"); + assert_eq!(err.kind(), ErrorKind::KeyUpdate); + } + + #[test] + fn normal_key_update_flow() { + let mut entries = VecDeque::new(); + entries.push_back(KeyEntry::new(0, dummy_pk(), dummy_pk())); + + let mut keys = test_keys_with_entries(entries); + + // Check initial state + let (generation, phase, _key) = keys.get_local_key(); + assert_eq!(generation, 0); + assert_eq!(phase, KeyPhaseBit::Zero); + + // Set acknowledged packet numbers to make key update ready + keys.sent_pn_ranges.insert(0, 0..10); + keys.largest_acked_pn = Some(5); + + // Note: Since rotate_to_next_key_generation() requires secrets, + // this test mainly verifies the preparation conditions for key update, + // not the actual key rotation + // Verify key update readiness status + let is_ready = keys.is_send_key_update_ready(); + assert!(is_ready); + + // Verify key selection logic + let candidates = keys.candidate_generations(KeyPhaseBit::Zero); + assert_eq!(candidates, [Some(0), None]); + } + + #[test] + fn key_generation_boundary_management() { + let mut entries = VecDeque::new(); + + // Create 4 generations of keys to test window management + for i in 0..4 { + entries.push_back(KeyEntry::new(i, dummy_pk(), dummy_pk())); + } + + let mut keys = test_keys_with_entries(entries); + keys.counter = 3; + + // Note: Since rotate_to_next_key_generation() requires secrets, + // this test mainly verifies the key generation boundary management logic + // Verify initial state + assert_eq!(keys.keys.len(), 4); + assert_eq!(keys.keys.front().unwrap().generation, 0); + assert_eq!(keys.keys.back().unwrap().generation, 3); + + // Test candidate generation selection logic + let candidates = keys.candidate_generations(KeyPhaseBit::Zero); + assert_eq!(candidates, [Some(3), Some(1)]); + + // Verify key generation boundary management (manually simulate window management) + while keys.keys.len() > 3 { + keys.keys.pop_front(); + } + assert_eq!(keys.keys.len(), 3); + assert_eq!(keys.keys.front().unwrap().generation, 1); + assert_eq!(keys.keys.back().unwrap().generation, 3); + } + + #[test] + fn packet_number_ordering_and_range_management() { + let mut entries = VecDeque::new(); + entries.push_back(KeyEntry::new(0, dummy_pk(), dummy_pk())); + + let mut keys = test_keys_with_entries(entries); + + // Track sent packets + keys.track_packet_sent(0, 10); + keys.track_packet_sent(0, 20); + keys.track_packet_sent(0, 15); + + // Verify packet number range + let range = keys.sent_pn_ranges.get(&0).unwrap(); + assert_eq!(range.start, 10); + assert_eq!(range.end, 21); // Packets 10, 15, 20 -> range 10..21 + + // Verify packet number to generation mapping + assert_eq!(keys.sent_pn_stage.get(&10), Some(&0)); + assert_eq!(keys.sent_pn_stage.get(&15), Some(&0)); + assert_eq!(keys.sent_pn_stage.get(&20), Some(&0)); + } + + #[test] + fn successful_decryption_flow() { + let mut entries = VecDeque::new(); + entries.push_back(KeyEntry::new(0, dummy_pk(), dummy_pk())); + + let mut keys = test_keys_with_entries(entries); + + // Successfully decrypt packet + let result = keys.apply_decrypt_attempt(0, 5, Ok(100)); + assert!(matches!(result, Ok(DecryptAttemptResult::Success(100)))); + + // Verify packet number range has been updated + let entry = keys.keys.front().unwrap(); + assert_eq!(entry.rcvd_pn_range, Some(5..6)); + + // Consecutive decryption failure counter should be reset + assert_eq!(keys.contiguous_decrypt_failures, 0); + } + + #[test] + fn candidate_key_selection_logic() { + let mut entries = VecDeque::new(); + + // Create multiple generations of keys + for i in 0..3 { + let mut entry = KeyEntry::new(i, dummy_pk(), dummy_pk()); + if i == 0 { + entry.rcvd_pn_range = Some(1..11); + } + entries.push_back(entry); + } + + let keys = test_keys_with_entries(entries); + + // Test selection when packet number is within range + let (candidates, next_gen) = keys.collect_decrypt_key_candidates(5, KeyPhaseBit::Zero); + assert_eq!(candidates.len(), 1); // Should only select one candidate + assert_eq!(candidates[0].0, 0); // Should select generation 0 + assert!(next_gen.is_none()); + + // Test selection when packet number is outside range + let (candidates, next_gen) = keys.collect_decrypt_key_candidates(20, KeyPhaseBit::Zero); + assert_eq!(candidates.len(), 2); // Should select two candidates + assert!(next_gen.is_none()); + } + + #[test] + fn key_phase_bit_management() { + let mut entries = VecDeque::new(); + entries.push_back(KeyEntry::new(0, dummy_pk(), dummy_pk())); + + let mut keys = test_keys_with_entries(entries); + + // Initial phase + assert_eq!(keys.cur_phase, KeyPhaseBit::Zero); + + // Test candidate generation selection + let candidates = keys.candidate_generations(KeyPhaseBit::Zero); + assert_eq!(candidates, [Some(0), None]); // Current generation and generation-2 + + // Switch phase + keys.cur_phase = KeyPhaseBit::One; + let candidates = keys.candidate_generations(KeyPhaseBit::Zero); + assert_eq!(candidates, [None, Some(1)]); // Generation-1 and generation+1 + } + + #[test] + fn decryption_failure_handling() { + let mut entries = VecDeque::new(); + entries.push_back(KeyEntry::new(0, dummy_pk(), dummy_pk())); + + let mut keys = test_keys_with_entries(entries); + + // Test consecutive decryption failures + for i in 0..3 { + let result = keys.apply_decrypt_attempt(0, i, Err(())); + assert!(matches!(result, Ok(DecryptAttemptResult::ContinueTrying))); + assert_eq!(keys.contiguous_decrypt_failures, i as u32 + 1); + } + + // The 4th failure should give up + let result = keys.apply_decrypt_attempt(0, 3, Err(())); + assert!(matches!(result, Ok(DecryptAttemptResult::GiveUp))); + } + + #[test] + fn arc_keys_concurrent_access_safety() { + use std::{sync::Arc, thread}; + + let keys = ArcKeys::new_pending(); + let keys_clone = keys.clone(); + + // Set keys in another thread + let handle = thread::spawn(move || { + let header_key = Arc::new(DummyPacketKey) as Arc; + let packet_key = dummy_pk(); + keys_clone.set_keys(Keys { + local: DirectionalKeys { + header: header_key.clone(), + packet: packet_key.clone(), + }, + remote: DirectionalKeys { + header: header_key, + packet: packet_key, + }, + }); + }); + + handle.join().unwrap(); + + // Verify keys have been properly set + let local_keys = keys.get_local_keys(); + assert!(local_keys.is_some()); + } + + #[test] + fn arc_one_rtt_keys_integration() { + let keys = ArcOneRttKeys::new_pending(); + + // Verify initial state + let local_keys = keys.get_local_keys(); + assert!(local_keys.is_none()); + + // Test key invalidation + let invalidated = keys.invalidate(); + assert!(invalidated.is_none()); + + // Verify state after invalidation + let local_keys = keys.get_local_keys(); + assert!(local_keys.is_none()); + } + + #[test] + fn keys_state_lifecycle() { + // Test Pending -> Ready transition + let mut state: KeysState = KeysState::Pending(None); + assert!(state.get().is_none()); + + state.set("test".to_string()); + assert_eq!(state.get(), Some(&"test".to_string())); + + // Test Invalid transition + let retrieved = state.invalid(); + assert_eq!(retrieved, Some("test".to_string())); + assert!(state.get().is_none()); + + // Test Ready -> Invalid transition + let mut state: KeysState = KeysState::Ready("ready".to_string()); + let retrieved = state.invalid(); + assert_eq!(retrieved, Some("ready".to_string())); + assert!(state.get().is_none()); + } + + #[test] + fn keys_state_waker_behavior() { + let mut state: KeysState = KeysState::Pending(None); + + // Setting keys should work even without waker + state.set("test".to_string()); + assert!(state.get().is_some()); + + // Invalid state should still work + state.invalid(); + assert!(state.get().is_none()); + } + + #[test] + fn directional_keys_conversion() { + let header_key = Arc::new(DummyPacketKey) as Arc; + let packet_key = dummy_pk(); + + let dir_keys = DirectionalKeys { + header: header_key.clone(), + packet: packet_key.clone(), + }; + + assert_eq!(dir_keys.header.sample_len(), 16); + assert_eq!(dir_keys.packet.tag_len(), 16); + } + + #[test] + fn keys_conversion() { + let header_key = Arc::new(DummyPacketKey) as Arc; + let packet_key = dummy_pk(); + + let keys = Keys { + local: DirectionalKeys { + header: header_key.clone(), + packet: packet_key.clone(), + }, + remote: DirectionalKeys { + header: header_key.clone(), + packet: packet_key.clone(), + }, + }; + + assert_eq!(keys.local.header.sample_len(), 16); + assert_eq!(keys.remote.packet.tag_len(), 16); + } + + #[test] + fn key_entry_update_rcvd_pn() { + let local_key = dummy_pk(); + let remote_key = dummy_pk(); + let mut entry = KeyEntry::new(0, local_key, remote_key); + + assert!(entry.rcvd_pn_range.is_none()); + + entry.update_rcvd_pn(5); + assert_eq!(entry.rcvd_pn_range, Some(5..6)); + + entry.update_rcvd_pn(3); + assert_eq!(entry.rcvd_pn_range, Some(3..6)); + + entry.update_rcvd_pn(10); + assert_eq!(entry.rcvd_pn_range, Some(3..11)); + } + + #[test] + fn get_key_error_variants() { + // Test GetKeyError variants are comparable + assert_eq!(GetKeyError::Retired, GetKeyError::Retired); + assert_eq!(GetKeyError::Unknown, GetKeyError::Unknown); + assert_ne!(GetKeyError::Retired, GetKeyError::Unknown); + } + + #[test] + fn decrypt_attempt_result_variants() { + assert_eq!( + DecryptAttemptResult::Success(100), + DecryptAttemptResult::Success(100) + ); + assert_eq!( + DecryptAttemptResult::ContinueTrying, + DecryptAttemptResult::ContinueTrying + ); + assert_eq!(DecryptAttemptResult::GiveUp, DecryptAttemptResult::GiveUp); + assert_ne!( + DecryptAttemptResult::Success(100), + DecryptAttemptResult::Success(200) + ); + } + + #[test] + fn key_entry_basic_properties() { + let local_key = dummy_pk(); + let remote_key = dummy_pk(); + let entry = KeyEntry::new(5, local_key, remote_key); + + assert_eq!(entry.generation, 5); + assert_eq!(entry.local_key.tag_len(), 16); + assert_eq!(entry.remote_key.tag_len(), 16); + assert!(entry.rcvd_pn_range.is_none()); + } +} diff --git a/qbase/src/packet/signal.rs b/qbase/src/packet/signal.rs index 1e87bd09f..96352ae43 100644 --- a/qbase/src/packet/signal.rs +++ b/qbase/src/packet/signal.rs @@ -43,14 +43,6 @@ impl Toggle { Toggle::One => *byte |= B, } } - - /// Treat Toggle as an index and get the index value it represents, i.e., 0 or 1 - pub(crate) fn as_index(&self) -> usize { - match self { - Toggle::Zero => 0, - Toggle::One => 1, - } - } } impl std::ops::Not for Toggle { diff --git a/qconnection/src/space.rs b/qconnection/src/space.rs index ec9167b6c..5e42ba96e 100644 --- a/qconnection/src/space.rs +++ b/qconnection/src/space.rs @@ -15,6 +15,7 @@ use qbase::{ AssemblePacket, Package, PacketContains, PacketSpace, PacketWriter, ProductHeader, header::{GetDcid, GetType, short::OneRttHeader}, io::{Packages, PadTo20}, + keys::ArcOneRttKeys, }, }; use qevent::{ @@ -259,6 +260,7 @@ impl ReceiveFrame for AckHandshakeSpace { struct AckDataSpace { send_journal: ArcSentJournal, + one_rtt_keys: ArcOneRttKeys, data_streams: DataStreams, crypto_stream_outgoing: CryptoStreamOutgoing, } @@ -266,21 +268,24 @@ struct AckDataSpace { impl AckDataSpace { fn new( journal: &Journal, + one_rtt_keys: ArcOneRttKeys, data_streams: &DataStreams, crypto_stream: &CryptoStream, ) -> Self { Self { send_journal: journal.of_sent_packets(), + one_rtt_keys, data_streams: data_streams.clone(), crypto_stream_outgoing: crypto_stream.outgoing(), } } } -impl ReceiveFrame for AckDataSpace { +impl ReceiveFrame<(AckFrame, Option)> for AckDataSpace { type Output = (); - fn recv_frame(&self, ack_frame: &AckFrame) -> Result { + fn recv_frame(&self, ack_input: &(AckFrame, Option)) -> Result { + let (ack_frame, rcvd_generation) = ack_input; let mut rotate_guard = self.send_journal.rotate(); rotate_guard.update_largest(ack_frame)?; @@ -289,7 +294,14 @@ impl ReceiveFrame for AckDataSpace { packet_number_space: qbase::Epoch::Data, packet_nubers: acked.clone(), }); + + let one_rtt_pk = self.one_rtt_keys.get_local_keys().map(|(_, pk)| pk); + for pn in acked { + if let Some(pk) = &one_rtt_pk { + pk.on_packet_acked(pn, ack_frame.largest(), *rcvd_generation)?; + } + for frame in rotate_guard.on_packet_acked(pn) { match frame { GuaranteedFrame::Stream(stream_frame) => { diff --git a/qconnection/src/space/data.rs b/qconnection/src/space/data.rs index ce2d7fd45..e0b73acd7 100644 --- a/qconnection/src/space/data.rs +++ b/qconnection/src/space/data.rs @@ -43,7 +43,9 @@ use crate::{ }, state, termination::Terminator, - tx::{PacketWriter, TrivialPacketWriter}, + tx::{ + OneRttPacketKeysMeta, OneRttPacketMeta, PacketTimeouts, PacketWriter, TrivialPacketWriter, + }, }; pub type CipherZeroRttPacket = CipherPacket; @@ -212,20 +214,28 @@ impl path::PacketSpace for DataSpace { cc: &ArcCC, buffer: &'b mut [u8], ) -> Result, Signals> { - let (hpk, pk) = self.one_rtt_keys.get_local_keys().ok_or(Signals::KEYS)?; - let (key_phase, pk) = pk.lock_guard().get_local(); + let (hpk, pk_keys) = self.one_rtt_keys.get_local_keys().ok_or(Signals::KEYS)?; + let (key_generation, key_phase, pk) = pk_keys.update_and_get_local_key(); let (retran_timeout, expire_timeout) = cc.retransmit_and_expire_time(Epoch::Data); PacketWriter::new_short( header, buffer, - DirectionalKeys { - header: hpk, - packet: pk, - }, - key_phase, self.journal.as_ref(), - retran_timeout, - expire_timeout, + OneRttPacketKeysMeta { + keys: DirectionalKeys { + header: hpk, + packet: pk, + }, + meta: OneRttPacketMeta { + key_phase, + key_generation, + packet_keys: pk_keys, + }, + }, + PacketTimeouts { + retran_timeout, + expire_timeout, + }, ) } } @@ -239,16 +249,22 @@ impl PacketSpace for DataSpace { header: OneRttHeader, buffer: &'a mut [u8], ) -> Result, Signals> { - let (hpk, pk) = self.one_rtt_keys.get_local_keys().ok_or(Signals::KEYS)?; - let (key_phase, pk) = pk.lock_guard().get_local(); + let (hpk, pk_keys) = self.one_rtt_keys.get_local_keys().ok_or(Signals::KEYS)?; + let (key_generation, key_phase, pk) = pk_keys.update_and_get_local_key(); TrivialPacketWriter::new_short( header, buffer, - DirectionalKeys { - header: hpk, - packet: pk, + OneRttPacketKeysMeta { + keys: DirectionalKeys { + header: hpk, + packet: pk, + }, + meta: OneRttPacketMeta { + key_phase, + key_generation, + packet_keys: pk_keys, + }, }, - key_phase, self.journal.as_ref(), ) } @@ -258,7 +274,7 @@ fn frame_dispathcer( space: &DataSpace, components: &Components, event_broker: &ArcEventBroker, -) -> impl for<'p> Fn(Frame, Type, &'p Path) + use<> { +) -> impl for<'p> Fn(Frame, Type, &'p Path, Option) + use<> { let (ack_frames_entry, rcvd_ack_frames) = mpsc::unbounded_channel(); // 连接级的 let (max_data_frames_entry, rcvd_max_data_frames) = mpsc::unbounded_channel(); @@ -333,6 +349,7 @@ fn frame_dispathcer( rcvd_ack_frames, AckDataSpace::new( &space.journal, + space.one_rtt_keys(), &components.data_streams, &components.crypto_streams[space.epoch()], ), @@ -351,33 +368,38 @@ fn frame_dispathcer( let event_broker = event_broker.clone(); let rcvd_joural = space.journal.of_rcvd_packets(); - let dispathc_v1_frame = move |frame: V1Frame, pty: packet::Type, path: &Path| match frame { - V1Frame::Ack(f) => { - path.cc().on_ack_rcvd(Epoch::Data, &f); - rcvd_joural.on_rcvd_ack(&f); - _ = ack_frames_entry.send(f) - } - V1Frame::NewToken(f) => _ = new_token_frames_entry.send(f), - V1Frame::MaxData(f) => _ = max_data_frames_entry.send(f), - V1Frame::NewConnectionId(f) => _ = new_cid_frames_entry.send(f), - V1Frame::RetireConnectionId(f) => _ = retire_cid_frames_entry.send(f), - V1Frame::HandshakeDone(f) => { - // See [Section 4.1.2](https://datatracker.ietf.org/doc/html/rfc9001#handshake-confirmed) - _ = handshake_done_frames_entry.send(f) - } - V1Frame::DataBlocked(f) => _ = data_blocked_frames_entry.send(f), - V1Frame::Challenge(f) => _ = path.recv_frame(&f), - V1Frame::Response(f) => _ = path.recv_frame(&f), - V1Frame::StreamCtl(f) => _ = stream_ctrl_frames_entry.send(f), - V1Frame::Stream(f, data) => _ = stream_frames_entry.send((f, data)), - V1Frame::Crypto(f, bytes) => _ = crypto_frames_entry.send((f, bytes)), - #[cfg(feature = "unreliable")] - V1Frame::Datagram(f, data) => _ = datagram_frames_entry.send((f, data)), - V1Frame::Close(f) if matches!(pty, Type::Short(_)) => event_broker.emit(Event::Closed(f)), - _ => {} - }; - move |frame, pty, path| match frame { - Frame::V1(frame) => dispathc_v1_frame(frame, pty, path), + let dispathc_v1_frame = + move |frame: V1Frame, pty: packet::Type, path: &Path, key_generation: Option| { + match frame { + V1Frame::Ack(f) => { + path.cc().on_ack_rcvd(Epoch::Data, &f); + rcvd_joural.on_rcvd_ack(&f); + _ = ack_frames_entry.send((f, key_generation)) + } + V1Frame::NewToken(f) => _ = new_token_frames_entry.send(f), + V1Frame::MaxData(f) => _ = max_data_frames_entry.send(f), + V1Frame::NewConnectionId(f) => _ = new_cid_frames_entry.send(f), + V1Frame::RetireConnectionId(f) => _ = retire_cid_frames_entry.send(f), + V1Frame::HandshakeDone(f) => { + // See [Section 4.1.2](https://datatracker.ietf.org/doc/html/rfc9001#handshake-confirmed) + _ = handshake_done_frames_entry.send(f) + } + V1Frame::DataBlocked(f) => _ = data_blocked_frames_entry.send(f), + V1Frame::Challenge(f) => _ = path.recv_frame(&f), + V1Frame::Response(f) => _ = path.recv_frame(&f), + V1Frame::StreamCtl(f) => _ = stream_ctrl_frames_entry.send(f), + V1Frame::Stream(f, data) => _ = stream_frames_entry.send((f, data)), + V1Frame::Crypto(f, bytes) => _ = crypto_frames_entry.send((f, bytes)), + #[cfg(feature = "unreliable")] + V1Frame::Datagram(f, data) => _ = datagram_frames_entry.send((f, data)), + V1Frame::Close(f) if matches!(pty, Type::Short(_)) => { + event_broker.emit(Event::Closed(f)) + } + _ => {} + } + }; + move |frame, pty, path, key_generation| match frame { + Frame::V1(frame) => dispathc_v1_frame(frame, pty, path, key_generation), Frame::Traversal(frame) => { _ = traversal_frames_entry.send(( path.bind_uri().clone(), @@ -393,7 +415,7 @@ async fn parse_normal_zero_rtt_packet( (packet, (bind_uri, pathway, link)): ReceivedZeroRttFrom, space: &DataSpace, components: &Components, - dispatch_frame: impl Fn(Frame, Type, &Path), + dispatch_frame: impl Fn(Frame, Type, &Path, Option), ) -> Result<(), Error> { let Some(packet) = space.decrypt_0rtt_packet(packet).await.transpose()? else { return Ok(()); @@ -416,7 +438,7 @@ async fn parse_normal_zero_rtt_packet( }; let packet_contains = read_plain_packet(&packet, |frame| { - dispatch_frame(frame, packet.get_type(), &path); + dispatch_frame(frame, packet.get_type(), &path, None); })?; space.journal.of_rcvd_packets().on_rcvd_pn( @@ -433,7 +455,7 @@ async fn parse_normal_one_rtt_packet( (packet, (bind_uri, pathway, link)): ReceivedOneRttFrom, space: &DataSpace, components: &Components, - dispatch_frame: impl Fn(Frame, Type, &Path), + dispatch_frame: impl Fn(Frame, Type, &Path, Option), ) -> Result<(), Error> { let Some(packet) = space.decrypt_1rtt_packet(packet).await.transpose()? else { return Ok(()); @@ -459,8 +481,9 @@ async fn parse_normal_one_rtt_packet( .quic_handshake .discard_spaces_on_server_handshake_done(&components.paths); + let key_generation = packet.key_generation(); let packet_contains = read_plain_packet(&packet, |frame| { - dispatch_frame(frame, packet.get_type(), &path); + dispatch_frame(frame, packet.get_type(), &path, key_generation); })?; space.journal.of_rcvd_packets().on_rcvd_pn( packet.pn(), diff --git a/qconnection/src/tls.rs b/qconnection/src/tls.rs index c3582c3f6..44c0c8dba 100644 --- a/qconnection/src/tls.rs +++ b/qconnection/src/tls.rs @@ -349,7 +349,7 @@ impl ServerTlsSession { match self.tls_conn.zero_rtt_keys() { Some(keys) => zero_rtt_keys.set_keys(keys.into()), - None => _ = zero_rtt_keys.invalid(), + None => _ = zero_rtt_keys.invalidate(), } Ok(()) diff --git a/qconnection/src/tx.rs b/qconnection/src/tx.rs index 8bebf6906..15cbe24ff 100644 --- a/qconnection/src/tx.rs +++ b/qconnection/src/tx.rs @@ -6,7 +6,7 @@ use qbase::{ packet::{ AssemblePacket, PacketProperties, PacketWriter as BasePacketWriter, RecordFrame, header::{EncodeHeader, GetType, io::WriteHeader, long::LongHeader, short::OneRttHeader}, - keys::DirectionalKeys, + keys::{ArcOneRttPacketKeys, DirectionalKeys}, signal::KeyPhaseBit, }, util::ContinuousData, @@ -21,8 +21,30 @@ pub struct PacketWriter<'b, 's, F> { writer: QEventPacketWriter<'b>, // 不同空间的send guard类型不一样 clerk: NewPacketGuard<'s, F>, - retran_timeout: Duration, - expire_timeout: Duration, + one_rtt_sent_record: Option, + timeouts: PacketTimeouts, +} + +pub struct PacketTimeouts { + pub retran_timeout: Duration, + pub expire_timeout: Duration, +} + +pub struct OneRttPacketMeta { + pub key_phase: KeyPhaseBit, + pub key_generation: u64, + pub packet_keys: ArcOneRttPacketKeys, +} + +pub struct OneRttPacketKeysMeta { + pub keys: DirectionalKeys, + pub meta: OneRttPacketMeta, +} + +struct OneRttSentRecord { + packet_keys: ArcOneRttPacketKeys, + key_generation: u64, + packet_number: u64, } impl<'b, F> AsRef> for PacketWriter<'b, '_, F> { @@ -58,27 +80,39 @@ impl<'b, 's, F> PacketWriter<'b, 's, F> { Ok(Self { clerk, writer: QEventPacketWriter::new_long(&header, buffer, pn, keys)?, - expire_timeout, - retran_timeout, + one_rtt_sent_record: None, + timeouts: PacketTimeouts { + retran_timeout, + expire_timeout, + }, }) } pub fn new_short( header: OneRttHeader, buffer: &'b mut [u8], - keys: DirectionalKeys, - key_phase: KeyPhaseBit, journal: &'s ArcSentJournal, - retran_timeout: Duration, - expire_timeout: Duration, + one_rtt: OneRttPacketKeysMeta, + timeouts: PacketTimeouts, ) -> Result { + let OneRttPacketKeysMeta { keys, meta } = one_rtt; + let OneRttPacketMeta { + key_phase, + key_generation, + packet_keys, + } = meta; let clerk = journal.new_packet(); let pn = clerk.pn(); + let actual_pn = pn.0; Ok(Self { clerk, writer: QEventPacketWriter::new_short(&header, buffer, pn, keys, key_phase)?, - expire_timeout, - retran_timeout, + one_rtt_sent_record: Some(OneRttSentRecord { + packet_keys, + key_generation, + packet_number: actual_pn, + }), + timeouts, }) } } @@ -109,8 +143,13 @@ unsafe impl<'b, 's, F> BufMut for PacketWriter<'b, 's, F> { impl AssemblePacket for PacketWriter<'_, '_, F> { #[inline] fn encrypt_and_protect_packet(self) -> (usize, PacketProperties) { + if let Some(record) = &self.one_rtt_sent_record { + record + .packet_keys + .on_packet_sent(record.key_generation, record.packet_number); + } self.clerk - .build_with_time(self.retran_timeout, self.expire_timeout); + .build_with_time(self.timeouts.retran_timeout, self.timeouts.expire_timeout); self.writer.encrypt_and_protect_packet() } } @@ -138,6 +177,7 @@ pub struct TrivialPacketWriter<'b, 's, F> { writer: QEventPacketWriter<'b>, // 不同空间的send guard类型不一样 clerk: NewPacketGuard<'s, F>, + one_rtt_sent_record: Option, } impl<'b, F> AsRef> for TrivialPacketWriter<'b, '_, F> { @@ -172,6 +212,7 @@ impl<'b, 's, F> TrivialPacketWriter<'b, 's, F> { Ok(Self { clerk, writer: QEventPacketWriter::new_long(&header, buffer, pn, keys)?, + one_rtt_sent_record: None, }) } @@ -179,15 +220,26 @@ impl<'b, 's, F> TrivialPacketWriter<'b, 's, F> { pub fn new_short( header: OneRttHeader, buffer: &'b mut [u8], - keys: DirectionalKeys, - key_phase: KeyPhaseBit, + one_rtt: OneRttPacketKeysMeta, journal: &'s ArcSentJournal, ) -> Result { + let OneRttPacketKeysMeta { keys, meta } = one_rtt; + let OneRttPacketMeta { + key_phase, + key_generation, + packet_keys, + } = meta; let clerk = journal.new_packet(); let pn = clerk.pn(); + let actual_pn = pn.0; Ok(Self { clerk, writer: QEventPacketWriter::new_short(&header, buffer, pn, keys, key_phase)?, + one_rtt_sent_record: Some(OneRttSentRecord { + packet_keys, + key_generation, + packet_number: actual_pn, + }), }) } } @@ -218,6 +270,11 @@ unsafe impl<'b, 's, F> BufMut for TrivialPacketWriter<'b, 's, F> { impl AssemblePacket for TrivialPacketWriter<'_, '_, F> { #[inline] fn encrypt_and_protect_packet(self) -> (usize, PacketProperties) { + if let Some(record) = &self.one_rtt_sent_record { + record + .packet_keys + .on_packet_sent(record.key_generation, record.packet_number); + } self.clerk.build_trivial(); self.writer.encrypt_and_protect_packet() } diff --git a/qinterface/src/component/route/packet.rs b/qinterface/src/component/route/packet.rs index 3a577cea2..ef5a41b9b 100644 --- a/qinterface/src/component/route/packet.rs +++ b/qinterface/src/component/route/packet.rs @@ -7,7 +7,7 @@ use qbase::{ decrypt_packet, remove_protection_of_long_packet, remove_protection_of_short_packet, }, header::long::InitialHeader, - keys::ArcOneRttPacketKeys, + keys::{ArcOneRttPacketKeys, DecryptAttemptResult}, number::{InvalidPacketNumber, PacketNumber}, }, }; @@ -158,6 +158,7 @@ where undecoded_pn, decoded_pn, body_len: body_length, + key_generation: None, })) } @@ -180,6 +181,7 @@ where return Some(Err(invalid_reverse_bits.into())); } }; + let decoded_pn = match pn_decoder(undecoded_pn) { Ok(pn) => pn, Err(invalid_pn) => { @@ -187,24 +189,88 @@ where return None; } }; - let pk = pk.lock_guard().get_remote(key_phase, decoded_pn); + let body_offset = self.payload_offset + undecoded_pn.size(); - let body_length = match decrypt_packet(pk.as_ref(), decoded_pn, pkt_buf, body_offset) { - Ok(body_length) => body_length, - Err(error) => { - self.drop_on_decryption_failure(error, decoded_pn); - return None; + + // Get candidate keys outside the decryption loop + let (candidates, next_generation) = + pk.get_candidate_keys_for_decryption(decoded_pn, key_phase); + + // Try decrypting with each candidate key + for (generation, key) in candidates { + let decrypt_result = + decrypt_packet(key.as_ref(), decoded_pn, pkt_buf, body_offset).map_err(|_| ()); + let attempt = pk.record_decrypt_attempt(generation, decoded_pn, decrypt_result); + + match attempt { + Ok(DecryptAttemptResult::Success(body_length)) => { + return Some(Ok(PlainPacket { + header: self.header, + plain: self.payload.freeze(), + payload_offset: self.payload_offset, + undecoded_pn, + decoded_pn, + body_len: body_length, + key_generation: Some(generation), + })); + } + Ok(DecryptAttemptResult::GiveUp) => { + self.drop_on_decryption_failure( + qbase::packet::error::Error::DecryptPacketFailure, + decoded_pn, + ); + return None; + } + Ok(DecryptAttemptResult::ContinueTrying) => continue, + Err(e) => return Some(Err(e)), } - }; + } - Some(Ok(PlainPacket { - header: self.header, - plain: self.payload.freeze(), - payload_offset: self.payload_offset, - undecoded_pn, + // Try next generation if available + if let Some(generation) = next_generation { + let next_key = match pk.resolve_remote_key_for_generation(generation) { + Ok(key) => key, + Err(e) => return Some(Err(e)), + }; + + if let Some(key) = next_key { + let decrypt_result = + decrypt_packet(key.as_ref(), decoded_pn, pkt_buf, body_offset).map_err(|_| ()); + let attempt = pk.record_decrypt_attempt(generation, decoded_pn, decrypt_result); + + match attempt { + Ok(DecryptAttemptResult::Success(body_length)) => { + return Some(Ok(PlainPacket { + header: self.header, + plain: self.payload.freeze(), + payload_offset: self.payload_offset, + undecoded_pn, + decoded_pn, + body_len: body_length, + key_generation: Some(generation), + })); + } + Ok(DecryptAttemptResult::GiveUp) => { + self.drop_on_decryption_failure( + qbase::packet::error::Error::DecryptPacketFailure, + decoded_pn, + ); + return None; + } + Ok(DecryptAttemptResult::ContinueTrying) => { + // Continue to final failure + } + Err(e) => return Some(Err(e)), + } + } + } + + // All decryption attempts failed + self.drop_on_decryption_failure( + qbase::packet::error::Error::DecryptPacketFailure, decoded_pn, - body_len: body_length, - })) + ); + None } } @@ -232,6 +298,7 @@ pub struct PlainPacket { plain: Bytes, payload_offset: usize, body_len: usize, + key_generation: Option, } impl PlainPacket { @@ -247,6 +314,10 @@ impl PlainPacket { self.undecoded_pn.size() + self.body_len } + pub fn key_generation(&self) -> Option { + self.key_generation + } + pub fn body(&self) -> Bytes { let packet_offset = self.payload_offset + self.undecoded_pn.size(); self.plain