ssec_core/
encrypt.rs

1use futures_core::Stream;
2use bytes::{Bytes, BytesMut, BufMut};
3use rand_core::TryCryptoRng;
4use ctr::cipher::{KeyIvInit, StreamCipher};
5use hmac::{Mac, KeyInit};
6use core::pin::Pin;
7use core::task::{Context, Poll, ready};
8use core::num::NonZeroUsize;
9use crate::util::{HmacSha3_256, new_arr, kdf, compute_verification_hash};
10use crate::{DEFAULT_BYTES_PER_POLL, Aes256Ctr};
11
12/// builder for arguments to [Encrypt::new] with default values, can be constructed with [EncryptArgs::default]
13#[derive(Debug, Clone, Copy)]
14pub struct EncryptArgs {
15	bytes_per_poll: NonZeroUsize
16}
17
18impl Default for EncryptArgs {
19	/// default settings are not part of semver contract
20	fn default() -> Self {
21		Self {
22			bytes_per_poll: DEFAULT_BYTES_PER_POLL
23		}
24	}
25}
26
27impl EncryptArgs {
28	/// sets the maximum number of bytes to encrypt before yielding to the executor
29	pub fn set_bytes_per_poll(&mut self, bytes_per_poll: NonZeroUsize) {
30		self.bytes_per_poll = bytes_per_poll;
31	}
32}
33
34enum EncryptState {
35	PreHeader,
36	PostHeader,
37	Finalizing,
38	Finished
39}
40
41pin_project_lite::pin_project! {
42	pub struct Encrypt<R> {
43		#[pin]
44		read: R,
45		aes: Aes256Ctr,
46		password_salt: Box<[u8; 32]>,
47		password_verification_hash: Box<[u8; 64]>,
48		integrity_code: Option<HmacSha3_256>,
49		state: EncryptState,
50		block_buffer: BytesMut,
51		iv: [u8; 16],
52		bytes_per_poll: NonZeroUsize
53	}
54}
55
56impl<R> Encrypt<R> {
57	/// This method is *very* blocking.
58	/// If you're using Tokio I advise that you wrap this call in a `spawn_blocking`.
59	///
60	/// SECURITY: It is advisable to zero out the memory containing the password after this method returns.
61	pub fn new<RNG: TryCryptoRng>(
62		additional_args: EncryptArgs,
63		rng: &mut RNG,
64		password: &[u8],
65		read: R
66	) -> Result<Self, RNG::Error> {
67		let mut password_salt = new_arr::<32>();
68		rng.try_fill_bytes(password_salt.as_mut())?;
69
70		let aes_key = kdf(password, password_salt.as_ref());
71		let password_verification_hash = compute_verification_hash(&aes_key);
72
73		let mut iv = [0; 16];
74		rng.try_fill_bytes(&mut iv)?;
75
76		let aes = Aes256Ctr::new(aes_key.as_ref().get_ref().into(), (&iv).into());
77
78		Ok(Self {
79			read,
80			aes,
81			password_salt,
82			password_verification_hash,
83			integrity_code: Some(HmacSha3_256::new_from_slice(aes_key.as_ref().get_ref()).unwrap()),
84			state: EncryptState::PreHeader,
85			block_buffer: BytesMut::new(),
86			iv,
87			bytes_per_poll: additional_args.bytes_per_poll
88		})
89	}
90}
91
92impl<E, R: Stream<Item = Result<Bytes, E>>> Stream for Encrypt<R> {
93	type Item = Result<Bytes, E>;
94
95	fn poll_next(
96		self: Pin<&mut Self>,
97		cx: &mut Context<'_>
98	) -> Poll<Option<Self::Item>> {
99		let mut this = self.project();
100
101		loop {
102			match this.state {
103				EncryptState::PreHeader => {
104					let mut buf = Vec::with_capacity(
105						4 + // magic
106						1 + // version number
107						1 + // compression algo
108						32 + // password salt
109						64 + // password verification hash
110						16 // IV
111					);
112
113					buf.extend_from_slice(b"SSEC");
114					buf.push(0x01);
115					buf.push(0x6e);
116					buf.extend_from_slice(this.password_salt.as_ref());
117					buf.extend_from_slice(this.password_verification_hash.as_ref());
118					buf.extend_from_slice(this.iv.as_ref());
119
120					// as per spec: first we add the version byte, compression algo, then iv before the data
121					let integrity_code = this.integrity_code.as_mut().unwrap();
122					integrity_code.update(&[0x01, 0x6e]);
123					integrity_code.update(this.iv.as_ref());
124
125					match this.read.poll_next(cx) {
126						Poll::Pending => *this.state = EncryptState::PostHeader,
127						Poll::Ready(None) => *this.state = EncryptState::Finalizing,
128						Poll::Ready(Some(Err(e))) => {
129							*this.state = EncryptState::Finished;
130							return Poll::Ready(Some(Err(e)));
131						},
132						Poll::Ready(Some(Ok(bytes))) => {
133							*this.state = EncryptState::PostHeader;
134							this.block_buffer.put(bytes);
135						}
136					}
137
138					return Poll::Ready(Some(Ok(Bytes::from_owner(buf))));
139				},
140				EncryptState::PostHeader => {
141					if this.block_buffer.len() >= this.bytes_per_poll.get() {
142						let mut data = this.block_buffer.split_to(this.bytes_per_poll.get());
143						this.aes.apply_keystream(&mut data);
144						this.integrity_code.as_mut().unwrap().update(&data);
145
146						return Poll::Ready(Some(Ok(data.freeze())));
147					} else {
148						match ready!(this.read.as_mut().poll_next(cx)) {
149							Some(Ok(bytes)) => {
150								this.block_buffer.put(bytes);
151								continue;
152							},
153							Some(Err(e)) => {
154								*this.state = EncryptState::Finished;
155								return Poll::Ready(Some(Err(e)));
156							},
157							None => {
158								*this.state = EncryptState::Finalizing;
159								continue;
160							}
161						}
162					}
163				},
164				EncryptState::Finalizing => {
165					debug_assert!(this.block_buffer.len() < this.bytes_per_poll.get());
166
167					let mut final_data = this.block_buffer.split();
168
169					let mut hmac = this.integrity_code.take()
170						.expect("integrity_code only taken here");
171
172					this.aes.apply_keystream(&mut final_data);
173
174					hmac.update(&final_data);
175					final_data.put(Bytes::from_owner(hmac.finalize().into_bytes()));
176
177					*this.state = EncryptState::Finished;
178
179					return Poll::Ready(Some(Ok(final_data.freeze())));
180				},
181				EncryptState::Finished => return Poll::Ready(None)
182			}
183		}
184	}
185}
ssec_core/encrypt.rs

ssec_core/
encrypt.rs
