Skip to main content

rw_builder/
stream_cipher.rs

1use std::{
2    io::{Read, Write},
3    marker::PhantomData,
4};
5
6use anyhow::Result;
7#[cfg(feature = "chacha20")]
8use chacha20::ChaCha20;
9use cipher::{KeyIvInit, StreamCipher};
10#[cfg(feature = "salsa20")]
11use salsa20::Salsa20;
12
13use crate::RwBuilder;
14
15/// Type returned by the `chacha20` and `salsa20` functions on the `RwBuilder`
16/// trait. It is itself an `RwBuilder` so can be chained further.
17#[derive(Debug)]
18pub struct Builder<B, C, K, N>
19where
20    B: RwBuilder,
21    C: StreamCipher,
22{
23    /// The inner builder it wraps
24    wrapped: B,
25    /// The key used for encryption and decryption
26    key: K,
27    /// The nonce used for encryption and decryption
28    nonce: N,
29    /// We need `Builder` to be generic over the `StreamCipher`
30    _marker: PhantomData<C>,
31}
32
33impl<B, C, K, N> Builder<B, C, K, N>
34where
35    B: RwBuilder,
36    C: StreamCipher,
37{
38    /// Create a new cipher builder from a key and a nonce
39    pub const fn new(builder: B, key: K, nonce: N) -> Self {
40        Self { wrapped: builder, key, nonce, _marker: PhantomData }
41    }
42}
43
44/// The key type for the chacha20 cipher
45#[cfg(feature = "chacha20")]
46pub type ChaCha20Key = chacha20::Key;
47
48/// The nonce type for the chacha20 cipher (12 bytes for IETF variant)
49#[cfg(feature = "chacha20")]
50pub type ChaCha20Nonce = [u8; 12];
51
52/// The type returned by the `chacha20` function in the `RwBuilder` trait
53#[cfg(feature = "chacha20")]
54pub type ChaCha20Builder<B> = Builder<B, ChaCha20, ChaCha20Key, ChaCha20Nonce>;
55
56/// The key type for the salsa20 cipher
57#[cfg(feature = "salsa20")]
58pub type Salsa20Key = salsa20::Key;
59
60/// The nonce type for the salsa20 cipher
61#[cfg(feature = "salsa20")]
62pub type Salsa20Nonce = salsa20::Nonce;
63
64/// The type returned by the `salsa20` function in the `RwBuilder` trait
65#[cfg(feature = "salsa20")]
66pub type Salsa20Builder<B> = Builder<B, Salsa20, Salsa20Key, Salsa20Nonce>;
67
68/// Recipe for how to create a cipher
69trait CipherFactory<C> {
70    /// Create the cipher from the key and the nonce stored in self
71    fn create_cipher(&self) -> C;
72}
73
74#[cfg(feature = "chacha20")]
75impl<B> CipherFactory<ChaCha20> for Builder<B, ChaCha20, ChaCha20Key, ChaCha20Nonce>
76where
77    B: RwBuilder,
78{
79    fn create_cipher(&self) -> ChaCha20 {
80        ChaCha20::new(&self.key, (&self.nonce).into())
81    }
82}
83
84#[cfg(feature = "salsa20")]
85impl<B> CipherFactory<Salsa20> for Builder<B, Salsa20, Salsa20Key, Salsa20Nonce>
86where
87    B: RwBuilder,
88{
89    fn create_cipher(&self) -> Salsa20 {
90        Salsa20::new(&self.key, &self.nonce)
91    }
92}
93
94impl<B, C, K, N> RwBuilder for Builder<B, C, K, N>
95where
96    B: RwBuilder,
97    C: StreamCipher,
98    Self: CipherFactory<C>,
99{
100    type Reader = Reader<B::Reader, C>;
101    type Writer = Writer<B::Writer, C>;
102
103    fn reader(&self) -> Result<Self::Reader> {
104        let reader = self.wrapped.reader()?;
105        let cipher = self.create_cipher();
106        Ok(Reader { cipher, reader })
107    }
108
109    fn writer(&self) -> Result<Self::Writer> {
110        let writer = self.wrapped.writer()?;
111        let cipher = self.create_cipher();
112        Ok(Writer { cipher, writer })
113    }
114}
115
116/// Generic Reader type for multiple ciphers
117#[derive(Debug)]
118pub struct Reader<R, C>
119where
120    R: Read,
121    C: StreamCipher,
122{
123    /// The cipher to use for reading
124    cipher: C,
125    /// The wrapped reader
126    reader: R,
127}
128
129impl<R, C> Read for Reader<R, C>
130where
131    R: Read,
132    C: StreamCipher,
133{
134    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
135        let bytes_read = self.reader.read(buf)?;
136        self.cipher.try_apply_keystream(buf).map_err(std::io::Error::other)?;
137        Ok(bytes_read)
138    }
139}
140
141/// Generic Writer type for multiple ciphers
142#[derive(Debug)]
143pub struct Writer<W, C>
144where
145    W: Write,
146    C: StreamCipher,
147{
148    /// The cipher to use for writing
149    cipher: C,
150    /// The wrapped writer
151    writer: W,
152}
153
154impl<W, C> Write for Writer<W, C>
155where
156    W: Write,
157    C: StreamCipher,
158{
159    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
160        let mut buffer = buf.to_owned();
161        self.cipher.try_apply_keystream(buffer.as_mut_slice()).map_err(std::io::Error::other)?;
162        self.writer.write(buffer.as_slice())
163    }
164
165    fn flush(&mut self) -> std::io::Result<()> {
166        self.writer.flush()
167    }
168}