noise.zig

const std = @import("std");
const crypto = std.crypto;
const X25519 = crypto.dh.X25519;
const ChaCha20Poly1305 = crypto.aead.chacha_poly.ChaCha20Poly1305;
const Blake2s256 = crypto.hash.blake2.Blake2s256;
const Blake2s128 = crypto.hash.blake2.Blake2s128;
const Hmac = crypto.auth.hmac.Hmac(Blake2s256);

const construction = "Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s";
pub const construction_hash = hash(construction);

pub const identifier = "WireGuard v1 zx2c4 [email protected]".*;

pub const label_mac1 = "mac1----";
pub const label_cookie = "cookie--";

pub const public_length = X25519.public_length;
pub const private_length = X25519.secret_length;
pub const PublicKey = [public_length]u8;
pub const PrivateKey = [private_length]u8;
pub const KeyPair = X25519.KeyPair;
pub const dh = X25519.scalarmult;

pub const digest_length = Blake2s256.digest_length;
pub const Digest = [digest_length]u8;

pub const mac_length = Blake2s128.digest_length;
pub const Mac = [mac_length]u8;
pub const Nonce = u64;

/// "Key" just means a data-type that gets passed to either mac, hmac, or kdf3. In all cases
/// it is 32-bytes long.
pub const Key = [key_length]u8;
pub const key_length = 32;
pub const tag_length = ChaCha20Poly1305.tag_length;
pub const Tag = [tag_length]u8;

pub fn Aead(comptime T: type) type {
    switch (@typeInfo(T)) {
        .array => |ti| {
            if (ti.child != u8) {
                @compileError("Aead(...) can only wrap [_]u8 or void");
            }
            var newtype = ti;
            newtype.len += tag_length;
            return @Type(.{ .array = newtype });
        },
        .void => {
            return Tag;
        },
        else => @compileError("Aead(...) can only wrap [_]u8 or void"),
    }
}

/// HKDF, but named based on the WireGuard whitepaper
pub fn kdf3(
    key: Key,
    input: []const u8,
) [3]Digest {
    const prk = hmac(key, input);
    var t: [3]Digest = undefined;
    t[0] = hmac(prk, &.{0x1});
    t[1] = hmac(prk, &(t[0] ++ .{0x2}));
    t[2] = hmac(prk, &(t[1] ++ .{0x3}));
    return t;
}

pub fn hmac(key: Key, input: []const u8) Digest {
    var output: Digest = undefined;
    var h = Hmac.init(&key);
    h.update(input);
    h.final(&output);
    return output;
}

pub fn hash(input: []const u8) Digest {
    @setEvalBranchQuota(20000);
    var output: Digest = undefined;
    Blake2s256.hash(input, &output, .{});
    return output;
}

pub fn mac(key: Key, input: []const u8) Mac {
    var output: Mac = undefined;
    Blake2s128.hash(input, &output, .{ .key = &key });
    return output;
}

pub fn aeadSeal(
    cipher: []u8,
    key: Key,
    counter: Nonce,
    plain: []const u8,
    auth: []const u8,
) void {
    var nonce = std.mem.zeroes([12]u8);
    std.mem.writeInt(Nonce, nonce[4..], counter, .little);
    ChaCha20Poly1305.encrypt(
        cipher[0..plain.len],
        @ptrCast(cipher[plain.len..]),
        plain,
        auth,
        nonce,
        key,
    );
}

pub fn aeadUnseal(
    plain: []u8,
    key: Key,
    counter: Nonce,
    cipher: []const u8,
    auth: []const u8,
) error{AuthenticationFailed}!void {
    var nonce = std.mem.zeroes([12]u8);
    std.mem.writeInt(Nonce, nonce[4..], counter, .little);
    var tag: Tag = undefined;
    @memcpy(&tag, cipher[cipher.len - tag_length ..]);
    try ChaCha20Poly1305.decrypt(
        plain,
        cipher[0 .. cipher.len - tag_length],
        tag,
        auth,
        nonce,
        key,
    );
}