subspace_proof_of_time/
aes.rs

1//! AES related functionality.
2
3#[cfg(target_arch = "aarch64")]
4mod aarch64;
5#[cfg(target_arch = "x86_64")]
6mod x86_64;
7
8use aes::Aes128;
9use aes::cipher::array::Array;
10use aes::cipher::{BlockCipherDecrypt, BlockCipherEncrypt, KeyInit};
11use subspace_core_primitives::pot::{PotCheckpoints, PotKey, PotOutput, PotSeed};
12
13/// Creates the AES based proof.
14#[inline(always)]
15pub(crate) fn create(seed: PotSeed, key: PotKey, checkpoint_iterations: u32) -> PotCheckpoints {
16    #[cfg(target_arch = "x86_64")]
17    {
18        cpufeatures::new!(has_aes, "aes");
19        if has_aes::get() {
20            // SAFETY: Checked `aes` feature
21            return unsafe { x86_64::create(seed.as_ref(), key.as_ref(), checkpoint_iterations) };
22        }
23    }
24    #[cfg(target_arch = "aarch64")]
25    {
26        cpufeatures::new!(has_aes, "aes");
27        if has_aes::get() {
28            // SAFETY: Checked `aes` feature
29            return unsafe { aarch64::create(seed.as_ref(), key.as_ref(), checkpoint_iterations) };
30        }
31    }
32
33    create_generic(seed, key, checkpoint_iterations)
34}
35
36fn create_generic(seed: PotSeed, key: PotKey, checkpoint_iterations: u32) -> PotCheckpoints {
37    let key = Array::from(*key);
38    let cipher = Aes128::new(&key);
39    let mut cur_block = Array::from(*seed);
40
41    let mut checkpoints = PotCheckpoints::default();
42    for checkpoint in checkpoints.iter_mut() {
43        for _ in 0..checkpoint_iterations {
44            // Encrypt in place to produce the next block.
45            cipher.encrypt_block(&mut cur_block);
46        }
47        checkpoint.copy_from_slice(&cur_block);
48    }
49
50    checkpoints
51}
52
53/// Verifies the AES based proof sequentially.
54///
55/// Panics if `checkpoint_iterations` is not a multiple of `2`.
56#[inline(always)]
57pub(crate) fn verify_sequential(
58    seed: PotSeed,
59    key: PotKey,
60    checkpoints: &PotCheckpoints,
61    checkpoint_iterations: u32,
62) -> bool {
63    assert_eq!(checkpoint_iterations % 2, 0);
64
65    #[cfg(target_arch = "x86_64")]
66    {
67        cpufeatures::new!(has_avx512f_vaes, "avx512f", "vaes");
68        if has_avx512f_vaes::get() {
69            // SAFETY: Checked `avx512f` and `vaes` features
70            return unsafe {
71                x86_64::verify_sequential_avx512f_vaes(
72                    &seed,
73                    &key,
74                    checkpoints,
75                    checkpoint_iterations,
76                )
77            };
78        }
79
80        cpufeatures::new!(has_avx2_vaes, "avx2", "vaes");
81        if has_avx2_vaes::get() {
82            // SAFETY: Checked `avx2` and `vaes` features
83            return unsafe {
84                x86_64::verify_sequential_avx2_vaes(&seed, &key, checkpoints, checkpoint_iterations)
85            };
86        }
87
88        cpufeatures::new!(has_aes_sse41, "aes", "sse4.1");
89        if has_aes_sse41::get() {
90            // SAFETY: Checked `aes` and `sse4.1` features
91            return unsafe {
92                x86_64::verify_sequential_aes_sse41(&seed, &key, checkpoints, checkpoint_iterations)
93            };
94        }
95    }
96    #[cfg(target_arch = "aarch64")]
97    {
98        cpufeatures::new!(has_aes, "aes");
99        if has_aes::get() {
100            // SAFETY: Checked `aes` feature
101            return unsafe {
102                aarch64::verify_sequential_aes(&seed, &key, checkpoints, checkpoint_iterations)
103            };
104        }
105    }
106
107    verify_sequential_generic(seed, key, checkpoints, checkpoint_iterations)
108}
109
110fn verify_sequential_generic(
111    seed: PotSeed,
112    key: PotKey,
113    checkpoints: &PotCheckpoints,
114    checkpoint_iterations: u32,
115) -> bool {
116    let key = Array::from(*key);
117    let cipher = Aes128::new(&key);
118
119    let mut inputs = [[0u8; 16]; PotCheckpoints::NUM_CHECKPOINTS.get() as usize];
120    inputs[0] = *seed;
121    inputs[1..].copy_from_slice(PotOutput::repr_from_slice(
122        &checkpoints[..PotCheckpoints::NUM_CHECKPOINTS.get() as usize - 1],
123    ));
124
125    let mut outputs = [[0u8; 16]; PotCheckpoints::NUM_CHECKPOINTS.get() as usize];
126    outputs.copy_from_slice(PotOutput::repr_from_slice(checkpoints.as_slice()));
127
128    for _ in 0..checkpoint_iterations / 2 {
129        cipher.encrypt_blocks(Array::cast_slice_from_core_mut(&mut inputs));
130        cipher.decrypt_blocks(Array::cast_slice_from_core_mut(&mut outputs));
131    }
132
133    inputs == outputs
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use subspace_core_primitives::pot::PotOutput;
140
141    const SEED: [u8; 16] = [
142        0xd6, 0x66, 0xcc, 0xd8, 0xd5, 0x93, 0xc2, 0x3d, 0xa8, 0xdb, 0x6b, 0x5b, 0x14, 0x13, 0xb1,
143        0x3a,
144    ];
145    const SEED_1: [u8; 16] = [
146        0xd7, 0xd6, 0xdc, 0xd8, 0xd5, 0x93, 0xc2, 0x3d, 0xa8, 0xdb, 0x6b, 0x5b, 0x14, 0x13, 0xb1,
147        0x3a,
148    ];
149    const KEY: [u8; 16] = [
150        0x9a, 0x84, 0x94, 0x0f, 0xfe, 0xf5, 0xb0, 0xd7, 0x01, 0x99, 0xfc, 0x67, 0xf4, 0x6e, 0xa2,
151        0x7a,
152    ];
153    const KEY_1: [u8; 16] = [
154        0x9b, 0x8b, 0x9b, 0x0f, 0xfe, 0xf5, 0xb0, 0xd7, 0x01, 0x99, 0xfc, 0x67, 0xf4, 0x6e, 0xa2,
155        0x7a,
156    ];
157    const BAD_CIPHER: [u8; 16] = [22; 16];
158
159    fn verify_test(
160        seed: PotSeed,
161        key: PotKey,
162        checkpoints: &PotCheckpoints,
163        checkpoint_iterations: u32,
164    ) -> bool {
165        let sequential = verify_sequential(seed, key, checkpoints, checkpoint_iterations);
166        let generic = verify_sequential_generic(seed, key, checkpoints, checkpoint_iterations);
167        assert_eq!(sequential, generic);
168
169        #[cfg(target_arch = "x86_64")]
170        {
171            cpufeatures::new!(has_avx512f_vaes, "avx512f", "vaes");
172            if has_avx512f_vaes::get() {
173                // SAFETY: Checked `avx512f` and `vaes` features
174                let avx512f_vaes = unsafe {
175                    x86_64::verify_sequential_avx512f_vaes(
176                        &seed,
177                        &key,
178                        checkpoints,
179                        checkpoint_iterations,
180                    )
181                };
182                assert_eq!(sequential, avx512f_vaes);
183            }
184
185            cpufeatures::new!(has_avx2_vaes, "avx2", "vaes");
186            if has_avx2_vaes::get() {
187                // SAFETY: Checked `avx2` and `vaes` features
188                let avx2_vaes = unsafe {
189                    x86_64::verify_sequential_avx2_vaes(
190                        &seed,
191                        &key,
192                        checkpoints,
193                        checkpoint_iterations,
194                    )
195                };
196                assert_eq!(sequential, avx2_vaes);
197            }
198
199            cpufeatures::new!(has_aes_sse41, "aes", "sse4.1");
200            if has_aes_sse41::get() {
201                // SAFETY: Checked `aes` and `sse4.1` features
202                let aes_sse41 = unsafe {
203                    x86_64::verify_sequential_aes_sse41(
204                        &seed,
205                        &key,
206                        checkpoints,
207                        checkpoint_iterations,
208                    )
209                };
210                assert_eq!(sequential, aes_sse41);
211            }
212        }
213        #[cfg(target_arch = "aarch64")]
214        {
215            cpufeatures::new!(has_aes, "aes");
216            if has_aes::get() {
217                // SAFETY: Checked `aes` feature
218                let aes = unsafe {
219                    aarch64::verify_sequential_aes(&seed, &key, checkpoints, checkpoint_iterations)
220                };
221                assert_eq!(sequential, aes);
222            }
223        }
224
225        sequential
226    }
227
228    #[test]
229    fn test_create_verify() {
230        let seed = PotSeed::from(SEED);
231        let key = PotKey::from(KEY);
232        let checkpoint_iterations = 20;
233
234        // Can encrypt/decrypt.
235        let checkpoints = create(seed, key, checkpoint_iterations);
236        {
237            let generic_checkpoints = create_generic(seed, key, checkpoint_iterations);
238            assert_eq!(checkpoints, generic_checkpoints);
239        }
240
241        assert!(verify_test(seed, key, &checkpoints, checkpoint_iterations,));
242
243        // Decryption of invalid cipher text fails.
244        let mut checkpoints_1 = checkpoints;
245        checkpoints_1[0] = PotOutput::from(BAD_CIPHER);
246        assert!(!verify_test(
247            seed,
248            key,
249            &checkpoints_1,
250            checkpoint_iterations,
251        ));
252
253        // Decryption with wrong number of iterations fails.
254        assert!(!verify_test(
255            seed,
256            key,
257            &checkpoints,
258            checkpoint_iterations + 2,
259        ));
260        assert!(!verify_test(
261            seed,
262            key,
263            &checkpoints,
264            checkpoint_iterations - 2,
265        ));
266
267        // Decryption with wrong seed fails.
268        assert!(!verify_test(
269            PotSeed::from(SEED_1),
270            key,
271            &checkpoints,
272            checkpoint_iterations,
273        ));
274
275        // Decryption with wrong key fails.
276        assert!(!verify_test(
277            seed,
278            PotKey::from(KEY_1),
279            &checkpoints,
280            checkpoint_iterations,
281        ));
282    }
283}
subspace_proof_of_time/aes.rs

subspace_proof_of_time/
aes.rs
