miri/intrinsics/x86/
aesni.rs

1use rustc_middle::ty::Ty;
2use rustc_span::Symbol;
3
4use crate::*;
5
6impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
7pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
8    fn emulate_x86_aesni_intrinsic(
9        &mut self,
10        link_name: Symbol,
11        args: &[OpTy<'tcx>],
12        dest: &MPlaceTy<'tcx>,
13    ) -> InterpResult<'tcx, EmulateItemResult> {
14        let this = self.eval_context_mut();
15        this.expect_target_feature_for_intrinsic(link_name, "aes")?;
16        // Prefix should have already been checked.
17        let unprefixed_name = link_name.as_str().strip_prefix("llvm.x86.aesni.").unwrap();
18
19        match unprefixed_name {
20            // Used to implement the _mm_aesdec_si128, _mm256_aesdec_epi128
21            // and _mm512_aesdec_epi128 functions.
22            // Performs one round of an AES decryption on each 128-bit word of
23            // `state` with the corresponding 128-bit key of `key`.
24            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesdec_si128
25            "aesdec" | "aesdec.256" | "aesdec.512" => {
26                let [state, key] = this.check_shim_sig_llvm_intrinsic(link_name, args)?;
27                aes_round(this, state, key, dest, |state, key| {
28                    let key = aes::Block::from(key.to_le_bytes());
29                    let mut state = aes::Block::from(state.to_le_bytes());
30                    // `aes::hazmat::equiv_inv_cipher_round` documentation states that
31                    // it performs the same operation as the x86 aesdec instruction.
32                    aes::hazmat::equiv_inv_cipher_round(&mut state, &key);
33                    u128::from_le_bytes(state.into())
34                })?;
35            }
36            // Used to implement the _mm_aesdeclast_si128, _mm256_aesdeclast_epi128
37            // and _mm512_aesdeclast_epi128 functions.
38            // Performs last round of an AES decryption on each 128-bit word of
39            // `state` with the corresponding 128-bit key of `key`.
40            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesdeclast_si128
41            "aesdeclast" | "aesdeclast.256" | "aesdeclast.512" => {
42                let [state, key] = this.check_shim_sig_llvm_intrinsic(link_name, args)?;
43
44                aes_round(this, state, key, dest, |state, key| {
45                    let mut state = aes::Block::from(state.to_le_bytes());
46                    // `aes::hazmat::equiv_inv_cipher_round` does the following operations:
47                    // state = InvShiftRows(state)
48                    // state = InvSubBytes(state)
49                    // state = InvMixColumns(state)
50                    // state = state ^ key
51                    // But we need to skip the InvMixColumns.
52                    // First, use a zeroed key to skip the XOR.
53                    aes::hazmat::equiv_inv_cipher_round(&mut state, &aes::Block::from([0; 16]));
54                    // Then, undo the InvMixColumns with MixColumns.
55                    aes::hazmat::mix_columns(&mut state);
56                    // Finally, do the XOR.
57                    u128::from_le_bytes(state.into()) ^ key
58                })?;
59            }
60            // Used to implement the _mm_aesenc_si128, _mm256_aesenc_epi128
61            // and _mm512_aesenc_epi128 functions.
62            // Performs one round of an AES encryption on each 128-bit word of
63            // `state` with the corresponding 128-bit key of `key`.
64            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesenc_si128
65            "aesenc" | "aesenc.256" | "aesenc.512" => {
66                let [state, key] = this.check_shim_sig_llvm_intrinsic(link_name, args)?;
67                aes_round(this, state, key, dest, |state, key| {
68                    let key = aes::Block::from(key.to_le_bytes());
69                    let mut state = aes::Block::from(state.to_le_bytes());
70                    // `aes::hazmat::cipher_round` documentation states that
71                    // it performs the same operation as the x86 aesenc instruction.
72                    aes::hazmat::cipher_round(&mut state, &key);
73                    u128::from_le_bytes(state.into())
74                })?;
75            }
76            // Used to implement the _mm_aesenclast_si128, _mm256_aesenclast_epi128
77            // and _mm512_aesenclast_epi128 functions.
78            // Performs last round of an AES encryption on each 128-bit word of
79            // `state` with the corresponding 128-bit key of `key`.
80            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesenclast_si128
81            "aesenclast" | "aesenclast.256" | "aesenclast.512" => {
82                let [state, key] = this.check_shim_sig_llvm_intrinsic(link_name, args)?;
83                aes_round(this, state, key, dest, |state, key| {
84                    let mut state = aes::Block::from(state.to_le_bytes());
85                    // `aes::hazmat::cipher_round` does the following operations:
86                    // state = ShiftRows(state)
87                    // state = SubBytes(state)
88                    // state = MixColumns(state)
89                    // state = state ^ key
90                    // But we need to skip the MixColumns.
91                    // First, use a zeroed key to skip the XOR.
92                    aes::hazmat::cipher_round(&mut state, &aes::Block::from([0; 16]));
93                    // Then, undo the MixColumns with InvMixColumns.
94                    aes::hazmat::inv_mix_columns(&mut state);
95                    // Finally, do the XOR.
96                    u128::from_le_bytes(state.into()) ^ key
97                })?;
98            }
99            // Used to implement the _mm_aesimc_si128 function.
100            // Performs the AES InvMixColumns operation on `op`
101            "aesimc" => {
102                let [op] = this.check_shim_sig_llvm_intrinsic(link_name, args)?;
103                // Transmute to `u128`
104                let op = op.transmute(this.machine.layouts.u128, this)?;
105                let dest = dest.transmute(this.machine.layouts.u128, this)?;
106
107                let state = this.read_scalar(&op)?.to_u128()?;
108                let mut state = aes::Block::from(state.to_le_bytes());
109                aes::hazmat::inv_mix_columns(&mut state);
110
111                this.write_scalar(Scalar::from_u128(u128::from_le_bytes(state.into())), &dest)?;
112            }
113            // Used to implement the _mm_aeskeygenassist_si128 function.
114            // Assists in expanding the AES cipher key.
115            "aeskeygenassist" => {
116                let [ckey, rcon] = this.check_shim_sig_llvm_intrinsic(link_name, args)?;
117                // Transmute `__m128i` to `u128`.
118                let ckey = ckey.transmute(this.machine.layouts.u128, this)?;
119                let dest = dest.transmute(this.machine.layouts.u128, this)?;
120
121                let rcon = this.read_scalar(rcon)?.to_u8()?;
122                let ckey = this.read_scalar(&ckey)?.to_u128()?;
123
124                let res = aeskeygenassist(ckey, rcon);
125
126                this.write_scalar(Scalar::from_u128(res), &dest)?;
127            }
128            _ => return interp_ok(EmulateItemResult::NotSupported),
129        }
130        interp_ok(EmulateItemResult::NeedsReturn)
131    }
132}
133
134// Performs an AES round (given by `f`) on each 128-bit word of
135// `state` with the corresponding 128-bit key of `key`.
136fn aes_round<'tcx>(
137    ecx: &mut crate::MiriInterpCx<'tcx>,
138    state: &OpTy<'tcx>,
139    key: &OpTy<'tcx>,
140    dest: &MPlaceTy<'tcx>,
141    f: impl Fn(u128, u128) -> u128,
142) -> InterpResult<'tcx, ()> {
143    assert_eq!(dest.layout.size, state.layout.size);
144    assert_eq!(dest.layout.size, key.layout.size);
145
146    // Transmute arguments to arrays of `u128`.
147    assert_eq!(dest.layout.size.bytes() % 16, 0);
148    let len = dest.layout.size.bytes() / 16;
149
150    let u128_array_layout = ecx.layout_of(Ty::new_array(ecx.tcx.tcx, ecx.tcx.types.u128, len))?;
151
152    let state = state.transmute(u128_array_layout, ecx)?;
153    let key = key.transmute(u128_array_layout, ecx)?;
154    let dest = dest.transmute(u128_array_layout, ecx)?;
155
156    for i in 0..len {
157        let state = ecx.read_scalar(&ecx.project_index(&state, i)?)?.to_u128()?;
158        let key = ecx.read_scalar(&ecx.project_index(&key, i)?)?.to_u128()?;
159        let dest = ecx.project_index(&dest, i)?;
160
161        let res = f(state, key);
162
163        ecx.write_scalar(Scalar::from_u128(res), &dest)?;
164    }
165
166    interp_ok(())
167}
168
169/// AES Key Generation Assist
170///
171/// From [Intel Intrinsics Guide][1]:
172/// ```text
173/// X3[31:0] := a[127:96]
174/// X2[31:0] := a[95:64]
175/// X1[31:0] := a[63:32]
176/// X0[31:0] := a[31:0]
177/// RCON[31:0] := ZeroExtend32(imm8[7:0])
178/// dst[31:0] := SubWord(X1)
179/// dst[63:32] := RotWord(SubWord(X1)) XOR RCON
180/// dst[95:64] := SubWord(X3)
181/// dst[127:96] := RotWord(SubWord(X3)) XOR RCON
182/// ```
183///
184/// [1]: https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aeskeygenassist_si128
185#[expect(clippy::as_conversions, reason = "deliberately truncating")]
186fn aeskeygenassist(a: u128, rcon: u8) -> u128 {
187    use crate::intrinsics::math::aes::sub_word;
188
189    let rcon = u32::from(rcon);
190    // TODO: use `truncate` method on stabilization
191    let x1 = (a >> 32) as u32;
192    let x3 = (a >> 96) as u32;
193
194    let x0 = sub_word(x1);
195    let x1 = x0.rotate_right(8) ^ rcon;
196    let x2 = sub_word(x3);
197    let x3 = x2.rotate_right(8) ^ rcon;
198
199    (u128::from(x3) << 96) | (u128::from(x2) << 64) | (u128::from(x1) << 32) | u128::from(x0)
200}
miri/intrinsics/x86/aesni.rs

miri/intrinsics/x86/
aesni.rs
