1use std::{
2 io::{Read, Write},
3 marker::PhantomData,
4};
5
6use crate::Result;
7#[cfg(feature = "aes_ctr")]
8use aes::{Aes128, Aes256};
9#[cfg(feature = "chacha20")]
10use chacha20::ChaCha20;
11use cipher::{KeyIvInit, StreamCipher};
12#[cfg(feature = "aes_ctr")]
13use ctr::Ctr128BE;
14#[cfg(feature = "salsa20")]
15use salsa20::Salsa20;
16
17use crate::RwBuilder;
18
19#[derive(Debug)]
22#[must_use]
23pub struct Builder<B, C, K, N>
24where
25 B: RwBuilder,
26 C: StreamCipher,
27{
28 wrapped: B,
30 key: K,
32 nonce: N,
34 _marker: PhantomData<C>,
36}
37
38impl<B, C, K, N> Builder<B, C, K, N>
39where
40 B: RwBuilder,
41 C: StreamCipher,
42{
43 pub const fn new(builder: B, key: K, nonce: N) -> Self {
45 Self { wrapped: builder, key, nonce, _marker: PhantomData }
46 }
47}
48
49#[cfg(feature = "chacha20")]
51pub type ChaCha20Key = chacha20::Key;
52
53#[cfg(feature = "chacha20")]
55pub type ChaCha20Nonce = [u8; 12];
56
57#[cfg(feature = "chacha20")]
59pub type ChaCha20Builder<B> = Builder<B, ChaCha20, ChaCha20Key, ChaCha20Nonce>;
60
61#[cfg(feature = "salsa20")]
63pub type Salsa20Key = salsa20::Key;
64
65#[cfg(feature = "salsa20")]
67pub type Salsa20Nonce = salsa20::Nonce;
68
69#[cfg(feature = "salsa20")]
71pub type Salsa20Builder<B> = Builder<B, Salsa20, Salsa20Key, Salsa20Nonce>;
72
73#[cfg(feature = "aes_ctr")]
75pub type Aes128Ctr = Ctr128BE<Aes128>;
76#[cfg(feature = "aes_ctr")]
78pub type Aes256Ctr = Ctr128BE<Aes256>;
79
80#[cfg(feature = "aes_ctr")]
82pub type Aes128Key = [u8; 16];
83#[cfg(feature = "aes_ctr")]
85pub type Aes256Key = [u8; 32];
86#[cfg(feature = "aes_ctr")]
88pub type AesNonce = [u8; 16];
89
90#[cfg(feature = "aes_ctr")]
92pub type Aes128CtrBuilder<B> = Builder<B, Aes128Ctr, Aes128Key, AesNonce>;
93#[cfg(feature = "aes_ctr")]
95pub type Aes256CtrBuilder<B> = Builder<B, Aes256Ctr, Aes256Key, AesNonce>;
96
97trait CipherFactory<C> {
99 fn create_cipher(&self) -> C;
101}
102
103#[cfg(feature = "chacha20")]
104impl<B> CipherFactory<ChaCha20> for Builder<B, ChaCha20, ChaCha20Key, ChaCha20Nonce>
105where
106 B: RwBuilder,
107{
108 fn create_cipher(&self) -> ChaCha20 {
109 ChaCha20::new(&self.key, (&self.nonce).into())
110 }
111}
112
113#[cfg(feature = "salsa20")]
114impl<B> CipherFactory<Salsa20> for Builder<B, Salsa20, Salsa20Key, Salsa20Nonce>
115where
116 B: RwBuilder,
117{
118 fn create_cipher(&self) -> Salsa20 {
119 Salsa20::new(&self.key, &self.nonce)
120 }
121}
122
123#[cfg(feature = "aes_ctr")]
124impl<B> CipherFactory<Aes128Ctr> for Builder<B, Aes128Ctr, Aes128Key, AesNonce>
125where
126 B: RwBuilder,
127{
128 fn create_cipher(&self) -> Aes128Ctr {
129 Aes128Ctr::new((&self.key).into(), (&self.nonce).into())
130 }
131}
132
133#[cfg(feature = "aes_ctr")]
134impl<B> CipherFactory<Aes256Ctr> for Builder<B, Aes256Ctr, Aes256Key, AesNonce>
135where
136 B: RwBuilder,
137{
138 fn create_cipher(&self) -> Aes256Ctr {
139 Aes256Ctr::new((&self.key).into(), (&self.nonce).into())
140 }
141}
142
143impl<B, C, K, N> RwBuilder for Builder<B, C, K, N>
144where
145 B: RwBuilder,
146 C: StreamCipher,
147 Self: CipherFactory<C>,
148{
149 type Reader = Reader<B::Reader, C>;
150 type Writer = Writer<B::Writer, C>;
151
152 fn reader(&self) -> Result<Self::Reader> {
153 let reader = self.wrapped.reader()?;
154 let cipher = self.create_cipher();
155 Ok(Reader { cipher, reader })
156 }
157
158 fn writer(&self) -> Result<Self::Writer> {
159 let writer = self.wrapped.writer()?;
160 let cipher = self.create_cipher();
161 Ok(Writer { cipher, writer })
162 }
163}
164
165#[derive(Debug)]
167pub struct Reader<R, C>
168where
169 R: Read,
170 C: StreamCipher,
171{
172 cipher: C,
174 reader: R,
176}
177
178impl<R, C> Read for Reader<R, C>
179where
180 R: Read,
181 C: StreamCipher,
182{
183 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
184 let bytes_read = self.reader.read(buf)?;
185 self.cipher.try_apply_keystream(buf).map_err(std::io::Error::other)?;
186 Ok(bytes_read)
187 }
188}
189
190#[derive(Debug)]
192pub struct Writer<W, C>
193where
194 W: Write,
195 C: StreamCipher,
196{
197 cipher: C,
199 writer: W,
201}
202
203impl<W, C> Write for Writer<W, C>
204where
205 W: Write,
206 C: StreamCipher,
207{
208 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
209 let mut buffer = buf.to_owned();
210 self.cipher.try_apply_keystream(buffer.as_mut_slice()).map_err(std::io::Error::other)?;
211 self.writer.write(buffer.as_slice())
212 }
213
214 fn flush(&mut self) -> std::io::Result<()> {
215 self.writer.flush()
216 }
217}