subspace_proof_of_time/aes/
x86_64.rs

1use core::arch::x86_64::*;
2use core::mem;
3use subspace_core_primitives::pot::PotCheckpoints;
4
5/// Create PoT proof with checkpoints
6#[target_feature(enable = "aes")]
7#[inline]
8pub(super) unsafe fn create(
9    seed: &[u8; 16],
10    key: &[u8; 16],
11    checkpoint_iterations: u32,
12) -> PotCheckpoints {
13    let mut checkpoints = PotCheckpoints::default();
14
15    let keys_reg = expand_key(key);
16    let xor_key = _mm_xor_si128(keys_reg[10], keys_reg[0]);
17    let mut seed_reg = _mm_loadu_si128(seed.as_ptr() as *const __m128i);
18    seed_reg = _mm_xor_si128(seed_reg, keys_reg[0]);
19    for checkpoint in checkpoints.iter_mut() {
20        for _ in 0..checkpoint_iterations {
21            seed_reg = _mm_aesenc_si128(seed_reg, keys_reg[1]);
22            seed_reg = _mm_aesenc_si128(seed_reg, keys_reg[2]);
23            seed_reg = _mm_aesenc_si128(seed_reg, keys_reg[3]);
24            seed_reg = _mm_aesenc_si128(seed_reg, keys_reg[4]);
25            seed_reg = _mm_aesenc_si128(seed_reg, keys_reg[5]);
26            seed_reg = _mm_aesenc_si128(seed_reg, keys_reg[6]);
27            seed_reg = _mm_aesenc_si128(seed_reg, keys_reg[7]);
28            seed_reg = _mm_aesenc_si128(seed_reg, keys_reg[8]);
29            seed_reg = _mm_aesenc_si128(seed_reg, keys_reg[9]);
30            seed_reg = _mm_aesenclast_si128(seed_reg, xor_key);
31        }
32
33        let checkpoint_reg = _mm_xor_si128(seed_reg, keys_reg[0]);
34        _mm_storeu_si128(
35            checkpoint.as_mut().as_mut_ptr() as *mut __m128i,
36            checkpoint_reg,
37        );
38    }
39
40    checkpoints
41}
42
43// Below code copied with minor changes from following place under MIT/Apache-2.0 license by Artyom
44// Pavlov:
45// https://github.com/RustCrypto/block-ciphers/blob/9413fcadd28d53854954498c0589b747d8e4ade2/aes/src/ni/aes128.rs
46
47/// AES-128 round keys
48type RoundKeys = [__m128i; 11];
49
50macro_rules! expand_round {
51    ($keys:expr, $pos:expr, $round:expr) => {
52        let mut t1 = $keys[$pos - 1];
53        let mut t2;
54        let mut t3;
55
56        t2 = _mm_aeskeygenassist_si128(t1, $round);
57        t2 = _mm_shuffle_epi32(t2, 0xff);
58        t3 = _mm_slli_si128(t1, 0x4);
59        t1 = _mm_xor_si128(t1, t3);
60        t3 = _mm_slli_si128(t3, 0x4);
61        t1 = _mm_xor_si128(t1, t3);
62        t3 = _mm_slli_si128(t3, 0x4);
63        t1 = _mm_xor_si128(t1, t3);
64        t1 = _mm_xor_si128(t1, t2);
65
66        $keys[$pos] = t1;
67    };
68}
69
70#[target_feature(enable = "aes")]
71#[inline]
72unsafe fn expand_key(key: &[u8; 16]) -> RoundKeys {
73    // SAFETY: `RoundKeys` is a `[__m128i; 11]` which can be initialized
74    // with all zeroes.
75    let mut keys: RoundKeys = mem::zeroed();
76
77    let k = _mm_loadu_si128(key.as_ptr() as *const __m128i);
78    keys[0] = k;
79
80    expand_round!(keys, 1, 0x01);
81    expand_round!(keys, 2, 0x02);
82    expand_round!(keys, 3, 0x04);
83    expand_round!(keys, 4, 0x08);
84    expand_round!(keys, 5, 0x10);
85    expand_round!(keys, 6, 0x20);
86    expand_round!(keys, 7, 0x40);
87    expand_round!(keys, 8, 0x80);
88    expand_round!(keys, 9, 0x1B);
89    expand_round!(keys, 10, 0x36);
90
91    keys
92}