ssec_core/
chaff.rs

1use rand_core::TryRng;
2use bytes::Bytes;
3use futures_core::Stream;
4use core::pin::Pin;
5use core::task::{Context, Poll};
6use crate::HEADER_LENGTH;
7
8enum ChaffState {
9	PreHeader,
10	Data,
11	Finished
12}
13
14pin_project_lite::pin_project! {
15	pub struct ChaffStream<RNG> {
16		rng: RNG,
17		state: ChaffState,
18		remaining_bytes: usize,
19		chunk_size: usize
20	}
21}
22
23const MIN_CHUNK_SIZE: usize = HEADER_LENGTH;
24
25#[derive(Debug)]
26pub enum NewChaffStreamError {
27	ChunkSizeTooSmall
28}
29
30impl std::fmt::Display for NewChaffStreamError {
31	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32		match self {
33			Self::ChunkSizeTooSmall => write!(f, "chunk size too small, must be at least {MIN_CHUNK_SIZE} bytes")
34		}
35	}
36}
37
38impl std::error::Error for NewChaffStreamError {}
39
40impl<RNG:TryRng> ChaffStream<RNG> {
41	/// The `output_length` parameter controls the size of the hypothetical chaff input file.
42	/// In other words, the length of the stream is the length of the headers plus `output_length`.
43	pub fn new(rng: RNG, output_length: usize, chunk_size: usize) -> Result<Self, NewChaffStreamError> {
44		if chunk_size < MIN_CHUNK_SIZE {
45			return Err(NewChaffStreamError::ChunkSizeTooSmall);
46		}
47
48		Ok(Self {
49			rng,
50			state: ChaffState::PreHeader,
51			remaining_bytes: output_length,
52			chunk_size
53		})
54	}
55}
56
57impl<RNG: TryRng> Stream for ChaffStream<RNG> {
58	type Item = Result<Bytes, RNG::Error>;
59
60	fn poll_next(
61		self: Pin<&mut Self>,
62		_cx: &mut Context<'_>
63	) -> Poll<Option<Self::Item>> {
64		let this = self.project();
65
66		match this.state {
67			ChaffState::PreHeader => {
68				let output_len: usize = (HEADER_LENGTH + *this.remaining_bytes).min(*this.chunk_size);
69				assert!(output_len >= HEADER_LENGTH);
70
71				let mut output = Vec::with_capacity(output_len);
72
73				output.extend_from_slice(b"SSEC");
74				output.push(0x01);
75				output.push(0x6e);
76				output.extend_from_slice(&vec![0u8; output_len - 6]);
77
78				if let Err(err) = this.rng.try_fill_bytes(&mut output[6..]) {
79					*this.state = ChaffState::Finished;
80					return Poll::Ready(Some(Err(err)));
81				};
82
83				let bytes_left = *this.remaining_bytes - (output.len() - HEADER_LENGTH);
84				if bytes_left == 0 {
85					*this.state = ChaffState::Finished
86				} else {
87					*this.state = ChaffState::Data;
88					*this.remaining_bytes = bytes_left;
89				}
90
91				Poll::Ready(Some(Ok(Bytes::from_owner(output))))
92			},
93			ChaffState::Data => {
94				let chaff_len: usize = (*this.chunk_size).min(*this.remaining_bytes);
95				let mut chaff_data = vec![0u8; chaff_len];
96
97				if let Err(err) = this.rng.try_fill_bytes(&mut chaff_data) {
98					*this.state = ChaffState::Finished;
99					return Poll::Ready(Some(Err(err)));
100				};
101
102				let bytes_left = *this.remaining_bytes - chaff_len;
103				if bytes_left == 0 {
104					*this.state = ChaffState::Finished
105				} else {
106					*this.remaining_bytes = bytes_left;
107				}
108
109				Poll::Ready(Some(Ok(Bytes::from_owner(chaff_data))))
110			},
111			ChaffState::Finished => Poll::Ready(None)
112		}
113	}
114}
ssec_core/chaff.rs

ssec_core/
chaff.rs
