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
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238

// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

mod sliding_window;
pub use sliding_window::{SlidingWindow, SlidingWindowError};

mod protected_packet_number;
pub use protected_packet_number::ProtectedPacketNumber;

//= https://www.rfc-editor.org/rfc/rfc9000#section-12.3
//# The packet number is an integer in the range 0 to 2^62-1.  This
//# number is used in determining the cryptographic nonce for packet
//# protection.  Each endpoint maintains a separate packet number for
//# sending and receiving.

use crate::varint::VarInt;

mod packet_number;
pub use packet_number::PacketNumber;

//= https://www.rfc-editor.org/rfc/rfc9000#section-12.3
//# Packet numbers are limited to this range because they need to be
//# representable in whole in the Largest Acknowledged field of an ACK
//# frame (Section 19.3).  When present in a long or short header
//# however, packet numbers are reduced and encoded in 1 to 4 bytes; see
//# Section 17.1.

mod truncated_packet_number;
pub use truncated_packet_number::TruncatedPacketNumber;

//= https://www.rfc-editor.org/rfc/rfc9000#section-12.3
//# Initial space:  All Initial packets (Section 17.2.2) are in this
//# space.
//#
//# Handshake space:  All Handshake packets (Section 17.2.4) are in this
//# space.
//#
//# Application data space:  All 0-RTT (Section 17.2.3) and 1-RTT
//# (Section 17.3.1) packets are in this space.

mod packet_number_space;
pub use packet_number_space::PacketNumberSpace;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.1
//# Packet numbers are integers in the range 0 to 2^62-1 (Section 12.3).
//# When present in long or short packet headers, they are encoded in 1
//# to 4 bytes.  The number of bits required to represent the packet
//# number is reduced by including only the least significant bits of the
//# packet number.

/// The packet number len is the two least significant bits of the packet tag
pub(crate) const PACKET_NUMBER_LEN_MASK: u8 = 0b11;

mod packet_number_len;
pub use packet_number_len::PacketNumberLen;

mod packet_number_range;
pub use packet_number_range::PacketNumberRange;

#[cfg(feature = "alloc")]
pub mod map;
#[cfg(feature = "alloc")]
pub use map::Map;

#[cfg(any(test, feature = "testing"))]
pub mod testing;
#[cfg(test)]
mod tests;

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.1
//# the sender MUST use a packet number size able to represent more than
//# twice as large a range as the difference between the largest
//# acknowledged packet number and the packet number being sent.  A peer
//# receiving the packet will then correctly decode the packet number,
//# unless the packet is delayed in transit such that it arrives after
//# many higher-numbered packets have been received.  An endpoint SHOULD
//# use a large enough packet number encoding to allow the packet number
//# to be recovered even if the packet arrives after packets that are
//# sent afterwards.

fn derive_truncation_range(
    largest_acknowledged_packet_number: PacketNumber,
    packet_number: PacketNumber,
) -> Option<PacketNumberLen> {
    let space = packet_number.space();
    space.assert_eq(largest_acknowledged_packet_number.space());
    packet_number
        .as_u64()
        .checked_sub(largest_acknowledged_packet_number.as_u64())
        .and_then(|value| value.checked_mul(2))
        .and_then(|value| VarInt::new(value).ok())
        .and_then(|value| PacketNumberLen::from_varint(value, space))
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.1
//# As a result, the size of the packet number encoding is at least one
//# bit more than the base-2 logarithm of the number of contiguous
//# unacknowledged packet numbers, including the new packet.

//= https://www.rfc-editor.org/rfc/rfc9000#appendix-A.2
//# For example, if an endpoint has received an acknowledgment for packet
//# 0xabe8b3 and is sending a packet with a number of 0xac5c02, there are
//# 29,519 (0x734f) outstanding packet numbers.  In order to represent at
//# least twice this range (59,038 packets, or 0xe69e), 16 bits are
//# required.
//#
//# In the same state, sending a packet with a number of 0xace8fe uses
//# the 24-bit encoding, because at least 18 bits are required to
//# represent twice the range (131,222 packets, or 0x020096).

#[test]
fn packet_number_len_example_test() {
    let largest_acknowledged_packet_number =
        PacketNumberSpace::default().new_packet_number(VarInt::from_u32(0x00ab_e8bc));

    assert_eq!(
        PacketNumberSpace::default()
            .new_packet_number(VarInt::from_u32(0x00ac_5c02))
            .truncate(largest_acknowledged_packet_number)
            .unwrap()
            .bitsize(),
        16,
    );

    assert_eq!(
        PacketNumberSpace::default()
            .new_packet_number(VarInt::from_u32(0x00ac_e8fe))
            .truncate(largest_acknowledged_packet_number)
            .unwrap()
            .bitsize(),
        24,
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.1
//# At a receiver, protection of the packet number is removed prior to
//# recovering the full packet number.  The full packet number is then
//# reconstructed based on the number of significant bits present, the
//# value of those bits, and the largest packet number received in a
//# successfully authenticated packet.  Recovering the full packet number
//# is necessary to successfully complete the removal of packet
//# protection.
//
//# Once header protection is removed, the packet number is decoded by
//# finding the packet number value that is closest to the next expected
//# packet.  The next expected packet is the highest received packet
//# number plus one.  Pseudocode and an example for packet number
//# decoding can be found in Appendix A.3.

#[test]
fn packet_decoding_example_test() {
    let space = PacketNumberSpace::default();
    let largest_packet_number = space.new_packet_number(VarInt::from_u32(0xa82f_30ea));
    let truncated_packet_number = TruncatedPacketNumber::new(0x9b32u16, space);
    let expected = space.new_packet_number(VarInt::from_u32(0xa82f_9b32));
    let actual = decode_packet_number(largest_packet_number, truncated_packet_number);
    assert_eq!(actual, expected);
    assert_eq!(
        expected.truncate(largest_packet_number).unwrap(),
        truncated_packet_number
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#appendix-A.3
//# DecodePacketNumber(largest_pn, truncated_pn, pn_nbits):
//#   expected_pn  = largest_pn + 1
//#   pn_win       = 1 << pn_nbits
//#   pn_hwin      = pn_win / 2
//#   pn_mask      = pn_win - 1
//#   // The incoming packet number should be greater than
//#   // expected_pn - pn_hwin and less than or equal to
//#   // expected_pn + pn_hwin
//#   //
//#   // This means we cannot just strip the trailing bits from
//#   // expected_pn and add the truncated_pn because that might
//#   // yield a value outside the window.
//#   //
//#   // The following code calculates a candidate value and
//#   // makes sure it's within the packet number window.
//#   // Note the extra checks to prevent overflow and underflow.
//#   candidate_pn = (expected_pn & ~pn_mask) | truncated_pn
//#   if candidate_pn <= expected_pn - pn_hwin and
//#     candidate_pn < (1 << 62) - pn_win:
//#     return candidate_pn + pn_win
//#   if candidate_pn > expected_pn + pn_hwin and
//#     candidate_pn >= pn_win:
//#     return candidate_pn - pn_win
//#   return candidate_pn

#[inline(never)] // prevent the compiler from optimizing call size values and making this non-constant time
fn decode_packet_number(
    largest_pn: PacketNumber,
    truncated_pn: TruncatedPacketNumber,
) -> PacketNumber {
    let space = largest_pn.space();
    space.assert_eq(truncated_pn.space());

    let pn_nbits = truncated_pn.bitsize();
    // deref to u64 so we have enough room
    let expected_pn = largest_pn.as_u64() + 1;
    let pn_win = 1 << pn_nbits;
    let pn_hwin = pn_win / 2;
    let pn_mask = pn_win - 1;
    let mut candidate_pn = (expected_pn & !pn_mask) | truncated_pn.into_u64();

    let a = expected_pn
        .checked_sub(pn_hwin)
        .filter(|v| candidate_pn <= *v)
        .is_some();
    let b = (1u64 << 62)
        .checked_sub(pn_win)
        .filter(|v| candidate_pn < *v)
        .is_some();
    let c = expected_pn
        .checked_add(pn_hwin)
        .filter(|v| candidate_pn > *v)
        .is_some();
    let d = candidate_pn >= pn_win;

    let ab = a && b;
    let cd = !ab && c && d;

    // Make sure the compiler doesn't try and optimize the if statements
    // See https://godbolt.org/z/348WbbbzM
    core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);

    if ab {
        candidate_pn += pn_win;
    }

    if cd {
        candidate_pn -= pn_win
    }

    let candidate_pn = VarInt::new(candidate_pn).unwrap_or(VarInt::MAX);

    PacketNumber::from_varint(candidate_pn, space)
}