trussed_staging/chunked/
mod.rs

1// Copyright (C) Nitrokey GmbH
2// SPDX-License-Identifier: Apache-2.0 or MIT
3
4mod store;
5use store::OpenSeekFrom;
6
7use chacha20poly1305::{
8    aead,
9    aead::stream::{DecryptorLE31, EncryptorLE31, Nonce as StreamNonce, StreamLE31},
10    ChaCha8Poly1305, KeyInit,
11};
12use littlefs2_core::{Path, PathBuf};
13use rand_core::RngCore;
14use trussed::{
15    config::MAX_MESSAGE_LENGTH,
16    key::{Kind, Secrecy},
17    serde_extensions::ExtensionImpl,
18    service::ServiceResources,
19    store::{Filestore, Keystore, Store},
20    types::CoreContext,
21};
22use trussed_chunked::{
23    reply, ChunkedExtension, ChunkedReply, ChunkedRequest, CHACHA8_STREAM_NONCE_LEN,
24};
25use trussed_core::{
26    types::{Bytes, Location, Message},
27    Error,
28};
29
30use crate::StagingContext;
31
32const POLY1305_TAG_LEN: usize = 16;
33const CHACHA8_KEY_LEN: usize = 32;
34
35struct HeaplessBuffer<'a, LenT: heapless::LenType>(&'a mut heapless_bytes::BytesView<LenT>);
36
37impl<'a, LenT: heapless::LenType, S: heapless_bytes::BytesStorage + ?Sized>
38    From<&'a mut heapless_bytes::BytesInner<LenT, S>> for HeaplessBuffer<'a, LenT>
39{
40    fn from(value: &'a mut heapless_bytes::BytesInner<LenT, S>) -> Self {
41        Self(value.as_mut_view())
42    }
43}
44
45impl<'a, LenT: heapless::LenType> AsMut<[u8]> for HeaplessBuffer<'a, LenT> {
46    fn as_mut(&mut self) -> &mut [u8] {
47        self.0
48    }
49}
50
51impl<'a, LenT: heapless::LenType> AsRef<[u8]> for HeaplessBuffer<'a, LenT> {
52    fn as_ref(&self) -> &[u8] {
53        self.0
54    }
55}
56
57impl<'a, LenT: heapless::LenType> aead::Buffer for HeaplessBuffer<'a, LenT> {
58    fn extend_from_slice(&mut self, other: &[u8]) -> aead::Result<()> {
59        self.0.extend_from_slice(other).map_err(|_| aead::Error)
60    }
61
62    fn truncate(&mut self, len: usize) {
63        self.0.truncate(len);
64    }
65}
66
67#[derive(Debug)]
68pub struct ChunkedReadState {
69    pub path: PathBuf,
70    pub location: Location,
71    pub offset: usize,
72}
73
74#[derive(Debug)]
75pub struct ChunkedWriteState {
76    pub path: PathBuf,
77    pub location: Location,
78}
79
80pub struct EncryptedChunkedReadState {
81    pub path: PathBuf,
82    pub location: Location,
83    pub offset: usize,
84    pub decryptor: DecryptorLE31<ChaCha8Poly1305>,
85}
86
87pub struct EncryptedChunkedWriteState {
88    pub path: PathBuf,
89    pub location: Location,
90    pub encryptor: EncryptorLE31<ChaCha8Poly1305>,
91}
92
93#[non_exhaustive]
94pub enum ChunkedIoState {
95    Read(ChunkedReadState),
96    Write(ChunkedWriteState),
97    EncryptedRead(EncryptedChunkedReadState),
98    EncryptedWrite(EncryptedChunkedWriteState),
99}
100
101impl ExtensionImpl<ChunkedExtension> for super::StagingBackend {
102    fn extension_request<P: trussed::Platform>(
103        &mut self,
104        core_ctx: &mut CoreContext,
105        backend_ctx: &mut Self::Context,
106        request: &ChunkedRequest,
107        resources: &mut ServiceResources<P>,
108    ) -> Result<ChunkedReply, Error> {
109        let rng = &mut resources.rng()?;
110        let keystore = &mut resources.keystore(core_ctx.path.clone())?;
111        let filestore = &mut resources.filestore(core_ctx.path.clone());
112        let client_id = &core_ctx.path;
113        let store = resources.platform_mut().store();
114        match request {
115            ChunkedRequest::ReadChunk(_) => {
116                let read_state = match &mut backend_ctx.chunked_io_state {
117                    Some(ChunkedIoState::Read(read_state)) => read_state,
118                    Some(ChunkedIoState::EncryptedRead(_)) => {
119                        return read_encrypted_chunk(&store, client_id, backend_ctx)
120                    }
121                    _ => return Err(Error::MechanismNotAvailable),
122                };
123                let (data, len) = store::filestore_read_chunk(
124                    &store,
125                    client_id,
126                    &read_state.path,
127                    read_state.location,
128                    OpenSeekFrom::Start(read_state.offset as u32),
129                )?;
130
131                read_state.offset += data.len();
132
133                Ok(reply::ReadChunk { data, len }.into())
134            }
135            ChunkedRequest::StartChunkedRead(request) => {
136                clear_chunked_state(&store, client_id, backend_ctx)?;
137                let (data, len) = store::filestore_read_chunk(
138                    &store,
139                    client_id,
140                    &request.path,
141                    request.location,
142                    OpenSeekFrom::Start(0),
143                )?;
144                backend_ctx.chunked_io_state = Some(ChunkedIoState::Read(ChunkedReadState {
145                    path: request.path.clone(),
146                    location: request.location,
147                    offset: data.len(),
148                }));
149                Ok(reply::StartChunkedRead { data, len }.into())
150            }
151            ChunkedRequest::WriteChunk(request) => {
152                let is_last = !request.data.is_full();
153                if is_last {
154                    write_last_chunk(&store, client_id, backend_ctx, &request.data)?;
155                } else {
156                    write_chunk(&store, client_id, backend_ctx, &request.data)?;
157                }
158                Ok(reply::WriteChunk {}.into())
159            }
160            ChunkedRequest::AbortChunkedWrite(_request) => {
161                let Some(ChunkedIoState::Write(ref write_state)) = backend_ctx.chunked_io_state
162                else {
163                    return Ok(reply::AbortChunkedWrite { aborted: false }.into());
164                };
165                let aborted = store::abort_chunked_write(
166                    &store,
167                    client_id,
168                    &write_state.path,
169                    write_state.location,
170                );
171                Ok(reply::AbortChunkedWrite { aborted }.into())
172            }
173            ChunkedRequest::StartChunkedWrite(request) => {
174                backend_ctx.chunked_io_state = Some(ChunkedIoState::Write(ChunkedWriteState {
175                    path: request.path.clone(),
176                    location: request.location,
177                }));
178                store::start_chunked_write(
179                    &store,
180                    client_id,
181                    &request.path,
182                    request.location,
183                    &[],
184                )?;
185                Ok(reply::StartChunkedWrite {}.into())
186            }
187            ChunkedRequest::PartialReadFile(request) => {
188                let (data, file_length) = store::partial_read_file(
189                    &store,
190                    client_id,
191                    &request.path,
192                    request.location,
193                    request.offset,
194                    request.length,
195                )?;
196                Ok(reply::PartialReadFile { data, file_length }.into())
197            }
198            ChunkedRequest::AppendFile(request) => {
199                let file_length = store::append_file(
200                    &store,
201                    client_id,
202                    &request.path,
203                    request.location,
204                    &request.data,
205                )?;
206                Ok(reply::AppendFile { file_length }.into())
207            }
208            ChunkedRequest::StartEncryptedChunkedWrite(request) => {
209                clear_chunked_state(&store, client_id, backend_ctx)?;
210                let key = keystore.load_key(
211                    Secrecy::Secret,
212                    Some(Kind::Symmetric(CHACHA8_KEY_LEN)),
213                    &request.key,
214                )?;
215                let nonce = request.nonce.map(|n| *n).unwrap_or_else(|| {
216                    let mut nonce = [0; CHACHA8_STREAM_NONCE_LEN];
217                    rng.fill_bytes(&mut nonce);
218                    nonce
219                });
220                let nonce: &StreamNonce<ChaCha8Poly1305, StreamLE31<ChaCha8Poly1305>> =
221                    (&nonce).into();
222                let aead = ChaCha8Poly1305::new((&*key.material).into());
223                let encryptor = EncryptorLE31::<ChaCha8Poly1305>::from_aead(aead, nonce);
224                store::start_chunked_write(
225                    &store,
226                    client_id,
227                    &request.path,
228                    request.location,
229                    nonce,
230                )?;
231                backend_ctx.chunked_io_state =
232                    Some(ChunkedIoState::EncryptedWrite(EncryptedChunkedWriteState {
233                        path: request.path.clone(),
234                        location: request.location,
235                        encryptor,
236                    }));
237                Ok(reply::StartEncryptedChunkedWrite {}.into())
238            }
239            ChunkedRequest::StartEncryptedChunkedRead(request) => {
240                clear_chunked_state(&store, client_id, backend_ctx)?;
241                let key = keystore.load_key(
242                    Secrecy::Secret,
243                    Some(Kind::Symmetric(CHACHA8_KEY_LEN)),
244                    &request.key,
245                )?;
246                let nonce: Bytes<CHACHA8_STREAM_NONCE_LEN> =
247                    filestore.read(&request.path, request.location)?;
248                let nonce: &StreamNonce<ChaCha8Poly1305, StreamLE31<ChaCha8Poly1305>> =
249                    (&*nonce).into();
250                let aead = ChaCha8Poly1305::new((&*key.material).into());
251                let decryptor = DecryptorLE31::<ChaCha8Poly1305>::from_aead(aead, nonce);
252                backend_ctx.chunked_io_state =
253                    Some(ChunkedIoState::EncryptedRead(EncryptedChunkedReadState {
254                        path: request.path.clone(),
255                        location: request.location,
256                        decryptor,
257                        offset: CHACHA8_STREAM_NONCE_LEN,
258                    }));
259                Ok(reply::StartEncryptedChunkedRead {}.into())
260            }
261        }
262    }
263}
264
265fn clear_chunked_state(
266    store: &impl Store,
267    client_id: &Path,
268    ctx: &mut StagingContext,
269) -> Result<(), Error> {
270    match ctx.chunked_io_state.take() {
271        Some(ChunkedIoState::Read(_)) | None => {}
272        Some(ChunkedIoState::Write(write_state)) => {
273            info!("Automatically cancelling write");
274            store::abort_chunked_write(store, client_id, &write_state.path, write_state.location);
275        }
276        Some(ChunkedIoState::EncryptedRead(_)) => {}
277        Some(ChunkedIoState::EncryptedWrite(write_state)) => {
278            info!("Automatically cancelling encrypted write");
279            store::abort_chunked_write(store, client_id, &write_state.path, write_state.location);
280        }
281    }
282    Ok(())
283}
284
285fn write_chunk(
286    store: &impl Store,
287    client_id: &Path,
288    ctx: &mut StagingContext,
289    data: &Message,
290) -> Result<(), Error> {
291    match ctx.chunked_io_state {
292        Some(ChunkedIoState::Write(ref write_state)) => {
293            store::filestore_write_chunk(
294                store,
295                client_id,
296                &write_state.path,
297                write_state.location,
298                data,
299            )?;
300        }
301        Some(ChunkedIoState::EncryptedWrite(ref mut write_state)) => {
302            let mut data =
303                Bytes::<{ MAX_MESSAGE_LENGTH + POLY1305_TAG_LEN }>::try_from(&**data).unwrap();
304            write_state
305                .encryptor
306                .encrypt_next_in_place(
307                    write_state.path.as_ref().as_bytes(),
308                    &mut HeaplessBuffer::from(&mut data),
309                )
310                .map_err(|_err| {
311                    error!("Failed to encrypt {:?}", _err);
312                    Error::AeadError
313                })?;
314            store::filestore_write_chunk(
315                store,
316                client_id,
317                &write_state.path,
318                write_state.location,
319                &data,
320            )?;
321        }
322        _ => return Err(Error::MechanismNotAvailable),
323    }
324    Ok(())
325}
326
327fn write_last_chunk(
328    store: &impl Store,
329    client_id: &Path,
330    ctx: &mut StagingContext,
331    data: &Message,
332) -> Result<(), Error> {
333    match ctx.chunked_io_state.take() {
334        Some(ChunkedIoState::Write(write_state)) => {
335            store::filestore_write_chunk(
336                store,
337                client_id,
338                &write_state.path,
339                write_state.location,
340                data,
341            )?;
342            store::flush_chunks(store, client_id, &write_state.path, write_state.location)?;
343        }
344        Some(ChunkedIoState::EncryptedWrite(write_state)) => {
345            let mut data =
346                Bytes::<{ MAX_MESSAGE_LENGTH + POLY1305_TAG_LEN }>::try_from(&**data).unwrap();
347            write_state
348                .encryptor
349                .encrypt_last_in_place(
350                    &[write_state.location as u8],
351                    &mut HeaplessBuffer::from(&mut data),
352                )
353                .map_err(|_err| {
354                    error!("Failed to encrypt {:?}", _err);
355                    Error::AeadError
356                })?;
357            store::filestore_write_chunk(
358                store,
359                client_id,
360                &write_state.path,
361                write_state.location,
362                &data,
363            )?;
364            store::flush_chunks(store, client_id, &write_state.path, write_state.location)?;
365        }
366        _ => return Err(Error::MechanismNotAvailable),
367    }
368
369    Ok(())
370}
371
372fn read_encrypted_chunk(
373    store: &impl Store,
374    client_id: &Path,
375    ctx: &mut StagingContext,
376) -> Result<ChunkedReply, Error> {
377    let Some(ChunkedIoState::EncryptedRead(ref mut read_state)) = ctx.chunked_io_state else {
378        unreachable!(
379            "Read encrypted chunk can only be called in the context encrypted chunk reads"
380        );
381    };
382    let (mut data, len): (Bytes<{ MAX_MESSAGE_LENGTH + POLY1305_TAG_LEN }>, usize) =
383        store::filestore_read_chunk(
384            store,
385            client_id,
386            &read_state.path,
387            read_state.location,
388            OpenSeekFrom::Start(read_state.offset as _),
389        )?;
390    read_state.offset += data.len();
391
392    let is_last = !data.is_full();
393    if is_last {
394        let Some(ChunkedIoState::EncryptedRead(read_state)) = ctx.chunked_io_state.take() else {
395            unreachable!();
396        };
397
398        read_state
399            .decryptor
400            .decrypt_last_in_place(
401                &[read_state.location as u8],
402                &mut HeaplessBuffer::from(&mut data),
403            )
404            .map_err(|_err| {
405                error!("Failed to decrypt {:?}", _err);
406                Error::AeadError
407            })?;
408        let data = Bytes::try_from(&*data).expect("decryptor removes the tag");
409        Ok(reply::ReadChunk {
410            data,
411            len: chunked_decrypted_len(len)?,
412        }
413        .into())
414    } else {
415        read_state
416            .decryptor
417            .decrypt_next_in_place(
418                read_state.path.as_ref().as_bytes(),
419                &mut HeaplessBuffer::from(&mut data),
420            )
421            .map_err(|_err| {
422                error!("Failed to decrypt {:?}", _err);
423                Error::AeadError
424            })?;
425        let data = Bytes::try_from(&*data).expect("decryptor removes the tag");
426        Ok(reply::ReadChunk {
427            data,
428            len: chunked_decrypted_len(len)?,
429        }
430        .into())
431    }
432}
433
434/// Calculate the decrypted length of a chunked encrypted file
435fn chunked_decrypted_len(len: usize) -> Result<usize, Error> {
436    let len = len.checked_sub(CHACHA8_STREAM_NONCE_LEN).ok_or_else(|| {
437        error!("File too small");
438        Error::FilesystemReadFailure
439    })?;
440    const CHUNK_LEN: usize = POLY1305_TAG_LEN + MAX_MESSAGE_LENGTH;
441    let chunk_count = len / CHUNK_LEN;
442    let last_chunk_len = (len % CHUNK_LEN)
443        .checked_sub(POLY1305_TAG_LEN)
444        .ok_or_else(|| {
445            error!("Incorrect last chunk length");
446            Error::FilesystemReadFailure
447        })?;
448
449    Ok(chunk_count * MAX_MESSAGE_LENGTH + last_chunk_len)
450}
trussed_staging/chunked/mod.rs

trussed_staging/chunked/
mod.rs
