zig: using pike for X25519 ECDH -> Blake2b 512-bit KDF -> AEAD

main.zig

const std = @import("std");
const pike = @import("pike/pike.zig");

const os = std.os;
const net = std.net;
const mem = std.mem;
const log = std.log;

const crypto = std.crypto;
const aead = crypto.aead;
const dh = crypto.dh;

const hash = crypto.hash;
const blake2 = hash.blake2;

pub fn main() !void {
    try pike.init();
    defer pike.deinit();

    const notifier = try pike.Notifier.init();
    defer notifier.deinit();

    var stopped = false;

    var server_frame = async doServer(&notifier, try net.Address.parseIp("0.0.0.0", 9000));
    var client_frame = async doClient(&notifier, try net.Address.parseIp("127.0.0.1", 9000), &stopped);

    while (!stopped) {
        try notifier.poll(10_000_000);
    }

    try nosuspend await client_frame;
    try nosuspend await server_frame;
}

pub fn doServer(notifier: *const pike.Notifier, address: net.Address) callconv(.Async) !void {
    var socket = try pike.Socket.init(address.any.family, os.SOCK_STREAM, os.IPPROTO_TCP, 0);
    defer socket.deinit();

    try socket.registerTo(notifier);
    try socket.bind(address);
    try socket.listen(128);

    var conn = try socket.accept();
    try conn.socket.registerTo(notifier);

    // Handshake (X25519 ECDH -> Blake2b 512-bit KDF)

    var ephemeral_peer_public_key: [dh.X25519.public_length]u8 = undefined;
    if ((try conn.socket.read(&ephemeral_peer_public_key)) != @sizeOf(@TypeOf(ephemeral_peer_public_key))) {
        return error.ShortRead;
    }

    var ephemeral_keypair = try dh.X25519.KeyPair.create(null);
    var ephemeral_shared_key = try dh.X25519.scalarmult(ephemeral_keypair.secret_key, ephemeral_peer_public_key);

    if ((try conn.socket.write(&ephemeral_keypair.public_key)) != @sizeOf(@TypeOf(ephemeral_keypair.public_key))) {
        return error.ShortWrite;
    }

    var key_storage: [blake2.Blake2b512.digest_length]u8 = undefined;
    var key_hasher = blake2.Blake2b512.init(.{});

    key_hasher.update(&ephemeral_shared_key);
    mem.secureZero(u8, &ephemeral_shared_key);

    key_hasher.update(&ephemeral_peer_public_key);
    mem.secureZero(u8, &ephemeral_peer_public_key);

    key_hasher.update(&ephemeral_keypair.public_key);
    mem.secureZero(dh.X25519.KeyPair, @as(*[1]dh.X25519.KeyPair, &ephemeral_keypair));

    key_hasher.final(&key_storage);

    var client_to_server_session_key = key_storage[0 .. @sizeOf(@TypeOf(key_storage)) / 2].*;
    var server_to_client_session_key = key_storage[@sizeOf(@TypeOf(key_storage)) / 2 ..].*;

    log.debug("[server] A -> B (encryption key): {x}", .{server_to_client_session_key});
    log.debug("[server] B -> A (decryption key): {x}", .{client_to_server_session_key});
}

pub fn doClient(notifier: *const pike.Notifier, address: net.Address, stopped: *bool) callconv(.Async) !void {
    defer stopped.* = true;

    var socket = try pike.Socket.init(address.any.family, os.SOCK_STREAM, os.IPPROTO_TCP, 0);
    defer socket.deinit();

    try socket.registerTo(notifier);
    try socket.connect(address);

    // Handshake (X25519 ECDH -> Blake2b 512-bit KDF)

    var ephemeral_keypair = try dh.X25519.KeyPair.create(null);
    if ((try socket.write(&ephemeral_keypair.public_key)) != @sizeOf(@TypeOf(ephemeral_keypair.public_key))) {
        return error.ShortWrite;
    }

    var ephemeral_peer_public_key: [dh.X25519.public_length]u8 = undefined;
    if ((try socket.read(&ephemeral_peer_public_key)) != @sizeOf(@TypeOf(ephemeral_peer_public_key))) {
        return error.ShortRead;
    }

    var ephemeral_shared_key = try dh.X25519.scalarmult(ephemeral_keypair.secret_key, ephemeral_peer_public_key);

    var key_storage: [blake2.Blake2b512.digest_length]u8 = undefined;
    var key_hasher = blake2.Blake2b512.init(.{});

    key_hasher.update(&ephemeral_shared_key);
    mem.secureZero(u8, &ephemeral_shared_key);

    key_hasher.update(&ephemeral_keypair.public_key);
    mem.secureZero(dh.X25519.KeyPair, @as(*[1]dh.X25519.KeyPair, &ephemeral_keypair));

    key_hasher.update(&ephemeral_peer_public_key);
    mem.secureZero(u8, &ephemeral_peer_public_key);

    key_hasher.final(&key_storage);

    var client_to_server_session_key = key_storage[0 .. @sizeOf(@TypeOf(key_storage)) / 2].*;
    var server_to_client_session_key = key_storage[@sizeOf(@TypeOf(key_storage)) / 2 ..].*;

    log.debug("[client] A -> B (encryption key): {x}", .{client_to_server_session_key});
    log.debug("[client] B -> A (decryption key): {x}", .{server_to_client_session_key});

    var client_to_server_nonce: u192 = 0;
    var server_to_client_nonce: u192 = 0;
}