miri/shims/x86/
aesni.rs

1use rustc_abi::CanonAbi;
2use rustc_middle::ty::Ty;
3use rustc_middle::ty::layout::LayoutOf as _;
4use rustc_span::Symbol;
5use rustc_target::callconv::FnAbi;
6
7use crate::*;
8
9impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
10pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
11    fn emulate_x86_aesni_intrinsic(
12        &mut self,
13        link_name: Symbol,
14        abi: &FnAbi<'tcx, Ty<'tcx>>,
15        args: &[OpTy<'tcx>],
16        dest: &MPlaceTy<'tcx>,
17    ) -> InterpResult<'tcx, EmulateItemResult> {
18        let this = self.eval_context_mut();
19        this.expect_target_feature_for_intrinsic(link_name, "aes")?;
20        // Prefix should have already been checked.
21        let unprefixed_name = link_name.as_str().strip_prefix("llvm.x86.aesni.").unwrap();
22
23        match unprefixed_name {
24            // Used to implement the _mm_aesdec_si128, _mm256_aesdec_epi128
25            // and _mm512_aesdec_epi128 functions.
26            // Performs one round of an AES decryption on each 128-bit word of
27            // `state` with the corresponding 128-bit key of `key`.
28            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesdec_si128
29            "aesdec" | "aesdec.256" | "aesdec.512" => {
30                let [state, key] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
31                aes_round(this, state, key, dest, |state, key| {
32                    let key = aes::Block::from(key.to_le_bytes());
33                    let mut state = aes::Block::from(state.to_le_bytes());
34                    // `aes::hazmat::equiv_inv_cipher_round` documentation states that
35                    // it performs the same operation as the x86 aesdec instruction.
36                    aes::hazmat::equiv_inv_cipher_round(&mut state, &key);
37                    u128::from_le_bytes(state.into())
38                })?;
39            }
40            // Used to implement the _mm_aesdeclast_si128, _mm256_aesdeclast_epi128
41            // and _mm512_aesdeclast_epi128 functions.
42            // Performs last round of an AES decryption on each 128-bit word of
43            // `state` with the corresponding 128-bit key of `key`.
44            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesdeclast_si128
45            "aesdeclast" | "aesdeclast.256" | "aesdeclast.512" => {
46                let [state, key] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
47
48                aes_round(this, state, key, dest, |state, key| {
49                    let mut state = aes::Block::from(state.to_le_bytes());
50                    // `aes::hazmat::equiv_inv_cipher_round` does the following operations:
51                    // state = InvShiftRows(state)
52                    // state = InvSubBytes(state)
53                    // state = InvMixColumns(state)
54                    // state = state ^ key
55                    // But we need to skip the InvMixColumns.
56                    // First, use a zeroed key to skip the XOR.
57                    aes::hazmat::equiv_inv_cipher_round(&mut state, &aes::Block::from([0; 16]));
58                    // Then, undo the InvMixColumns with MixColumns.
59                    aes::hazmat::mix_columns(&mut state);
60                    // Finally, do the XOR.
61                    u128::from_le_bytes(state.into()) ^ key
62                })?;
63            }
64            // Used to implement the _mm_aesenc_si128, _mm256_aesenc_epi128
65            // and _mm512_aesenc_epi128 functions.
66            // Performs one round of an AES encryption on each 128-bit word of
67            // `state` with the corresponding 128-bit key of `key`.
68            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesenc_si128
69            "aesenc" | "aesenc.256" | "aesenc.512" => {
70                let [state, key] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
71                aes_round(this, state, key, dest, |state, key| {
72                    let key = aes::Block::from(key.to_le_bytes());
73                    let mut state = aes::Block::from(state.to_le_bytes());
74                    // `aes::hazmat::cipher_round` documentation states that
75                    // it performs the same operation as the x86 aesenc instruction.
76                    aes::hazmat::cipher_round(&mut state, &key);
77                    u128::from_le_bytes(state.into())
78                })?;
79            }
80            // Used to implement the _mm_aesenclast_si128, _mm256_aesenclast_epi128
81            // and _mm512_aesenclast_epi128 functions.
82            // Performs last round of an AES encryption on each 128-bit word of
83            // `state` with the corresponding 128-bit key of `key`.
84            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_aesenclast_si128
85            "aesenclast" | "aesenclast.256" | "aesenclast.512" => {
86                let [state, key] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
87                aes_round(this, state, key, dest, |state, key| {
88                    let mut state = aes::Block::from(state.to_le_bytes());
89                    // `aes::hazmat::cipher_round` does the following operations:
90                    // state = ShiftRows(state)
91                    // state = SubBytes(state)
92                    // state = MixColumns(state)
93                    // state = state ^ key
94                    // But we need to skip the MixColumns.
95                    // First, use a zeroed key to skip the XOR.
96                    aes::hazmat::cipher_round(&mut state, &aes::Block::from([0; 16]));
97                    // Then, undo the MixColumns with InvMixColumns.
98                    aes::hazmat::inv_mix_columns(&mut state);
99                    // Finally, do the XOR.
100                    u128::from_le_bytes(state.into()) ^ key
101                })?;
102            }
103            // Used to implement the _mm_aesimc_si128 function.
104            // Performs the AES InvMixColumns operation on `op`
105            "aesimc" => {
106                let [op] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
107                // Transmute to `u128`
108                let op = op.transmute(this.machine.layouts.u128, this)?;
109                let dest = dest.transmute(this.machine.layouts.u128, this)?;
110
111                let state = this.read_scalar(&op)?.to_u128()?;
112                let mut state = aes::Block::from(state.to_le_bytes());
113                aes::hazmat::inv_mix_columns(&mut state);
114
115                this.write_scalar(Scalar::from_u128(u128::from_le_bytes(state.into())), &dest)?;
116            }
117            // TODO: Implement the `llvm.x86.aesni.aeskeygenassist` when possible
118            // with an external crate.
119            _ => return interp_ok(EmulateItemResult::NotSupported),
120        }
121        interp_ok(EmulateItemResult::NeedsReturn)
122    }
123}
124
125// Performs an AES round (given by `f`) on each 128-bit word of
126// `state` with the corresponding 128-bit key of `key`.
127fn aes_round<'tcx>(
128    ecx: &mut crate::MiriInterpCx<'tcx>,
129    state: &OpTy<'tcx>,
130    key: &OpTy<'tcx>,
131    dest: &MPlaceTy<'tcx>,
132    f: impl Fn(u128, u128) -> u128,
133) -> InterpResult<'tcx, ()> {
134    assert_eq!(dest.layout.size, state.layout.size);
135    assert_eq!(dest.layout.size, key.layout.size);
136
137    // Transmute arguments to arrays of `u128`.
138    assert_eq!(dest.layout.size.bytes() % 16, 0);
139    let len = dest.layout.size.bytes() / 16;
140
141    let u128_array_layout = ecx.layout_of(Ty::new_array(ecx.tcx.tcx, ecx.tcx.types.u128, len))?;
142
143    let state = state.transmute(u128_array_layout, ecx)?;
144    let key = key.transmute(u128_array_layout, ecx)?;
145    let dest = dest.transmute(u128_array_layout, ecx)?;
146
147    for i in 0..len {
148        let state = ecx.read_scalar(&ecx.project_index(&state, i)?)?.to_u128()?;
149        let key = ecx.read_scalar(&ecx.project_index(&key, i)?)?.to_u128()?;
150        let dest = ecx.project_index(&dest, i)?;
151
152        let res = f(state, key);
153
154        ecx.write_scalar(Scalar::from_u128(res), &dest)?;
155    }
156
157    interp_ok(())
158}