tiny.simd.crypto

1const std = @import("std");23pub 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}78pub 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}1213pub 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}1718pub 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}2223pub fn aesInvMixColumns(comptime D: type, state: D.Vector) D.Vector {24    validateAes(D);25    return mixColumns(D, state, true);26}2728pub 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}4243pub fn clMulLower(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {44    return clMul(D, a, b, 0);45}4647pub fn clMulUpper(comptime D: type, a: D.Vector, b: D.Vector) D.Vector {48    return clMul(D, a, b, 1);49}5051fn addRoundKey(comptime D: type, state: D.Vector, key: D.Vector) D.Vector {52    return state ^ key;53}5455fn 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) & 363                else64                    (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                else69                    sbox(input[source]);70            }71        }72    }73    return result;74}7576fn 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}99100fn 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}105106fn inverseSbox(value: u8) u8 {107    return gfInverse(rotateByte(value, 1) ^ rotateByte(value, 3) ^108        rotateByte(value, 6) ^ 0x05);109}110111fn 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}120121fn 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}134135fn rotateByte(value: u8, comptime amount: u3) u8 {136    return std.math.rotl(u8, value, amount);137}138139fn 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}159160fn 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}168169fn 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}174175test "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}193194test "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}213214test "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}227228test "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}238239test "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}

1pub const crypto = @import("crypto.zig");