trivium/
trivium.rs

1#[allow(dead_code)]
2const KEY_BYTES: usize = 10;
3#[allow(dead_code)]
4const IV_BYTES: usize = 10;
5
6#[allow(dead_code)]
7#[derive(Debug, Clone, Copy)]
8pub enum BitOrder {
9    Msb,
10    Lsb,
11}
12
13#[allow(dead_code)]
14#[derive(Debug, Clone, Copy)]
15pub enum PackOrder {
16    Lsb,
17} 
18
19// `parse_bit_order` removed: key/IV bits are interpreted MSB-first by default.
20
21
22fn normalize_to_fixed(input: &[u8], target_len: usize) -> [u8; KEY_BYTES] {
23    // NOTE: target_len is always 10 in our usage; keep small and simple.
24    let mut out = [0u8; KEY_BYTES];
25    let n = input.len().min(target_len);
26    out[..n].copy_from_slice(&input[..n]);
27    out
28}
29
30fn get_bit(bytes: &[u8], bit_index: usize, order: BitOrder) -> u8 {
31    let byte_index = bit_index >> 3;
32    let bit_in_byte = bit_index & 7;
33    if byte_index >= bytes.len() {
34        return 0;
35    }
36    let b = bytes[byte_index];
37    match order {
38        // bit_index 0 is the MSB of the first byte
39        BitOrder::Msb => (b >> (7 - bit_in_byte)) & 1,
40        // bit_index 0 is the LSB of the first byte
41        BitOrder::Lsb => (b >> bit_in_byte) & 1,
42    }
43}
44
45#[derive(Clone)]
46pub struct Trivium {
47    s1: [u8; 93],
48    s2: [u8; 84],
49    s3: [u8; 111],
50    pack_order: PackOrder,
51}
52
53impl Trivium {
54    /// Create a new `Trivium` instance.
55    ///
56    /// `load_order` controls how key/IV bits are read within each input byte (MSB or LSB).
57    /// `pack_order` controls how keystream bits are packed into output bytes.
58    pub fn new(key_raw: &[u8], iv_raw: &[u8], load_order: BitOrder, pack_order: PackOrder) -> Self {
59        let key = normalize_to_fixed(key_raw, KEY_BYTES);
60        let iv = normalize_to_fixed(iv_raw, IV_BYTES);
61
62        let mut s1 = [0u8; 93];
63        let mut s2 = [0u8; 84];
64        let mut s3 = [0u8; 111];
65
66        // s1: 80-bit key then 13 zeros
67        for (i, b) in s1.iter_mut().enumerate().take(80) {
68            *b = get_bit(&key, i, load_order);
69        }
70
71        // s2: 80-bit IV then 4 zeros
72        for (i, b) in s2.iter_mut().enumerate().take(80) {
73            *b = get_bit(&iv, i, load_order);
74        }
75
76        // s3: 108 zeros then 3 ones
77        s3[108] = 1;
78        s3[109] = 1;
79        s3[110] = 1;
80
81        let mut st = Self {
82            s1,
83            s2,
84            s3,
85            pack_order,
86        };
87
88        // Warm-up: 4 * 288 = 1152 steps
89        for _ in 0..(4 * 288) {
90            let _ = st.next_bit();
91        }
92
93        st
94    }
95
96    #[inline]
97    fn shift_in(reg: &mut [u8], new_bit: u8) {
98        for i in (1..reg.len()).rev() {
99            reg[i] = reg[i - 1];
100        }
101        reg[0] = new_bit & 1;
102    }
103
104    fn next_bit(&mut self) -> u8 {
105        // Indices are 0-based (spec positions minus 1)
106        let t1 = self.s1[65] ^ self.s1[92];
107        let t2 = self.s2[68] ^ self.s2[83];
108        let t3 = self.s3[65] ^ self.s3[110];
109        let z = t1 ^ t2 ^ t3;
110
111        let t1n = t1 ^ (self.s1[90] & self.s1[91]) ^ self.s2[77];
112        let t2n = t2 ^ (self.s2[81] & self.s2[82]) ^ self.s3[86];
113        let t3n = t3 ^ (self.s3[108] & self.s3[109]) ^ self.s1[68];
114
115        Self::shift_in(&mut self.s1, t3n);
116        Self::shift_in(&mut self.s2, t1n);
117        Self::shift_in(&mut self.s3, t2n);
118
119        z & 1
120    }
121
122    fn keystream_byte(&mut self) -> u8 {
123        match self.pack_order {
124            PackOrder::Lsb => {
125                // Pack bits LSB-first within each output byte: first output bit becomes bit 0.
126                let mut b = 0u8;
127                for j in 0..8 {
128                    b |= (self.next_bit() & 1) << j;
129                }
130                b
131            }
132        }
133    }
134
135    /// XOR `data` with the Trivium keystream produced by this instance.
136    ///
137    /// Consumes the `Trivium` instance (it owns the state progression).
138    pub fn xor_bytes(mut self, data: &[u8]) -> Vec<u8> {
139        let mut out = Vec::with_capacity(data.len());
140        for &v in data {
141            let k = self.keystream_byte();
142            out.push(v ^ k);
143        }
144        out
145    }
146}
147
148/// Trivium stream cipher XOR (encrypt/decrypt)
149///
150/// - `key` / `iv`: arbitrary-length bytes; will be truncated/padded to 10 bytes (80-bit)
151/// - `data`: input bytes
152/// - Key/IV bits are read MSB-first within each byte.
153///
154/// # Examples
155/// ```
156/// let key = vec![0u8; 10];
157/// let iv = vec![0u8; 10];
158/// let data = b"hello".to_vec();
159/// let ct = trivium::trivium_xor(key.clone(), iv.clone(), data.clone()).unwrap();
160/// let pt = trivium::trivium_xor(key, iv, ct).unwrap();
161/// assert_eq!(pt, data);
162/// ```
163pub fn trivium_xor(
164    key: Vec<u8>,
165    iv: Vec<u8>,
166    data: Vec<u8>,
167) -> Result<Vec<u8>, String> {
168    // Default to MSB when reading key/IV bits within each byte.
169    let trivium = Trivium::new(&key, &iv, BitOrder::Msb, PackOrder::Lsb);
170    Ok(trivium.xor_bytes(&data))
171} 
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    fn hex_to_bytes(s: &str) -> Vec<u8> {
178        let clean: String = s
179            .chars()
180            .filter(|c| !c.is_whitespace())
181            .collect::<String>()
182            .trim_start_matches("0x")
183            .to_string();
184        assert!(clean.len().is_multiple_of(2), "hex must have even length");
185        (0..clean.len())
186            .step_by(2)
187            .map(|i| u8::from_str_radix(&clean[i..i + 2], 16).expect("hex byte"))
188            .collect()
189    }
190
191    fn make_prng(seed: u32) -> impl FnMut() -> u32 {
192        let mut x = seed;
193        move || {
194            x ^= x << 13;
195            x ^= x >> 17;
196            x ^= x << 5;
197            x
198        }
199    }
200
201    fn random_bytes(mut prng: impl FnMut() -> u32, len: usize) -> Vec<u8> {
202        let mut out = vec![0u8; len];
203        for b in &mut out {
204            *b = (prng() & 0xff) as u8;
205        }
206        out
207    }
208
209    #[test]
210    fn xor_is_symmetric_msb() {
211        let mut prng = make_prng(0x1234_5678);
212
213        for _ in 0..64 {
214            let key = random_bytes(&mut prng, 10);
215            let iv = random_bytes(&mut prng, 10);
216            let msg_len = (prng() as usize) % 512;
217            let msg = random_bytes(&mut prng, msg_len);
218
219            let ct = trivium_xor(key.clone(), iv.clone(), msg.clone()).expect("encrypt");
220            let pt = trivium_xor(key, iv, ct).expect("decrypt");
221
222            assert_eq!(pt, msg);
223        }
224    }
225
226    #[test]
227    fn xor_is_symmetric_lsb() {
228        let mut prng = make_prng(0x9e37_79b9);
229
230        for _ in 0..64 {
231            let key = random_bytes(&mut prng, 10);
232            let iv = random_bytes(&mut prng, 10);
233            let msg_len = (prng() as usize) % 512;
234            let msg = random_bytes(&mut prng, msg_len);
235
236            let ct = Trivium::new(&key, &iv, BitOrder::Lsb, PackOrder::Lsb).xor_bytes(&msg);
237            let pt = Trivium::new(&key, &iv, BitOrder::Lsb, PackOrder::Lsb).xor_bytes(&ct);
238
239            assert_eq!(pt, msg);
240        }
241    }
242
243    #[test]
244    fn accepts_short_key_iv_by_padding_with_zeros() {
245        let key = vec![1, 2, 3];
246        let iv = vec![4, 5];
247        let msg = b"hello".to_vec();
248
249        let ct = trivium_xor(key.clone(), iv.clone(), msg.clone()).unwrap();
250        let pt = trivium_xor(key, iv, ct).unwrap();
251
252        assert_eq!(pt, msg);
253    }
254
255    #[test]
256    fn known_vector_hello_key_iv() {
257        // User-provided known-answer test:
258        // plaintext: "hello"
259        // key (hex): 00000010000000000000
260        // iv  (hex): 00000000000000000000
261        // expected ciphertext (hex): 9f804f6861
262        let key = hex_to_bytes("00000010000000000000");
263        let iv = hex_to_bytes("00000000000000000000");
264        let msg = b"hello".to_vec();
265
266        let expected = hex_to_bytes("9f804f6861");
267
268        // This project interprets HEX key/IV as big-endian 80-bit values in the UI,
269        // then reverses the 10-byte sequence before sending to Rust.
270        // The backend reads bits MSB-first (bitOrder=msb) and packs keystream bytes LSB-first.
271        let mut key_reversed = key.clone();
272        key_reversed.reverse();
273
274        let ct = trivium_xor(key_reversed, iv, msg).unwrap();
275        assert_eq!(ct, expected);
276    }
277}
trivium/trivium.rs

trivium/
trivium.rs
