ssb_box_stream/
decrypt.rs

1use futures::prelude::*;
2use std::pin::Pin;
3use std::task::{Context, Poll};
4
5use crate::utils::ReadBuffer;
6
7/// A [Stream] of `Vec<u8>` that decrypts and authenticates data from the underlying `Reader`.
8#[pin_project::pin_project]
9pub struct Decrypt<Reader: AsyncRead> {
10    #[pin]
11    reader: Reader,
12    params: crate::cipher::Params,
13    state: DecryptState,
14}
15
16impl<Reader: AsyncRead> Decrypt<Reader> {
17    pub fn new(reader: Reader, params: crate::cipher::Params) -> Self {
18        Decrypt {
19            reader,
20            params,
21            state: DecryptState::init(),
22        }
23    }
24}
25
26/// Error when decrypting and authenticating data.
27#[derive(thiserror::Error, Debug)]
28pub enum DecryptError {
29    /// The underlying source errored.
30    #[error(transparent)]
31    Io(#[from] std::io::Error),
32
33    /// Failed to decrypt and authenticate packet body
34    #[error("Failed to decrypt and authenticate packet body")]
35    UnboxBody,
36
37    /// Failed to decrypt and authenticate packet header
38    #[error("Failed to decrypt and authenticate packet header")]
39    UnboxHeader,
40
41    /// Received packet that exceeds maximum packet size
42    #[error("Received packet that exceeds maximum packet size")]
43    ExceededMaxPacketSize,
44}
45
46enum DecryptState {
47    Closed,
48    ReadingHeader {
49        buffer: ReadBuffer,
50    },
51    ReadingBody {
52        auth_tag: sodiumoxide::crypto::secretbox::Tag,
53        buffer: ReadBuffer,
54    },
55}
56
57impl DecryptState {
58    fn init() -> Self {
59        DecryptState::ReadingHeader {
60            buffer: ReadBuffer::new(crate::cipher::BOXED_HEADER_SIZE),
61        }
62    }
63}
64
65impl<Reader: AsyncRead> Stream for Decrypt<Reader> {
66    type Item = Result<Vec<u8>, DecryptError>;
67
68    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
69        let result = futures::ready!(self.as_mut().poll_next_inner(cx));
70        match result {
71            Some(Err(_)) | None => *self.project().state = DecryptState::Closed,
72            _ => (),
73        }
74        Poll::Ready(result)
75    }
76}
77
78impl<Reader: AsyncRead> Decrypt<Reader> {
79    fn poll_next_inner(
80        mut self: Pin<&mut Self>,
81        cx: &mut Context,
82    ) -> Poll<Option<Result<Vec<u8>, DecryptError>>> {
83        loop {
84            let mut this = self.as_mut().project();
85            match &mut this.state {
86                DecryptState::Closed => return Poll::Ready(None),
87                DecryptState::ReadingHeader { buffer } => {
88                    let boxed_header = futures::ready!(buffer.poll_read(cx, this.reader))?;
89                    let mut boxed_header_array = [0u8; crate::cipher::BOXED_HEADER_SIZE];
90                    boxed_header_array.copy_from_slice(&boxed_header);
91                    match this
92                        .params
93                        .decrypt_header(&boxed_header_array)
94                        .map_err(|()| DecryptError::UnboxHeader)?
95                    {
96                        Some((len, auth_tag)) => {
97                            if len >= crate::cipher::MAX_PACKET_SIZE_BYTES {
98                                *this.state = DecryptState::Closed;
99                                return Poll::Ready(Some(Err(DecryptError::ExceededMaxPacketSize)));
100                            }
101                            *this.state = DecryptState::ReadingBody {
102                                auth_tag,
103                                buffer: ReadBuffer::new(len as usize),
104                            }
105                        }
106                        None => {
107                            *this.state = DecryptState::Closed;
108                        }
109                    }
110                }
111                DecryptState::ReadingBody { auth_tag, buffer } => {
112                    let boxed_body = futures::ready!(buffer.poll_read(cx, this.reader))?;
113                    let body = this
114                        .params
115                        .decrypt_body(auth_tag, &boxed_body)
116                        .map_err(|()| DecryptError::UnboxBody)?;
117                    *this.state = DecryptState::init();
118                    return Poll::Ready(Some(Ok(body)));
119                }
120            };
121        }
122    }
123}