Skip to main content

aes_elk/
elk.rs

1use cipher::{
2    Block, BlockCipherEncrypt, BlockSizeUser, InnerIvInit, Iv,
3    IvSizeUser, typenum::{U16, U32, U64}, common::InnerUser,
4};
5use core::fmt;
6
7#[cfg(feature = "zeroize")]
8use zeroize::ZeroizeOnDrop;
9
10/// Error types for ELK operations.
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum Error {
13    /// The IV provided to ELK mode cannot be all zeroes, as the LFSR would never advance.
14    ZeroIv,
15}
16
17impl fmt::Display for Error {
18    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
19        match self {
20            Self::ZeroIv => write!(f, "IV cannot be all zeroes in ELK mode"),
21        }
22    }
23}
24
25impl core::error::Error for Error {}
26
27/// Trait to define the Galois field feedback polynomial mask for different block sizes.
28pub trait ElkPolynomial {
29    fn apply_mask(block: &mut [u8]);
30}
31
32impl ElkPolynomial for U16 {
33    #[inline(always)]
34    fn apply_mask(block: &mut [u8]) {
35        // SELF-COMPUTED. NOT IN NIST SP 800-197!
36        // Polynomial for n = 128:
37        // X^128 + X^113 + X^97 + X^89 + X^71 + X^59 + X^37 + X^23 + 1
38
39        block[1]  ^= 1 << 1; // X^113
40        block[3]  ^= 1 << 1; // X^97
41        block[4]  ^= 1 << 1; // X^89
42        block[7]  ^= 1 << 7; // X^71
43        block[8]  ^= 1 << 3; // X^59
44        block[11] ^= 1 << 5; // X^37
45        block[13] ^= 1 << 7; // X^23
46        block[15] ^= 1 << 0; // 1 (X^0)
47    }
48}
49
50impl ElkPolynomial for U32 {
51    #[inline(always)]
52    fn apply_mask(block: &mut [u8]) {
53        // §4.4 (ELK: The Encrypted LFSR Keystream Mode of Operation) https://csrc.nist.gov/files/pubs/sp/800/197/iprd/docs/3_samvadini.pdf
54        // Polynomial for n = 256:
55        // X^256 + X^229 + X^193 + X^167 + X^139 + X^109 + X^71 + X^37 + 1
56
57        block[3]  ^= 1 << 5; // X^229
58        block[7]  ^= 1 << 1; // X^193
59        block[11] ^= 1 << 7; // X^167
60        block[14] ^= 1 << 3; // X^139
61        block[18] ^= 1 << 5; // X^109
62        block[23] ^= 1 << 7; // X^71
63        block[27] ^= 1 << 5; // X^37
64        block[31] ^= 1 << 0; // 1 (X^0)
65    }
66}
67
68impl ElkPolynomial for U64 {
69    #[inline(always)]
70    fn apply_mask(block: &mut [u8]) {
71        // §4.4 (ELK: The Encrypted LFSR Keystream Mode of Operation) https://csrc.nist.gov/files/pubs/sp/800/197/iprd/docs/3_samvadini.pdf
72        // Polynomial for n = 512:
73        // X^512 + X^461 + X^397 + X^331 + X^281 + X^211 + X^139 + X^71 + 1
74
75        block[6]  ^= 1 << 5; // X^461
76        block[14] ^= 1 << 5; // X^397
77        block[22] ^= 1 << 3; // X^331
78        block[28] ^= 1 << 1; // X^281
79        block[37] ^= 1 << 3; // X^211
80        block[46] ^= 1 << 3; // X^139
81        block[55] ^= 1 << 7; // X^71
82        block[63] ^= 1 << 0; // 1 (X^0)
83    }
84}
85
86/// Advances the state array using a left-shifting LFSR.
87#[inline]
88fn step_lfsr<C>(state: &mut Block<C>)
89where
90    C: BlockSizeUser,
91    C::BlockSize: ElkPolynomial,
92{
93    let overflow = state[0] >> 7;
94    let len = state.len();
95    
96    // Shift left by 1 bit across the entire array
97    for i in 0..(len - 1) {
98        state[i] = (state[i] << 1) | (state[i + 1] >> 7);
99    }
100    state[len - 1] <<= 1;
101    
102    // Apply Galois Field primitive polynomial if an overflow occurred
103    if overflow != 0 {
104        C::BlockSize::apply_mask(state.as_mut_slice());
105    }
106}
107
108/// Universal ELK (Encrypted LFSR Keystream) mode.
109/// Supports any block cipher with 128, 256, 512, or 1024 bit block sizes.
110#[derive(Clone)]
111#[cfg_attr(feature = "zeroize", derive(ZeroizeOnDrop))]
112pub struct Elk<C>
113where
114    C: BlockCipherEncrypt + BlockSizeUser,
115    C::BlockSize: ElkPolynomial,
116{
117    #[cfg_attr(feature = "zeroize", zeroize(skip))]
118    cipher: C,
119    state: Block<C>,
120}
121
122impl<C> BlockSizeUser for Elk<C>
123where
124    C: BlockCipherEncrypt + BlockSizeUser,
125    C::BlockSize: ElkPolynomial,
126{
127    type BlockSize = C::BlockSize;
128}
129
130impl<C> InnerUser for Elk<C>
131where
132    C: BlockCipherEncrypt + BlockSizeUser,
133    C::BlockSize: ElkPolynomial,
134{
135    type Inner = C;
136}
137
138impl<C> IvSizeUser for Elk<C>
139where
140    C: BlockCipherEncrypt + BlockSizeUser,
141    C::BlockSize: ElkPolynomial,
142{
143    type IvSize = C::BlockSize;
144}
145
146impl<C> InnerIvInit for Elk<C>
147where
148    C: BlockCipherEncrypt + BlockSizeUser,
149    C::BlockSize: ElkPolynomial,
150{
151    #[inline]
152    fn inner_iv_init(cipher: C, iv: &Iv<Self>) -> Self {
153        Self {
154            cipher,
155            state: iv.clone(),
156        }
157    }
158}
159
160impl<C> cipher::AlgorithmName for Elk<C>
161where
162    C: BlockCipherEncrypt + BlockSizeUser + cipher::AlgorithmName,
163    C::BlockSize: ElkPolynomial,
164{
165    fn write_alg_name(f: &mut fmt::Formatter<'_>) -> fmt::Result {
166        f.write_str("Elk<")?;
167        <C as cipher::AlgorithmName>::write_alg_name(f)?;
168        f.write_str(">")
169    }
170}
171
172impl<C> fmt::Debug for Elk<C>
173where
174    C: BlockCipherEncrypt + BlockSizeUser + cipher::AlgorithmName,
175    C::BlockSize: ElkPolynomial,
176{
177    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178        f.write_str("Elk<")?;
179        <C as cipher::AlgorithmName>::write_alg_name(f)?;
180        f.write_str("> { ... }")
181    }
182}
183
184impl<C> Elk<C>
185where
186    C: BlockCipherEncrypt + BlockSizeUser,
187    C::BlockSize: ElkPolynomial,
188{
189    /// Core ELK operation: apply the keystream to the `data` in-place.
190    /// Because ELK is a stream cipher, encryption and decryption are identical.
191    ///
192    /// Must be called exactly once with the complete message. Splitting a
193    /// single logical message across multiple calls does not match the
194    /// ELK specification and will produce incorrect keystream output.
195    pub fn apply_keystream(&mut self, data: &mut [u8]) -> Result<(), Error> {
196        // Prevent catastrophic infinite loop of zeroes (checked first)
197        if self.state.iter().all(|&b| b == 0) {
198            return Err(Error::ZeroIv);
199        }
200
201        if data.is_empty() {
202            return Ok(());
203        }
204
205        let block_size = self.state.len();
206        let m = data.len();
207
208        if m < block_size {
209            // Short message case: m < n.
210            // Step the LFSR before generating the keystream.
211            step_lfsr::<C>(&mut self.state);
212            let mut keystream = self.state.clone();
213            self.cipher.encrypt_block(&mut keystream);
214
215            for (byte, key_byte) in data.iter_mut().zip(keystream.iter()) {
216                *byte ^= key_byte;
217            }
218
219            #[cfg(feature = "zeroize")]
220            {
221                use zeroize::Zeroize;
222                keystream.zeroize();
223            }
224        } else {
225            // Long message case: m >= n.
226            let mut keystream = self.state.clone();
227            let mut chunks = data.chunks_mut(block_size);
228
229            // 1. Process first full block using the unstepped state S
230            if let Some(chunk) = chunks.next() {
231                keystream.clone_from(&self.state);
232                self.cipher.encrypt_block(&mut keystream);
233                for (byte, key_byte) in chunk.iter_mut().zip(keystream.iter()) {
234                    *byte ^= key_byte;
235                }
236            }
237
238            // 2. Process subsequent full/partial blocks
239            let mut next_chunk = chunks.next();
240            while let Some(chunk) = next_chunk {
241                // Determine if this is the last chunk
242                next_chunk = chunks.next();
243
244                // We always step before processing a subsequent block
245                step_lfsr::<C>(&mut self.state);
246                keystream.clone_from(&self.state);
247                self.cipher.encrypt_block(&mut keystream);
248
249                for (byte, key_byte) in chunk.iter_mut().zip(keystream.iter()) {
250                    *byte ^= key_byte;
251                }
252            }
253
254            #[cfg(feature = "zeroize")]
255            {
256                use zeroize::Zeroize;
257                keystream.zeroize();
258            }
259        }
260
261        Ok(())
262    }
263}
264
265#[cfg(kani)]
266mod verification {
267    use super::*;
268    use cipher::{Block, BlockCipherEncrypt, BlockSizeUser};
269    use cipher::typenum::U16;
270
271    #[derive(Clone)]
272    struct MockCipher<Size> {
273        _marker: core::marker::PhantomData<Size>,
274    }
275
276    impl<Size> BlockSizeUser for MockCipher<Size>
277    where
278        Size: hybrid_array::ArraySize,
279    {
280        type BlockSize = Size;
281    }
282
283    impl<Size> BlockCipherEncrypt for MockCipher<Size>
284    where
285        Size: hybrid_array::ArraySize,
286    {
287        fn encrypt_block(&self, _block: &mut Block<Self>) {
288            // Dummy encryption that does nothing
289        }
290
291        fn encrypt_with_backend(&self, _backend: impl cipher::BlockCipherEncClosure<BlockSize = Self::BlockSize>) {
292            // Dummy implementation
293        }
294    }
295
296    #[kani::proof]
297    #[kani::unwind(22)]
298    fn verify_elk_16_multi_block() {
299        let state_bytes: [u8; 16] = kani::any();
300        let state: Block<MockCipher<U16>> = hybrid_array::Array::from(state_bytes);
301        let cipher = MockCipher { _marker: core::marker::PhantomData };
302        let mut elk = Elk { cipher, state };
303
304        // Test up to 3 full blocks (48 bytes) to force the LFSR to step multiple times
305        let mut buffer = [0u8; 48];
306        let data_len: usize = kani::any();
307        kani::assume(data_len <= 48);
308        let data = &mut buffer[..data_len];
309
310        let _ = elk.apply_keystream(data);
311    }
312}