dkls23/proto/
encrypted.rs

1// Copyright (c) Silence Laboratories Pte. Ltd. All Rights Reserved.
2// This software is licensed under the Silence Laboratories License Agreement.
3
4//! Module for handling encrypted messages in the protocol.
5//! This module provides functionality for encrypting and decrypting messages
6//! with support for additional data and trailers. It uses a pluggable encryption
7//! scheme interface to allow for different encryption implementations.
8
9use std::marker::PhantomData;
10
11use bytemuck::{AnyBitPattern, NoUninit};
12use chacha20poly1305::ChaCha20Poly1305;
13
14use sl_mpc_mate::message::*;
15
16pub use crate::proto::scheme::EncryptionScheme;
17
18/// Default encryption scheme using X25519 key exchange and ChaCha20Poly1305 for encryption.
19pub type Scheme = crate::proto::scheme::AeadX25519<ChaCha20Poly1305>;
20
21/// A wrapper for a message of type T with support for in-place encryption/decryption.
22///
23/// This struct provides functionality for encrypting and decrypting messages while
24/// maintaining a specific format:
25///
26/// ```text
27/// [ msg-hdr | additional-data | payload | trailer | tag + nonce ]
28/// ```
29///
30/// Where:
31/// - `msg-hdr`: Message header containing ID, TTL, and flags
32/// - `additional-data`: Optional unencrypted data
33/// - `payload`: The encrypted external representation of type T
34/// - `trailer`: Optional encrypted variable-sized data
35/// - `tag + nonce`: Authentication tag and nonce for the encryption scheme
36///
37/// The `payload` and `trailer` sections are encrypted, while the header and
38/// additional data remain in plaintext.
39pub struct EncryptedMessage<T> {
40    buffer: Vec<u8>,
41    additional_data: usize, // size of additional-data
42    marker: PhantomData<T>,
43}
44
45impl<T: AnyBitPattern + NoUninit> EncryptedMessage<T> {
46    const T_SIZE: usize = core::mem::size_of::<T>();
47
48    /// Calculates the total size of an encrypted message.
49    ///
50    /// # Arguments
51    /// * `ad` - Size of additional data in bytes
52    /// * `trailer` - Size of trailer data in bytes
53    /// * `scheme` - The encryption scheme to use
54    ///
55    /// # Returns
56    /// The total size in bytes needed for the encrypted message
57    pub fn size(
58        ad: usize,
59        trailer: usize,
60        scheme: &dyn EncryptionScheme,
61    ) -> usize {
62        MESSAGE_HEADER_SIZE + ad + Self::T_SIZE + trailer + scheme.overhead()
63    }
64
65    /// Creates a new encrypted message with the specified parameters.
66    ///
67    /// # Arguments
68    /// * `id` - Message identifier
69    /// * `ttl` - Time-to-live value
70    /// * `flags` - Message flags
71    /// * `trailer` - Size of trailer data in bytes
72    /// * `scheme` - The encryption scheme to use
73    ///
74    /// # Returns
75    /// A new `EncryptedMessage` instance
76    pub fn new(
77        id: &MsgId,
78        ttl: u32,
79        flags: u16,
80        trailer: usize,
81        scheme: &dyn EncryptionScheme,
82    ) -> Self {
83        let buffer = vec![0u8; Self::size(0, trailer, scheme)];
84
85        Self::from_buffer(buffer, id, ttl, flags, 0, trailer, scheme)
86    }
87
88    /// Creates a new encrypted message with additional data.
89    ///
90    /// # Arguments
91    /// * `id` - Message identifier
92    /// * `ttl` - Time-to-live value
93    /// * `flags` - Message flags
94    /// * `additional_data` - Size of additional data in bytes
95    /// * `trailer` - Size of trailer data in bytes
96    /// * `scheme` - The encryption scheme to use
97    ///
98    /// # Returns
99    /// A new `EncryptedMessage` instance with space for additional data
100    pub fn new_with_ad(
101        id: &MsgId,
102        ttl: u32,
103        flags: u16,
104        additional_data: usize,
105        trailer: usize,
106        scheme: &dyn EncryptionScheme,
107    ) -> Self {
108        let buffer = vec![0u8; Self::size(additional_data, trailer, scheme)];
109
110        Self::from_buffer(
111            buffer,
112            id,
113            ttl,
114            flags,
115            additional_data,
116            trailer,
117            scheme,
118        )
119    }
120
121    /// Creates an encrypted message from an existing buffer.
122    ///
123    /// # Arguments
124    /// * `buffer` - Existing buffer to use
125    /// * `id` - Message identifier
126    /// * `ttl` - Time-to-live value
127    /// * `flags` - Message flags
128    /// * `additional_data` - Size of additional data in bytes
129    /// * `trailer` - Size of trailer data in bytes
130    /// * `scheme` - The encryption scheme to use
131    ///
132    /// # Returns
133    /// A new `EncryptedMessage` instance using the provided buffer
134    pub fn from_buffer(
135        mut buffer: Vec<u8>,
136        id: &MsgId,
137        ttl: u32,
138        flags: u16,
139        additional_data: usize,
140        trailer: usize,
141        scheme: &dyn EncryptionScheme,
142    ) -> Self {
143        buffer.resize(Self::size(additional_data, trailer, scheme), 0);
144
145        if let Some(hdr) = buffer.first_chunk_mut::<MESSAGE_HEADER_SIZE>() {
146            MsgHdr::encode(hdr, id, ttl, flags);
147        }
148
149        Self {
150            buffer,
151            additional_data,
152            marker: PhantomData,
153        }
154    }
155
156    /// Returns mutable references to the message payload, trailer, and additional data.
157    ///
158    /// # Arguments
159    /// * `scheme` - The encryption scheme to use
160    ///
161    /// # Returns
162    /// A tuple containing:
163    /// - Mutable reference to the payload object
164    /// - Mutable reference to the trailer bytes
165    /// - Mutable reference to the additional data bytes
166    pub fn payload_with_ad(
167        &mut self,
168        scheme: &dyn EncryptionScheme,
169    ) -> (&mut T, &mut [u8], &mut [u8]) {
170        let tag_offset = self.buffer.len() - scheme.overhead();
171
172        // body = ad | payload | trailer
173        let body = &mut self.buffer[MESSAGE_HEADER_SIZE..tag_offset];
174
175        let (additional_data, msg_and_trailer) =
176            body.split_at_mut(self.additional_data);
177
178        let (msg, trailer) = msg_and_trailer.split_at_mut(Self::T_SIZE);
179
180        (bytemuck::from_bytes_mut(msg), trailer, additional_data)
181    }
182
183    /// Returns mutable references to the message payload and trailer.
184    ///
185    /// # Arguments
186    /// * `scheme` - The encryption scheme to use
187    ///
188    /// # Returns
189    /// A tuple containing:
190    /// - Mutable reference to the payload object
191    /// - Mutable reference to the trailer bytes
192    pub fn payload(
193        &mut self,
194        scheme: &dyn EncryptionScheme,
195    ) -> (&mut T, &mut [u8]) {
196        let (msg, trailer, _) = self.payload_with_ad(scheme);
197
198        (msg, trailer)
199    }
200
201    /// Encrypts the message using the provided encryption scheme.
202    ///
203    /// # Arguments
204    /// * `scheme` - The encryption scheme to use
205    /// * `receiver` - The ID of the intended receiver
206    ///
207    /// # Returns
208    /// The encrypted message as a byte vector, or `None` if encryption failed
209    pub fn encrypt(
210        self,
211        scheme: &mut dyn EncryptionScheme,
212        receiver: usize,
213    ) -> Option<Vec<u8>> {
214        let mut buffer = self.buffer;
215
216        let last = buffer.len() - scheme.overhead();
217        let (msg, tail) = buffer.split_at_mut(last);
218
219        let (associated_data, plaintext) =
220            msg.split_at_mut(MESSAGE_HEADER_SIZE + self.additional_data);
221
222        scheme
223            .encrypt(associated_data, plaintext, tail, receiver)
224            .ok()?;
225
226        Some(buffer)
227    }
228
229    /// Decrypts a message and returns references to the payload, trailer, and additional data.
230    ///
231    /// # Arguments
232    /// * `buffer` - The encrypted message buffer
233    /// * `additional_data` - Size of additional data in bytes
234    /// * `trailer` - Size of trailer data in bytes
235    /// * `scheme` - The encryption scheme to use
236    /// * `sender` - The ID of the message sender
237    ///
238    /// # Returns
239    /// A tuple containing references to the decrypted payload, trailer, and additional data,
240    /// or `None` if decryption failed
241    pub fn decrypt_with_ad<'msg>(
242        buffer: &'msg mut [u8],
243        additional_data: usize,
244        trailer: usize,
245        scheme: &dyn EncryptionScheme,
246        sender: usize,
247    ) -> Option<(&'msg T, &'msg [u8], &'msg [u8])> {
248        if buffer.len() != Self::size(additional_data, trailer, scheme) {
249            return None;
250        }
251
252        let (associated_data, body) =
253            buffer.split_at_mut(MESSAGE_HEADER_SIZE + additional_data);
254
255        let (ciphertext, tail) =
256            body.split_at_mut(body.len() - scheme.overhead());
257
258        scheme
259            .decrypt(associated_data, ciphertext, tail, sender)
260            .ok()?;
261
262        let (msg, trailer) = ciphertext.split_at_mut(Self::T_SIZE);
263
264        Some((
265            bytemuck::from_bytes_mut(msg),
266            trailer,
267            &associated_data[MESSAGE_HEADER_SIZE..],
268        ))
269    }
270
271    /// Decrypts a message and returns references to the payload and trailer.
272    ///
273    /// # Arguments
274    /// * `buffer` - The encrypted message buffer
275    /// * `trailer` - Size of trailer data in bytes
276    /// * `scheme` - The encryption scheme to use
277    /// * `sender` - The ID of the message sender
278    ///
279    /// # Returns
280    /// A tuple containing references to the decrypted payload and trailer,
281    /// or `None` if decryption failed
282    pub fn decrypt<'msg>(
283        buffer: &'msg mut [u8],
284        trailer: usize,
285        scheme: &dyn EncryptionScheme,
286        sender: usize,
287    ) -> Option<(&'msg T, &'msg [u8])> {
288        Self::decrypt_with_ad(buffer, 0, trailer, scheme, sender)
289            .map(|(msg, trailer, _)| (msg, trailer))
290    }
291}
dkls23/proto/encrypted.rs

dkls23/proto/
encrypted.rs
