1
use std::{
2
    io::{Read, Write},
3
    marker::PhantomData,
4
};
5

            
6
use anyhow::Result;
7
#[cfg(feature = "chacha20")]
8
use chacha20::ChaCha20;
9
use cipher::{KeyIvInit, StreamCipher};
10
#[cfg(feature = "salsa20")]
11
use salsa20::Salsa20;
12

            
13
use 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)]
18
pub struct Builder<B, C, K, N>
19
where
20
    B: RwBuilder,
21
    C: StreamCipher,
22
{
23
    /// The inner builder it wraps
24
    builder: 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

            
33
impl<B, C, K, N> Builder<B, C, K, N>
34
where
35
    B: RwBuilder,
36
    C: StreamCipher,
37
{
38
    /// Create a new cipher builder from a key and a nonce
39
2
    pub const fn new(builder: B, key: K, nonce: N) -> Self {
40
2
        Self { builder, key, nonce, _marker: PhantomData }
41
2
    }
42
}
43

            
44
/// The key type for the chacha20 cipher
45
#[cfg(feature = "chacha20")]
46
pub type ChaCha20Key = chacha20::Key;
47

            
48
/// The nonce type for the chacha20 cipher
49
#[cfg(feature = "chacha20")]
50
pub type ChaCha20Nonce = chacha20::Nonce;
51

            
52
/// The type returned by the `chacha20` function in the `RwBuilder` trait
53
#[cfg(feature = "chacha20")]
54
pub type ChaCha20Builder<B> = Builder<B, ChaCha20, ChaCha20Key, ChaCha20Nonce>;
55

            
56
/// The key type for the salsa20 cipher
57
#[cfg(feature = "salsa20")]
58
pub type Salsa20Key = salsa20::Key;
59

            
60
/// The nonce type for the salsa20 cipher
61
#[cfg(feature = "salsa20")]
62
pub type Salsa20Nonce = salsa20::Nonce;
63

            
64
/// The type returned by the `salsa20` function in the `RwBuilder` trait
65
#[cfg(feature = "salsa20")]
66
pub type Salsa20Builder<B> = Builder<B, Salsa20, Salsa20Key, Salsa20Nonce>;
67

            
68
/// Recipe for how to create a cipher
69
trait 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")]
75
impl<B> CipherFactory<ChaCha20> for Builder<B, ChaCha20, ChaCha20Key, ChaCha20Nonce>
76
where
77
    B: RwBuilder,
78
{
79
2
    fn create_cipher(&self) -> ChaCha20 {
80
2
        ChaCha20::new(&self.key, &self.nonce)
81
2
    }
82
}
83

            
84
#[cfg(feature = "salsa20")]
85
impl<B> CipherFactory<Salsa20> for Builder<B, Salsa20, Salsa20Key, Salsa20Nonce>
86
where
87
    B: RwBuilder,
88
{
89
2
    fn create_cipher(&self) -> Salsa20 {
90
2
        Salsa20::new(&self.key, &self.nonce)
91
2
    }
92
}
93

            
94
impl<B, C, K, N> RwBuilder for Builder<B, C, K, N>
95
where
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
2
    fn reader(&self) -> Result<Self::Reader> {
104
2
        let reader = self.builder.reader()?;
105
2
        let cipher = self.create_cipher();
106
2
        Ok(Reader { cipher, reader })
107
2
    }
108

            
109
2
    fn writer(&self) -> Result<Self::Writer> {
110
2
        let writer = self.builder.writer()?;
111
2
        let cipher = self.create_cipher();
112
2
        Ok(Writer { cipher, writer })
113
2
    }
114
}
115

            
116
/// Generic Reader type for multiple ciphers
117
#[derive(Debug)]
118
pub struct Reader<R, C>
119
where
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

            
129
impl<R, C> Read for Reader<R, C>
130
where
131
    R: Read,
132
    C: StreamCipher,
133
{
134
2
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
135
2
        let bytes_read = self.reader.read(buf)?;
136
2
        self.cipher
137
2
            .try_apply_keystream(buf)
138
2
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
139
2
        Ok(bytes_read)
140
2
    }
141
}
142

            
143
/// Generic Writer type for multiple ciphers
144
#[derive(Debug)]
145
pub struct Writer<W, C>
146
where
147
    W: Write,
148
    C: StreamCipher,
149
{
150
    /// The cipher to use for writing
151
    cipher: C,
152
    /// The wrapped writer
153
    writer: W,
154
}
155

            
156
impl<W, C> Write for Writer<W, C>
157
where
158
    W: Write,
159
    C: StreamCipher,
160
{
161
2
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
162
2
        let mut buffer = buf.to_owned();
163
2
        self.cipher
164
2
            .try_apply_keystream(buffer.as_mut_slice())
165
2
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
166
2
        self.writer.write(buffer.as_slice())
167
2
    }
168

            
169
2
    fn flush(&mut self) -> std::io::Result<()> {
170
2
        self.writer.flush()
171
2
    }
172
}