lib/simd/src/crypto.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub fn aesRound(comptime D: type, state: D.Vector, round_key: D.Vector) D.Vector {
4 validateAes(D);
5 return addRoundKey(D, mixColumns(D, shiftSubBytes(D, state, false), false), round_key);
6 }
7
8 pub fn aesLastRound(comptime D: type, state: D.Vector, round_key: D.Vector) D.Vector {
9 validateAes(D);
10 return addRoundKey(D, shiftSubBytes(D, state, false), round_key);
11 }
12
13 pub fn aesRoundInv(comptime D: type, state: D.Vector, round_key: D.Vector) D.Vector {
14 validateAes(D);
15 return addRoundKey(D, mixColumns(D, shiftSubBytes(D, state, true), true), round_key);
16 }
17
18 pub fn aesLastRoundInv(comptime D: type, state: D.Vector, round_key: D.Vector) D.Vector {
19 validateAes(D);
20 return addRoundKey(D, shiftSubBytes(D, state, true), round_key);
21 }
22
23 pub fn aesInvMixColumns(comptime D: type, state: D.Vector) D.Vector {
24 validateAes(D);
25 return mixColumns(D, state, true);
26 }
27
28 pub fn aesKeyGenAssist(comptime D: type, comptime rcon: u8, value: D.Vector) D.Vector {
29 validateAes(D);
30 const input: [D.lane_count]u8 = value;
31 var result: [D.lane_count]u8 = undefined;
32 const indices = [16]usize{ 4, 5, 6, 7, 5, 6, 7, 4, 12, 13, 14, 15, 13, 14, 15, 12 };
33 inline for (0..D.lane_count / 16) |block| {
34 inline for (0..16) |index| {
35 const source = block * 16 + indices[index];
36 const constant: u8 = if (index == 4 or index == 12) rcon else 0;
37 result[block * 16 + index] = sbox(input[source]) ^ constant;
38 }
39 }
40 return result;
41 }
42
43 pub fn clMulLower(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {
44 return clMul(D, a, b, 0);
45 }
46
47 pub fn clMulUpper(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {
48 return clMul(D, a, b, 1);
49 }
50
51 fn addRoundKey(comptime D: type, state: D.Vector, key: D.Vector) D.Vector {
52 return state ^ key;
53 }
54
55 fn shiftSubBytes(comptime D: type, state: D.Vector, comptime inverse: bool) D.Vector {
56 const input: [D.lane_count]u8 = state;
57 var result: [D.lane_count]u8 = undefined;
58 inline for (0..D.lane_count / 16) |block| {
59 inline for (0..4) |column| {
60 inline for (0..4) |row| {
61 const source_column = if (inverse)
62 (column + 4 - row) & 3
63 else
64 (column + row) & 3;
65 const source = block * 16 + source_column * 4 + row;
66 result[block * 16 + column * 4 + row] = if (inverse)
67 inverseSbox(input[source])
68 else
69 sbox(input[source]);
70 }
71 }
72 }
73 return result;
74 }
75
76 fn mixColumns(comptime D: type, state: D.Vector, comptime inverse: bool) D.Vector {
77 const input: [D.lane_count]u8 = state;
78 var result: [D.lane_count]u8 = undefined;
79 inline for (0..D.lane_count / 4) |column| {
80 const offset = column * 4;
81 const a0 = input[offset];
82 const a1 = input[offset + 1];
83 const a2 = input[offset + 2];
84 const a3 = input[offset + 3];
85 if (inverse) {
86 result[offset] = gfMul(a0, 14) ^ gfMul(a1, 11) ^ gfMul(a2, 13) ^ gfMul(a3, 9);
87 result[offset + 1] = gfMul(a0, 9) ^ gfMul(a1, 14) ^ gfMul(a2, 11) ^ gfMul(a3, 13);
88 result[offset + 2] = gfMul(a0, 13) ^ gfMul(a1, 9) ^ gfMul(a2, 14) ^ gfMul(a3, 11);
89 result[offset + 3] = gfMul(a0, 11) ^ gfMul(a1, 13) ^ gfMul(a2, 9) ^ gfMul(a3, 14);
90 } else {
91 result[offset] = gfMul(a0, 2) ^ gfMul(a1, 3) ^ a2 ^ a3;
92 result[offset + 1] = a0 ^ gfMul(a1, 2) ^ gfMul(a2, 3) ^ a3;
93 result[offset + 2] = a0 ^ a1 ^ gfMul(a2, 2) ^ gfMul(a3, 3);
94 result[offset + 3] = gfMul(a0, 3) ^ a1 ^ a2 ^ gfMul(a3, 2);
95 }
96 }
97 return result;
98 }
99
100 fn sbox(value: u8) u8 {
101 const inverse = gfInverse(value);
102 return inverse ^ rotateByte(inverse, 1) ^ rotateByte(inverse, 2) ^
103 rotateByte(inverse, 3) ^ rotateByte(inverse, 4) ^ 0x63;
104 }
105
106 fn inverseSbox(value: u8) u8 {
107 return gfInverse(rotateByte(value, 1) ^ rotateByte(value, 3) ^
108 rotateByte(value, 6) ^ 0x05);
109 }
110
111 fn gfInverse(value: u8) u8 {
112 var result: u8 = 1;
113 var base = value;
114 inline for (0..8) |bit| {
115 if ((254 >> bit) & 1 != 0) result = gfMul(result, base);
116 base = gfMul(base, base);
117 }
118 return result;
119 }
120
121 fn gfMul(a: u8, b: u8) u8 {
122 var product: u8 = 0;
123 var multiplicand = a;
124 var multiplier = b;
125 inline for (0..8) |_| {
126 const selected = 0 -% (multiplier & 1);
127 product ^= multiplicand & selected;
128 const high = multiplicand >> 7;
129 multiplicand = (multiplicand << 1) ^ (0x1b & (0 -% high));
130 multiplier >>= 1;
131 }
132 return product;
133 }
134
135 fn rotateByte(value: u8, comptime amount: u3) u8 {
136 return std.math.rotl(u8, value, amount);
137 }
138
139 fn clMul(
140 comptime D: type,
141 a: D.Vector,
142 b: D.Vector,
143 comptime selected_lane: usize,
144 ) D.Vector {
145 if (comptime D.Lane != u64 or D.lane_count < 2 or D.lane_count & 1 != 0) {
146 @compileError("carryless multiplication requires an even number of u64 lanes");
147 }
148 const a_lanes: [D.lane_count]u64 = a;
149 const b_lanes: [D.lane_count]u64 = b;
150 var result: [D.lane_count]u64 = undefined;
151 inline for (0..D.lane_count / 2) |block| {
152 const index = block * 2 + selected_lane;
153 const product = carrylessProduct(a_lanes[index], b_lanes[index]);
154 result[block * 2] = @truncate(product);
155 result[block * 2 + 1] = @truncate(product >> 64);
156 }
157 return result;
158 }
159
160 fn carrylessProduct(a: u64, b: u64) u128 {
161 var product: u128 = 0;
162 inline for (0..64) |bit| {
163 const selected = 0 -% @as(u128, (b >> bit) & 1);
164 product ^= (@as(u128, a) << bit) & selected;
165 }
166 return product;
167 }
168
169 fn validateAes(comptime D: type) void {
170 if (comptime D.Lane != u8 or D.lane_count < 16 or D.lane_count % 16 != 0) {
171 @compileError("AES operations require whole 128-bit blocks of u8 lanes");
172 }
173 }
174
175 test "Highway AES encryption rounds match the NIST core fixture" {
176 const simd = @import("root.zig");
177 const D = simd.FixedTag(u8, 16);
178 const state: D.Vector = .{
179 0x40, 0xbf, 0xab, 0xf4, 0x06, 0xee, 0x4d, 0x30,
180 0x42, 0xca, 0x6b, 0x99, 0x7a, 0x5c, 0x58, 0x16,
181 };
182 const mixed: D.Vector = .{
183 0x52, 0x9f, 0x16, 0xc2, 0x97, 0x86, 0x15, 0xca,
184 0xe0, 0x1a, 0xae, 0x54, 0xba, 0x1a, 0x26, 0x59,
185 };
186 const shifted: D.Vector = .{
187 0x09, 0x28, 0x7f, 0x47, 0x6f, 0x74, 0x6a, 0xbf,
188 0x2c, 0x4a, 0x62, 0x04, 0xda, 0x08, 0xe3, 0xee,
189 };
190 try std.testing.expect(@reduce(.And, aesRound(D, state, @splat(0)) == mixed));
191 try std.testing.expect(@reduce(.And, aesLastRound(D, state, @splat(0)) == shifted));
192 }
193
194 test "Highway AES inverse rounds and inverse mixing match FIPS fixtures" {
195 const simd = @import("root.zig");
196 const D = simd.FixedTag(u8, 16);
197 const state: D.Vector = .{
198 0x7a, 0xd5, 0xfd, 0xa7, 0x89, 0xef, 0x4e, 0x27,
199 0x2b, 0xca, 0x10, 0x0b, 0x3d, 0x9f, 0xf5, 0x9f,
200 };
201 const shifted: D.Vector = .{
202 0xbd, 0x6e, 0x7c, 0x3d, 0xf2, 0xb5, 0x77, 0x9e,
203 0x0b, 0x61, 0x21, 0x6e, 0x8b, 0x10, 0xb6, 0x89,
204 };
205 const mixed: D.Vector = .{
206 0x47, 0x73, 0xb9, 0x1f, 0xf7, 0x2f, 0x35, 0x43,
207 0x61, 0xcb, 0x01, 0x8e, 0xa1, 0xe6, 0xcf, 0x2c,
208 };
209 try std.testing.expect(@reduce(.And, aesLastRoundInv(D, state, @splat(0)) == shifted));
210 try std.testing.expect(@reduce(.And, aesRoundInv(D, state, @splat(0)) == mixed));
211 try std.testing.expect(@reduce(.And, aesInvMixColumns(D, shifted) == mixed));
212 }
213
214 test "Highway AES key generation assist matches x86 semantics" {
215 const simd = @import("root.zig");
216 const D = simd.FixedTag(u8, 16);
217 const value: D.Vector = .{
218 0x27, 0xcf, 0x73, 0xc3, 0x27, 0xcf, 0x73, 0xc3,
219 0x74, 0x01, 0x90, 0x5a, 0x74, 0x01, 0x90, 0x5a,
220 };
221 const expected: D.Vector = .{
222 0xcc, 0x8a, 0x8f, 0x2e, 0xaa, 0x8f, 0x2e, 0xcc,
223 0x92, 0x7c, 0x60, 0xbe, 0x5c, 0x60, 0xbe, 0x92,
224 };
225 try std.testing.expect(@reduce(.And, aesKeyGenAssist(D, 0x20, value) == expected));
226 }
227
228 test "Highway AES S-box permutations cover all byte values" {
229 var seen: [256]bool = @splat(false);
230 for (0..256) |index| {
231 const value: u8 = @intCast(index);
232 const substituted = sbox(value);
233 try std.testing.expectEqual(value, inverseSbox(substituted));
234 seen[substituted] = true;
235 }
236 for (seen) |value| try std.testing.expect(value);
237 }
238
239 test "Highway carryless multiplication selects lower and upper block lanes" {
240 const simd = @import("root.zig");
241 const D = simd.FixedTag(u64, 4);
242 const a: D.Vector = .{ 2, 4, 5, 7 };
243 const b: D.Vector = .{ 3, 9, 11, 13 };
244 try std.testing.expect(@reduce(.And, clMulLower(D, a, b) ==
245 @as(D.Vector, .{ 6, 0, 0x27, 0 })));
246 try std.testing.expect(@reduce(.And, clMulUpper(D, a, b) ==
247 @as(D.Vector, .{ 0x24, 0, 0x23, 0 })));
248 }