pb_mapper_protocol/secure/
frame.rs

1//! Protocol-v2 key derivation and directional authenticated frame codecs.
2//!
3//! ```text
4//! credential + connection salt -> HKDF -> c2s key / s2c key
5//! plaintext -> counter + length + AEAD(AAD) -> encrypted frame
6//! encrypted frame -> bound length -> verify counter/tag -> plaintext
7//! ```
8//!
9//! Counters are monotonic per direction and are included in both the nonce and AAD.
10//! The initial reader can impose a smaller pre-authentication limit before allocating
11//! a body; continuation frames retain the normal protocol maximum.
12
13use pb_mapper_auth::KeyId;
14
15use super::*;
16
17#[derive(Clone)]
18pub(super) struct V2Material {
19    pub(super) key_id: KeyId,
20    pub(super) flags: u8,
21    pub(super) salt: [u8; CONNECTION_SALT_LEN],
22    pub(super) client_to_server: AesKeyType,
23    pub(super) server_to_client: AesKeyType,
24}
25
26pub struct V2MessageReader<'a, T: AsyncReadExt + Unpin> {
27    reader: &'a mut T,
28    material: V2Material,
29    key: LessSafeKey,
30    direction: u8,
31    expected_counter: u64,
32    buffer: Vec<u8>,
33    frame: crate::frame_read::FrameRead<FRAME_HEADER_LEN>,
34}
35
36impl<'a, T: AsyncReadExt + Unpin> V2MessageReader<'a, T> {
37    pub(super) fn new(
38        reader: &'a mut T,
39        material: V2Material,
40        direction: u8,
41        expected_counter: u64,
42    ) -> Result<Self> {
43        let key_bytes = direction_key(&material, direction);
44        let key = LessSafeKey::new(
45            UnboundKey::new(&AES_256_GCM, key_bytes)
46                .map_err(|_| protocol_error("invalid protocol-v2 read key"))?,
47        );
48        Ok(Self {
49            reader,
50            material,
51            key,
52            direction,
53            expected_counter,
54            buffer: Vec::new(),
55            frame: crate::frame_read::FrameRead::new(),
56        })
57    }
58
59    pub(super) async fn read_msg_with_limit(&mut self, max_plaintext_len: u32) -> Result<&'_ [u8]> {
60        self.frame
61            .header(self.reader)
62            .await
63            .map_err(|error| protocol_error(format!("failed to read v2 header: {error}")))?;
64        let (counter, datalen) =
65            validate_frame_header(&self.frame.header, self.expected_counter, max_plaintext_len)?;
66        self.buffer.resize(datalen as usize, 0);
67        self.frame
68            .body(self.reader, &mut self.buffer)
69            .await
70            .map_err(|error| protocol_error(format!("failed to read v2 payload: {error}")))?;
71        let datalen = u32::try_from(self.buffer.len())
72            .map_err(|_| protocol_error("protocol-v2 payload exceeds u32 length"))?;
73        let aad = frame_aad(&self.material, self.direction, counter, datalen);
74        let plain = self
75            .key
76            .open_in_place(nonce(counter), Aad::from(aad.as_slice()), &mut self.buffer)
77            .map_err(|_| protocol_error("protocol-v2 payload authentication failed"))?;
78        let plain_len = plain.len();
79        self.buffer.truncate(plain_len);
80        self.expected_counter = self
81            .expected_counter
82            .checked_add(1)
83            .ok_or_else(|| protocol_error("protocol-v2 receive counter exhausted"))?;
84        self.frame.finish();
85        Ok(&self.buffer)
86    }
87}
88
89pub(super) async fn read_v2_frame<T: AsyncReadExt + Unpin>(
90    reader: &mut T,
91    expected_counter: u64,
92    max_plaintext_len: u32,
93) -> Result<(u64, Vec<u8>)> {
94    let mut frame = crate::frame_read::FrameRead::<FRAME_HEADER_LEN>::new();
95    frame
96        .header(reader)
97        .await
98        .map_err(|error| protocol_error(format!("failed to read v2 header: {error}")))?;
99    let (counter, datalen) =
100        validate_frame_header(&frame.header, expected_counter, max_plaintext_len)?;
101    let mut ciphertext = vec![0_u8; datalen as usize];
102    reader
103        .read_exact(&mut ciphertext)
104        .await
105        .map_err(|error| protocol_error(format!("failed to read v2 payload: {error}")))?;
106    Ok((counter, ciphertext))
107}
108
109fn validate_frame_header(
110    header: &[u8; FRAME_HEADER_LEN],
111    expected_counter: u64,
112    max_plaintext_len: u32,
113) -> Result<(u64, u32)> {
114    let counter = u64::from_be_bytes([
115        header[0], header[1], header[2], header[3], header[4], header[5], header[6], header[7],
116    ]);
117    if counter != expected_counter {
118        return Err(protocol_error(format!(
119            "protocol-v2 counter mismatch: expected {expected_counter}, got {counter}"
120        )));
121    }
122    let datalen = u32::from_be_bytes([header[8], header[9], header[10], header[11]]);
123    let max_encrypted_len = max_plaintext_len.saturating_add(AES_256_GCM.tag_len() as u32);
124    if datalen < AES_256_GCM.tag_len() as u32 || datalen > max_encrypted_len {
125        return Err(protocol_error(format!(
126            "protocol-v2 payload length {datalen} exceeds the {max_plaintext_len}-byte limit"
127        )));
128    }
129    Ok((counter, datalen))
130}
131
132impl<T: AsyncReadExt + Unpin> MessageReader for V2MessageReader<'_, T> {
133    async fn read_msg(&mut self) -> Result<&'_ [u8]> {
134        self.read_msg_with_limit(MAX_MSG_LEN - AES_256_GCM.tag_len() as u32)
135            .await
136    }
137}
138
139pub struct V2MessageWriter<'a, T: AsyncWriteExt + Unpin> {
140    writer: &'a mut T,
141    material: V2Material,
142    key: LessSafeKey,
143    direction: u8,
144    counter: u64,
145    poisoned: bool,
146}
147
148impl<'a, T: AsyncWriteExt + Unpin> V2MessageWriter<'a, T> {
149    pub(super) fn new(
150        writer: &'a mut T,
151        material: V2Material,
152        direction: u8,
153        counter: u64,
154    ) -> Result<Self> {
155        let key_bytes = direction_key(&material, direction);
156        let key = LessSafeKey::new(
157            UnboundKey::new(&AES_256_GCM, key_bytes)
158                .map_err(|_| protocol_error("invalid protocol-v2 write key"))?,
159        );
160        Ok(Self {
161            writer,
162            material,
163            key,
164            direction,
165            counter,
166            poisoned: false,
167        })
168    }
169}
170
171impl<T: AsyncWriteExt + Unpin> MessageWriter for V2MessageWriter<'_, T> {
172    async fn write_msg(&mut self, message: &[u8]) -> Result<()> {
173        if self.poisoned {
174            return Err(protocol_error(
175                "protocol-v2 writer interrupted; close the connection",
176            ));
177        }
178        let encrypted_len = message
179            .len()
180            .checked_add(AES_256_GCM.tag_len())
181            .and_then(|len| DataLenType::try_from(len).ok())
182            .ok_or_else(|| protocol_error("protocol-v2 message is too large"))?;
183        if encrypted_len > MAX_MSG_LEN {
184            return Err(protocol_error(
185                "protocol-v2 message exceeds the maximum length",
186            ));
187        }
188        let counter = self.counter;
189        // Reserve before sealing or awaiting I/O. Cancellation leaves a
190        // terminal writer; restarting a frame would desynchronize the stream.
191        self.counter = self
192            .counter
193            .checked_add(1)
194            .ok_or_else(|| protocol_error("protocol-v2 send counter exhausted"))?;
195        self.poisoned = true;
196        let aad = frame_aad(&self.material, self.direction, counter, encrypted_len);
197        let mut encrypted = message.to_vec();
198        self.key
199            .seal_in_place_append_tag(nonce(counter), Aad::from(aad.as_slice()), &mut encrypted)
200            .map_err(|_| protocol_error("failed to encrypt protocol-v2 message"))?;
201        self.writer
202            .write_u64(counter)
203            .await
204            .map_err(|error| protocol_error(format!("failed to write v2 frame header: {error}")))?;
205        self.writer
206            .write_u32(encrypted_len)
207            .await
208            .map_err(|error| protocol_error(format!("failed to write v2 frame header: {error}")))?;
209        self.writer
210            .write_all(&encrypted)
211            .await
212            .map_err(|error| protocol_error(format!("failed to write v2 frame body: {error}")))?;
213        self.poisoned = false;
214        Ok(())
215    }
216}
217
218pub(super) fn open_v2_payload(
219    material: &V2Material,
220    direction: u8,
221    counter: u64,
222    ciphertext: &mut [u8],
223) -> Result<Vec<u8>> {
224    let key_bytes = direction_key(material, direction);
225    let key = LessSafeKey::new(
226        UnboundKey::new(&AES_256_GCM, key_bytes)
227            .map_err(|_| protocol_error("invalid protocol-v2 read key"))?,
228    );
229    let datalen = u32::try_from(ciphertext.len())
230        .map_err(|_| protocol_error("protocol-v2 payload length is invalid"))?;
231    let aad = frame_aad(material, direction, counter, datalen);
232    let plain = key
233        .open_in_place(nonce(counter), Aad::from(aad.as_slice()), ciphertext)
234        .map_err(|_| protocol_error("protocol-v2 payload authentication failed"))?;
235    Ok(plain.to_vec())
236}
237
238pub(super) fn derive_material(
239    key_id: KeyId,
240    credential_key: &AesKeyType,
241    salt_bytes: [u8; CONNECTION_SALT_LEN],
242) -> Result<V2Material> {
243    let salt = Salt::new(HKDF_SHA256, &salt_bytes);
244    let pseudo_random_key = salt.extract(credential_key);
245    let client_to_server = expand_direction(&pseudo_random_key, b"pb-mapper-v2-c2s")?;
246    let server_to_client = expand_direction(&pseudo_random_key, b"pb-mapper-v2-s2c")?;
247    Ok(V2Material {
248        key_id,
249        flags: 0,
250        salt: salt_bytes,
251        client_to_server,
252        server_to_client,
253    })
254}
255
256fn expand_direction(
257    pseudo_random_key: &ring::hkdf::Prk,
258    label: &'static [u8],
259) -> Result<AesKeyType> {
260    let info = [label];
261    let output = pseudo_random_key
262        .expand(&info, HkdfLen(32))
263        .map_err(|_| protocol_error("failed to derive protocol-v2 direction key"))?;
264    let mut key = [0_u8; 32];
265    output
266        .fill(&mut key)
267        .map_err(|_| protocol_error("failed to fill protocol-v2 direction key"))?;
268    Ok(key)
269}
270
271struct HkdfLen(usize);
272
273impl ring::hkdf::KeyType for HkdfLen {
274    fn len(&self) -> usize {
275        self.0
276    }
277}
278
279fn direction_key(material: &V2Material, direction: u8) -> &AesKeyType {
280    if direction == DIRECTION_CLIENT_TO_SERVER {
281        &material.client_to_server
282    } else {
283        &material.server_to_client
284    }
285}
286
287pub(super) fn first_prefix(material: &V2Material) -> Vec<u8> {
288    let mut prefix = Vec::with_capacity(PROTOCOL_V2_MAGIC.len() + FIRST_PREFIX_REMAINDER_LEN);
289    prefix.extend_from_slice(&PROTOCOL_V2_MAGIC);
290    prefix.push(PROTOCOL_V2_VERSION);
291    prefix.push(material.flags);
292    prefix.extend_from_slice(&0_u16.to_be_bytes());
293    prefix.extend_from_slice(&material.key_id.to_be_bytes());
294    prefix.extend_from_slice(&material.salt);
295    prefix
296}
297
298fn frame_aad(material: &V2Material, direction: u8, counter: u64, datalen: u32) -> Vec<u8> {
299    let mut aad = Vec::with_capacity(
300        PROTOCOL_V2_MAGIC.len() + FIRST_PREFIX_REMAINDER_LEN + 1 + FRAME_HEADER_LEN,
301    );
302    aad.extend_from_slice(&first_prefix(material));
303    aad.push(direction);
304    aad.extend_from_slice(&counter.to_be_bytes());
305    aad.extend_from_slice(&datalen.to_be_bytes());
306    aad
307}
308
309fn nonce(counter: u64) -> Nonce {
310    let mut bytes = [0_u8; 12];
311    bytes[4..].copy_from_slice(&counter.to_be_bytes());
312    Nonce::assume_unique_for_key(bytes)
313}
pb_mapper_protocol/secure/frame.rs

pb_mapper_protocol/secure/
frame.rs
