1use crate::{Cipher, CryptoError, CryptoErrorKind};
4use crate::cipher_mode::{Counter, EncryptStream, Pond, DecryptStream};
5use std::marker::PhantomData;
6use std::cell::Cell;
7
8pub struct CTR<C, T> {
9 buf: Cell<Vec<u8>>,
10 cipher: C,
11 counter: Cell<T>,
12 phd: PhantomData<*const u8>,
13}
14
15impl<C, T> CTR<C, T>
16 where C: Cipher, T: Counter {
17 pub fn new(cipher: C, counter: T) -> Result<Self, CryptoError> {
18 let block_len = cipher.block_size().unwrap_or(1);
19
20 if counter.bits_len() < (block_len << 3) {
21 Err(CryptoError::new(CryptoErrorKind::InnerErr, format!("The length of counter value is too short: {}<{} in bits", counter.bits_len(), block_len << 3)))
22 } else {
23 Ok(
24 Self {
25 buf: Cell::new(Vec::with_capacity(block_len)),
26 cipher,
27 counter: Cell::new(counter),
28 phd: PhantomData,
29 }
30 )
31 }
32 }
33
34 pub fn set_counter(&mut self, counter: T) -> Result<(), CryptoError> {
35 let block_len = self.cipher.block_size().unwrap_or(1);
36 if counter.bits_len() < (block_len << 3) {
37 Err(CryptoError::new(CryptoErrorKind::InnerErr, format!("The length of counter value is too short: {}<{} in bits", counter.bits_len(), block_len << 3)))
38 } else {
39 self.counter.set(counter);
40 Ok(())
41 }
42 }
43
44 #[inline]
45 fn get_buf(&self) -> &mut Vec<u8> {
46 unsafe {
47 &mut (*self.buf.as_ptr())
48 }
49 }
50
51 #[inline]
52 fn get_counter(&self) -> &mut T {
53 unsafe {
54 &mut (*self.counter.as_ptr())
55 }
56 }
57
58 fn encrypt_inner(&self, mut data: &[u8], dst: &mut Vec<u8>) -> Result<usize, CryptoError> {
59 let block_len = self.cipher.block_size().unwrap_or(1);
60 let oj = self.get_buf();
61 while !data.is_empty() {
62 match self.get_counter().next() {
63 Some(c) => {
64 match self.cipher.encrypt(oj, &c.as_slice()[..block_len]) {
65 Ok(_) => {
66 let len = std::cmp::min(block_len,data.len());
67 let block = &data[..len];
68 block.iter().zip(oj.iter()).for_each(|(&a, &b)| {
69 dst.push(a ^ b);
70 });
71 data = &data[len..];
72 },
73 Err(e) => {
74 return Err(e);
75 }
76 }
77 },
78 None => {
79 return Err(CryptoError::new(CryptoErrorKind::InnerErr,
80 format!("counter next is none")));
81 }
82 }
83 }
84
85 Ok(dst.len())
86 }
87
88 pub fn encrypt_stream(self) -> CTREncrypt<C, T> {
89 let len = self.cipher.block_size().unwrap_or(1);
90 self.get_counter().reset();
91 CTREncrypt {
92 ctr: self,
93 data: Vec::with_capacity(len),
94 pond: Vec::with_capacity(len),
95 }
96 }
97
98 pub fn decrypt_stream(self) -> CTRDecrypt<C, T> {
99 CTRDecrypt {
100 ctr: self.encrypt_stream()
101 }
102 }
103}
104
105impl<C, T> Cipher for CTR<C, T>
106 where C: Cipher, T: Counter {
107 type Output = usize;
108
109 fn block_size(&self) -> Option<usize> {
110 self.cipher.block_size()
111 }
112
113 fn encrypt(&self, dst: &mut Vec<u8>, plaintext_block: &[u8]) -> Result<usize, CryptoError> {
114 dst.clear();
115 self.encrypt_inner(plaintext_block, dst)
116 }
117
118 fn decrypt(&self, dst: &mut Vec<u8>, cipher_block: &[u8]) -> Result<usize, CryptoError> {
119 dst.clear();
120
121 self.encrypt_inner(cipher_block, dst)
122 }
123}
124
125impl<C, T> Clone for CTR<C, T>
126 where C: Cipher + Clone, T: Counter + Clone {
127 fn clone(&self) -> Self {
128 Self {
129 buf: Cell::new(self.get_buf().clone()),
130 cipher: self.cipher.clone(),
131 counter: Cell::new(self.get_counter().clone()),
132 phd: PhantomData,
133 }
134 }
135}
136
137pub struct CTREncrypt<C, T> {
138 ctr: CTR<C, T>,
139 data: Vec<u8>,
140 pond: Vec<u8>,
141}
142
143pub struct CTRDecrypt<C, T> {
144 ctr: CTREncrypt<C, T>,
145}
146
147impl<C, T> CTREncrypt<C, T>
148 where C: Cipher, T: Counter {
149 pub fn reset(&mut self) {
150 self.data.clear();
151 self.pond.clear();
152 self.ctr.get_counter().reset();
153 }
154}
155
156impl<C, T> CTRDecrypt<C, T>
157 where C: Cipher, T: Counter {
158 pub fn reset(&mut self) {
159 self.ctr.reset();
160 }
161}
162
163impl<C, T> Cipher for CTREncrypt<C, T>
164 where C: Cipher, T: Counter {
165 type Output = usize;
166
167 fn block_size(&self) -> Option<usize> {
168 self.ctr.block_size()
169 }
170
171 fn encrypt(&self, dst: &mut Vec<u8>, plaintext_block: &[u8]) -> Result<usize, CryptoError> {
172 self.ctr.encrypt(dst, plaintext_block)
173 }
174
175 fn decrypt(&self, dst: &mut Vec<u8>, cipher_block: &[u8]) -> Result<usize, CryptoError> {
176 self.ctr.decrypt(dst, cipher_block)
177 }
178}
179
180impl<C, T> Cipher for CTRDecrypt<C, T>
181 where C: Cipher, T: Counter {
182 type Output = usize;
183
184 fn block_size(&self) -> Option<usize> {
185 self.ctr.block_size()
186 }
187
188 fn encrypt(&self, dst: &mut Vec<u8>, plaintext_block: &[u8]) -> Result<usize, CryptoError> {
189 self.ctr.encrypt(dst, plaintext_block)
190 }
191
192 fn decrypt(&self, dst: &mut Vec<u8>, cipher_block: &[u8]) -> Result<usize, CryptoError> {
193 self.ctr.decrypt(dst, cipher_block)
194 }
195}
196
197impl<C, T> EncryptStream for CTREncrypt<C, T>
198 where C: Cipher, T: Counter {
199 fn write(&mut self, data: &[u8]) -> Result<Pond, CryptoError> {
200 let block_len = self.ctr.block_size().unwrap_or(1);
201 if data.is_empty() {
202 Ok(Pond::new(&mut self.pond, false))
203 } else {
204 self.data.extend(data.iter());
205
206 let remain = self.data.len() % block_len;
207 if let Err(e) = self.ctr.encrypt_inner(&self.data.as_slice()[..(self.data.len() - remain)], &mut self.pond) {
208 Err(e)
209 } else {
210 let tmp = self.ctr.get_buf();
211 tmp.clear();
212 tmp.extend(self.data.iter().skip(self.data.len() - remain));
213 self.data.clear();
214 self.data.append(tmp);
215 Ok(Pond::new(&mut self.pond, false))
216 }
217 }
218 }
219
220 fn finish(&mut self) -> Result<Pond, CryptoError> {
221 match self.ctr.encrypt_inner(self.data.as_slice(), &mut self.pond) {
222 Ok(_) => {
223 self.data.clear();
224 Ok(Pond::new(&mut self.pond, true))
225 },
226 Err(e) => {
227 Err(e)
228 }
229 }
230 }
231}
232
233impl<C, T> DecryptStream for CTRDecrypt<C, T>
234 where C: Cipher, T: Counter {
235 fn write(&mut self, data: &[u8]) -> Result<Pond, CryptoError> {
236 self.ctr.write(data)
237 }
238
239 fn finish(&mut self) -> Result<Pond, CryptoError> {
240 self.ctr.finish()
241 }
242}