1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use crate::credentials::{Credentials, KeyId};
use bitvec::BitArr;
use std::sync::{
atomic::{AtomicU64, Ordering},
Mutex,
};
const WINDOW: usize = 896;
type Seen = BitArr!(for WINDOW);
#[derive(Debug)]
pub struct State {
// This is the maximum ID we've seen so far. This is sent to peers for when we cannot determine
// if the packet sent is replayed as it falls outside our replay window. Peers use this
// information to resynchronize on the latest state.
max_seen_key_id: AtomicU64,
seen: Mutex<Seen>,
}
impl super::map::SizeOf for Mutex<Seen> {
fn size(&self) -> usize {
// If we don't need drop, it's very likely that this type is fully contained in size_of
// Self. This simplifies implementing this trait for e.g. std types.
//
// Mutex on macOS (at least) has a more expensive, pthread-based impl that allocates. But
// on Linux there's no extra allocation.
if cfg!(target_os = "linux") {
assert!(
!std::mem::needs_drop::<Self>(),
"{:?} requires custom SizeOf impl",
std::any::type_name::<Self>()
);
}
std::mem::size_of::<Self>()
}
}
impl super::map::SizeOf for State {
fn size(&self) -> usize {
let State {
max_seen_key_id,
seen,
} = self;
max_seen_key_id.size() + seen.size()
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
/// This indicates that we know about this element and it *definitely* already exists.
#[error("packet definitely already seen before")]
AlreadyExists,
/// We don't know whether we've seen this element before. It may or may not have already been
/// received.
#[error("packet may have been seen before")]
Unknown,
}
impl State {
pub fn new() -> State {
State {
max_seen_key_id: AtomicU64::new(u64::MAX),
seen: Default::default(),
}
}
pub fn pre_authentication(&self, identity: &Credentials) -> Result<(), Error> {
// Bail if we get the max key ID. This is not practically reachable on well-behaved senders
// (see sender.rs for comments), and lets us always return a valid KeyId from
// `minimum_unseen_key_id` even with non well-behaved peers.
if identity.key_id == KeyId::MAX {
return Err(Error::Unknown);
}
Ok(())
}
pub fn minimum_unseen_key_id(&self) -> KeyId {
KeyId::try_from(
self.max_seen_key_id
.load(Ordering::Relaxed)
// Initial u64::MAX wraps to zero, which is the correct answer for the initial
// state. After that just +1 consistently.
.wrapping_add(1),
)
.unwrap_or(
// Saturate if we've exhausted the key ID space. Should be unreachable in practice due
// to the pre_authentication check above, but avoid a panic by handling it here too.
KeyId::MAX,
)
}
/// Called after decryption has been performed
#[expect(
clippy::unwrap_used,
clippy::unwrap_in_result,
reason = "lock is only poisoned if another thread already panicked while holding it"
)]
pub fn post_authentication(&self, identity: &Credentials) -> Result<(), Error> {
// Duplicate since it's cheap right now, can be refined in the future.
// In practice callers should have already run this early in the receiving process.
self.pre_authentication(identity)?;
let mut seen = self.seen.lock().unwrap();
let key_id = *identity.key_id;
let mut previous_max = self.max_seen_key_id.load(Ordering::Relaxed);
let new_max = if previous_max == u64::MAX {
previous_max = 0;
key_id
} else {
previous_max.max(key_id)
};
self.max_seen_key_id.store(new_max, Ordering::Relaxed);
let delta = new_max - previous_max;
if delta > seen.len() as u64 {
// not yet seen since we shifted forward by more than the bitset's size.
seen.fill(false);
} else {
// Even on a 32-bit platform we'd hit the check above (since seen is way smaller than
// 2^32).
seen.shift_end(delta as usize);
}
let Ok(idx) = usize::try_from(new_max - key_id) else {
// We'd never store more than usize bits, so treat this as too old as well.
return Err(Error::Unknown);
};
let ret = if let Some(mut entry) = seen.get_mut(idx) {
if *entry {
return Err(Error::AlreadyExists);
}
entry.set(true);
Ok(())
} else {
// Too old -- no longer in memory.
return Err(Error::Unknown);
};
ret
}
}
impl Default for State {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests;