Merge pull request 'std.crypto.aes: expose the inverse MixColumns operation' (#30052) from jedisct1/zig:invmixcolumns into master

Reviewed-on: https://codeberg.org/ziglang/zig/pulls/30052

lib/std/crypto/aes.zig

@@ -108,6 +108,36 @@ test "expand 128-bit key" {
     }
 }
 
+test "invMixColumns" {
+    const key = [_]u8{ 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c };
+    const enc_ctx = Aes128.initEnc(key);
+    const dec_ctx = Aes128.initDec(key);
+
+    for (1..10) |i| {
+        const enc_rk = enc_ctx.key_schedule.round_keys[10 - i];
+        const dec_rk = dec_ctx.key_schedule.round_keys[i];
+        const computed = enc_rk.invMixColumns();
+        try testing.expectEqualSlices(u8, &dec_rk.toBytes(), &computed.toBytes());
+    }
+}
+
+test "BlockVec invMixColumns" {
+    const input = [_]u8{
+        0x5f, 0x57, 0xf7, 0x1d, 0x72, 0xf5, 0xbe, 0xb9, 0x64, 0xbc, 0x3b, 0xf9, 0x15, 0x92, 0x29, 0x1a,
+        0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf, 0x4f, 0x3c,
+    };
+
+    const vec2 = BlockVec(2).fromBytes(&input);
+    const result_vec = vec2.invMixColumns();
+    const result_bytes = result_vec.toBytes();
+
+    for (0..2) |i| {
+        const block = Block.fromBytes(input[i * 16 ..][0..16]);
+        const expected = block.invMixColumns().toBytes();
+        try testing.expectEqualSlices(u8, &expected, result_bytes[i * 16 ..][0..16]);
+    }
+}
+
 test "expand 256-bit key" {
     const key = [_]u8{
         0x60, 0x3d, 0xeb, 0x10,

lib/std/crypto/aes/aesni.zig

@@ -96,6 +96,17 @@ pub const Block = struct {
         return Block{ .repr = block1.repr | block2.repr };
     }
 
+    /// Apply the inverse MixColumns operation to a block.
+    pub fn invMixColumns(block: Block) Block {
+        return Block{
+            .repr = asm (
+                \\ vaesimc %[in], %[out]
+                : [out] "=x" (-> Repr),
+                : [in] "x" (block.repr),
+            ),
+        };
+    }
+
     /// Perform operations on multiple blocks in parallel.
     pub const parallel = struct {
         const cpu = std.Target.x86.cpu;
@@ -308,6 +319,17 @@ pub fn BlockVec(comptime blocks_count: comptime_int) type {
             }
             return out;
         }
+
+        /// Apply the inverse MixColumns operation to each block in the vector.
+        pub fn invMixColumns(block_vec: Self) Self {
+            var out_bytes: [blocks_count * 16]u8 = undefined;
+            const in_bytes = block_vec.toBytes();
+            inline for (0..blocks_count) |i| {
+                const block = Block.fromBytes(in_bytes[i * 16 ..][0..16]);
+                out_bytes[i * 16 ..][0..16].* = block.invMixColumns().toBytes();
+            }
+            return fromBytes(&out_bytes);
+        }
     };
 }
 

lib/std/crypto/aes/armcrypto.zig

@@ -99,6 +99,17 @@ pub const Block = struct {
         return Block{ .repr = block1.repr | block2.repr };
     }
 
+    /// Apply the inverse MixColumns operation to a block.
+    pub fn invMixColumns(block: Block) Block {
+        return Block{
+            .repr = asm (
+                \\ aesimc %[out].16b, %[in].16b
+                : [out] "=x" (-> Repr),
+                : [in] "x" (block.repr),
+            ),
+        };
+    }
+
     /// Perform operations on multiple blocks in parallel.
     pub const parallel = struct {
         /// The recommended number of AES encryption/decryption to perform in parallel for the chosen implementation.
@@ -275,6 +286,15 @@ pub fn BlockVec(comptime blocks_count: comptime_int) type {
             }
             return out;
         }
+
+        /// Apply the inverse MixColumns operation to each block in the vector.
+        pub fn invMixColumns(block_vec: Self) Self {
+            var out: Self = undefined;
+            inline for (0..native_words) |i| {
+                out.repr[i] = block_vec.repr[i].invMixColumns();
+            }
+            return out;
+        }
     };
 }
 

lib/std/crypto/aes/soft.zig

@@ -265,6 +265,26 @@ pub const Block = struct {
         return Block{ .repr = x };
     }
 
+    /// Apply the inverse MixColumns operation to a block.
+    pub fn invMixColumns(block: Block) Block {
+        var out: Repr = undefined;
+        inline for (0..4) |i| {
+            const col = block.repr[i];
+            const b0: u8 = @truncate(col);
+            const b1: u8 = @truncate(col >> 8);
+            const b2: u8 = @truncate(col >> 16);
+            const b3: u8 = @truncate(col >> 24);
+
+            const r0 = mul(0x0e, b0) ^ mul(0x0b, b1) ^ mul(0x0d, b2) ^ mul(0x09, b3);
+            const r1 = mul(0x09, b0) ^ mul(0x0e, b1) ^ mul(0x0b, b2) ^ mul(0x0d, b3);
+            const r2 = mul(0x0d, b0) ^ mul(0x09, b1) ^ mul(0x0e, b2) ^ mul(0x0b, b3);
+            const r3 = mul(0x0b, b0) ^ mul(0x0d, b1) ^ mul(0x09, b2) ^ mul(0x0e, b3);
+
+            out[i] = @as(u32, r0) | (@as(u32, r1) << 8) | (@as(u32, r2) << 16) | (@as(u32, r3) << 24);
+        }
+        return Block{ .repr = out };
+    }
+
     /// Perform operations on multiple blocks in parallel.
     pub const parallel = struct {
         /// The recommended number of AES encryption/decryption to perform in parallel for the chosen implementation.
@@ -441,6 +461,15 @@ pub fn BlockVec(comptime blocks_count: comptime_int) type {
             }
             return out;
         }
+
+        /// Apply the inverse MixColumns operation to each block in the vector.
+        pub fn invMixColumns(block_vec: Self) Self {
+            var out: Self = undefined;
+            for (0..native_words) |i| {
+                out.repr[i] = block_vec.repr[i].invMixColumns();
+            }
+            return out;
+        }
     };
 }