//! Shared-secret MACs for permissioned-space repos (Spaces alpha 2026-10-22).
//!
//! A syncer sends `Atproto-Space-Want-Mac` with an ephemeral P-256 key share.
//! The repo host derives a shared secret against the author's `#atproto_kem`
//! key (DHKEM(P-256, HKDF-SHA256), RFC 9180 §4.1) and returns an HMAC over the
//! commit or record in `Atproto-Space-Mac`. A leaked response proves nothing to
//! a third party, since the syncer could have computed the same MAC.
//! Mirrors `@atproto/space` `kem.ts` and `mac.ts`.

const std = @import("std");
const zat = @import("zat");

const P256 = std.crypto.ecc.P256;
const HkdfSha256 = std.crypto.kdf.hkdf.HkdfSha256;
const HmacSha256 = std.crypto.auth.hmac.sha2.HmacSha256;

pub const want_mac_header = "atproto-space-want-mac";
pub const mac_header = "atproto-space-mac";
pub const commit_version = 1;

pub const MacType = enum {
    commit,
    record,

    pub fn name(self: MacType) []const u8 {
        return switch (self) {
            .commit => "atproto-space-repo",
            .record => "atproto-space-record",
        };
    }

    // Commits share their MAC version with the commit format.
    pub fn version(self: MacType) u8 {
        return switch (self) {
            .commit => commit_version,
            .record => 1,
        };
    }
};

pub const Commit = struct {
    ver: u8 = commit_version,
    repo: []const u8,
    space: []const u8,
    rev: []const u8,
    data: [32]u8,
};

pub const Error = error{
    MissingMacHeader,
    MalformedMacHeader,
    UnsupportedMacType,
    UnsupportedMacVersion,
    MissingKeyShare,
    InvalidKemKey,
};

// ---- KEM ----

/// Repo host side: derive the shared secret for a syncer's `key_share` with the
/// author's KEM secret key.
pub fn decapsulate(author_secret: [32]u8, key_share: []const u8) Error![32]u8 {
    var buf: [512]u8 = undefined;
    var fba = std.heap.FixedBufferAllocator.init(&buf);
    const syncer = try p256KeyPoint(fba.allocator(), key_share);
    const author = P256.basePoint.mul(author_secret, .big) catch return error.InvalidKemKey;
    const shared = syncer.mul(author_secret, .big) catch return error.InvalidKemKey;
    const dh = shared.affineCoordinates().x.toBytes(.big);
    return deriveSharedSecret(&dh, &syncer.toUncompressedSec1(), &author.toUncompressedSec1());
}

/// Syncer side, for tests: a fresh ephemeral key share against an author key.
pub fn encapsulate(io: std.Io, allocator: std.mem.Allocator, author_kem_key: []const u8) !struct { key_share: []const u8, shared_secret: [32]u8 } {
    const author = try p256KeyPoint(allocator, author_kem_key);
    const ephemeral = std.crypto.sign.ecdsa.EcdsaP256Sha256.KeyPair.generate(io);
    const secret = ephemeral.secret_key.toBytes();
    const dh = (author.mul(secret, .big) catch return error.InvalidKemKey).affineCoordinates().x.toBytes(.big);
    const ephemeral_point = ephemeral.public_key.p;
    return .{
        .key_share = try zat.multicodec.formatDidKey(allocator, .p256, &ephemeral_point.toCompressedSec1()),
        .shared_secret = deriveSharedSecret(&dh, &ephemeral_point.toUncompressedSec1(), &author.toUncompressedSec1()),
    };
}

// "HPKE-v1" || "KEM" || u16be(0x0010), the suite ID for DHKEM(P-256, HKDF-SHA256).
const kem_prefix = "HPKE-v1" ++ "KEM" ++ "\x00\x10";

fn deriveSharedSecret(dh: *const [32]u8, syncer_key: *const [65]u8, author_key: *const [65]u8) [32]u8 {
    var extract = HkdfSha256.extractInit("");
    extract.update(kem_prefix ++ "eae_prk");
    extract.update(dh);
    var prk: [HkdfSha256.prk_length]u8 = undefined;
    extract.final(&prk);

    const info = "\x00\x20" ++ kem_prefix ++ "shared_secret";
    var info_buf: [info.len + 65 + 65]u8 = undefined;
    @memcpy(info_buf[0..info.len], info);
    @memcpy(info_buf[info.len..][0..65], syncer_key);
    @memcpy(info_buf[info.len + 65 ..][0..65], author_key);
    var out: [32]u8 = undefined;
    HkdfSha256.expand(&out, &info_buf, prk);
    return out;
}

fn p256KeyPoint(allocator: std.mem.Allocator, did_key: []const u8) Error!P256 {
    const parsed = zat.multicodec.parseDidKey(allocator, did_key) catch return error.InvalidKemKey;
    defer allocator.free(parsed.raw);
    if (parsed.key_type != .p256) return error.InvalidKemKey;
    return P256.fromSec1(parsed.raw) catch error.InvalidKemKey;
}

// ---- MACs ----

/// HMAC-SHA256(secret, "atproto-space-repo-v1" || lp(space) || lp(repo) || lp(rev) || data)
pub fn computeCommitMac(shared_secret: *const [32]u8, commit: Commit) [32]u8 {
    var hmac = HmacSha256.init(shared_secret);
    hmac.update("atproto-space-repo-v1");
    updateLengthPrefixed(&hmac, commit.space);
    updateLengthPrefixed(&hmac, commit.repo);
    updateLengthPrefixed(&hmac, commit.rev);
    hmac.update(&commit.data);
    return hmac.finalResult();
}

pub fn verifyCommitMac(shared_secret: *const [32]u8, commit: Commit, mac: []const u8) bool {
    if (mac.len != 32) return false;
    return std.crypto.timing_safe.eql([32]u8, computeCommitMac(shared_secret, commit), mac[0..32].*);
}

/// HMAC-SHA256(secret, "atproto-space-record-v1" || lp(uri) || cid_bytes)
pub fn computeRecordMac(shared_secret: *const [32]u8, uri: []const u8, cid_bytes: []const u8) [32]u8 {
    var hmac = HmacSha256.init(shared_secret);
    hmac.update("atproto-space-record-v1");
    updateLengthPrefixed(&hmac, uri);
    hmac.update(cid_bytes);
    return hmac.finalResult();
}

fn updateLengthPrefixed(hmac: *HmacSha256, value: []const u8) void {
    var len: [2]u8 = undefined;
    std.mem.writeInt(u16, &len, @intCast(value.len), .big);
    hmac.update(&len);
    hmac.update(value);
}

// ---- Headers ----

/// `type="atproto-space-repo", ver=1, mac=:<padded base64>:`
pub fn formatMacHeader(allocator: std.mem.Allocator, mac_type: MacType, mac: *const [32]u8) ![]const u8 {
    var encoded: [std.base64.standard.Encoder.calcSize(32)]u8 = undefined;
    _ = std.base64.standard.Encoder.encode(&encoded, mac);
    return std.fmt.allocPrint(allocator, "type=\"{s}\", ver={d}, mac=:{s}:", .{ mac_type.name(), mac_type.version(), encoded });
}

/// `type="atproto-space-repo", ver=1, key_share="did:key:..."`
pub fn formatWantMacHeader(allocator: std.mem.Allocator, mac_type: MacType, key_share: []const u8) ![]const u8 {
    return std.fmt.allocPrint(allocator, "type=\"{s}\", ver={d}, key_share=\"{s}\"", .{ mac_type.name(), mac_type.version(), key_share });
}

/// Parse a request header and return its `key_share`.
pub fn parseWantMacHeader(allocator: std.mem.Allocator, value: ?[]const u8, mac_type: MacType) Error![]const u8 {
    const fields = try parseMacDictionary(allocator, value, mac_type);
    return switch (fields.key_share orelse return error.MissingKeyShare) {
        .string => |key_share| key_share,
        else => error.MissingKeyShare,
    };
}

/// Parse a response header and return its MAC.
pub fn parseMacHeader(allocator: std.mem.Allocator, value: ?[]const u8, mac_type: MacType) Error![]const u8 {
    const fields = try parseMacDictionary(allocator, value, mac_type);
    return switch (fields.mac orelse return error.MalformedMacHeader) {
        .bytes => |mac| mac,
        else => error.MalformedMacHeader,
    };
}

const MacFields = struct {
    key_share: ?Item = null,
    mac: ?Item = null,
};

fn parseMacDictionary(allocator: std.mem.Allocator, value: ?[]const u8, mac_type: MacType) Error!MacFields {
    var parser: DictionaryParser = .{ .input = value orelse return error.MissingMacHeader, .allocator = allocator };
    var type_item: ?Item = null;
    var ver_item: ?Item = null;
    var fields: MacFields = .{};
    while (try parser.next()) |member| {
        const slot: *?Item = if (std.mem.eql(u8, member.key, "type"))
            &type_item
        else if (std.mem.eql(u8, member.key, "ver"))
            &ver_item
        else if (std.mem.eql(u8, member.key, "key_share"))
            &fields.key_share
        else if (std.mem.eql(u8, member.key, "mac"))
            &fields.mac
        else
            continue;
        slot.* = member.item;
    }
    const type_ok = if (type_item) |item| item == .string and std.mem.eql(u8, item.string, mac_type.name()) else false;
    if (!type_ok) return error.UnsupportedMacType;
    const ver_ok = if (ver_item) |item| switch (item) {
        .integer => |ver| ver == mac_type.version(),
        .decimal => |ver| ver == @as(f64, @floatFromInt(mac_type.version())),
        else => false,
    } else false;
    if (!ver_ok) return error.UnsupportedMacVersion;
    return fields;
}

const Item = union(enum) {
    integer: i64,
    decimal: f64,
    string: []const u8,
    token: []const u8,
    bytes: []const u8,
    boolean: bool,
    inner_list,
};

/// RFC 9651 dictionary parsing, sufficient to read the bare items these headers
/// carry. Parameters are parsed and dropped; a later duplicate key wins.
const DictionaryParser = struct {
    input: []const u8,
    allocator: std.mem.Allocator,
    pos: usize = 0,
    started: bool = false,

    fn next(self: *DictionaryParser) Error!?struct { key: []const u8, item: Item } {
        if (!self.started) {
            self.started = true;
            self.skipSp();
            if (self.pos == self.input.len) return error.MalformedMacHeader;
        } else {
            self.skipOws();
            if (self.pos == self.input.len) return null;
            if (self.input[self.pos] != ',') return error.MalformedMacHeader;
            self.pos += 1;
            self.skipOws();
            if (self.pos == self.input.len) return error.MalformedMacHeader;
        }
        const key = try self.parseKey();
        var item: Item = .{ .boolean = true };
        if (self.peek() == '=') {
            self.pos += 1;
            item = if (self.peek() == '(') try self.parseInnerList() else try self.parseBareItem();
        }
        try self.parseParameters();
        return .{ .key = key, .item = item };
    }

    fn peek(self: *const DictionaryParser) ?u8 {
        return if (self.pos < self.input.len) self.input[self.pos] else null;
    }

    fn skipSp(self: *DictionaryParser) void {
        while (self.peek() == ' ') self.pos += 1;
    }

    fn skipOws(self: *DictionaryParser) void {
        while (self.peek()) |c| {
            if (c != ' ' and c != '\t') break;
            self.pos += 1;
        }
    }

    fn parseKey(self: *DictionaryParser) Error![]const u8 {
        const start = self.pos;
        const first = self.peek() orelse return error.MalformedMacHeader;
        if (!std.ascii.isLower(first) and first != '*') return error.MalformedMacHeader;
        self.pos += 1;
        while (self.peek()) |c| {
            if (!(std.ascii.isLower(c) or std.ascii.isDigit(c) or c == '_' or c == '-' or c == '.' or c == '*')) break;
            self.pos += 1;
        }
        return self.input[start..self.pos];
    }

    fn parseParameters(self: *DictionaryParser) Error!void {
        while (self.peek() == ';') {
            self.pos += 1;
            self.skipSp();
            _ = try self.parseKey();
            if (self.peek() == '=') {
                self.pos += 1;
                _ = try self.parseBareItem();
            }
        }
    }

    fn parseInnerList(self: *DictionaryParser) Error!Item {
        self.pos += 1;
        while (true) {
            self.skipSp();
            if (self.peek() == ')') {
                self.pos += 1;
                try self.parseParameters();
                return .inner_list;
            }
            _ = try self.parseBareItem();
            try self.parseParameters();
            const c = self.peek() orelse return error.MalformedMacHeader;
            if (c != ' ' and c != ')') return error.MalformedMacHeader;
        }
    }

    fn parseBareItem(self: *DictionaryParser) Error!Item {
        const c = self.peek() orelse return error.MalformedMacHeader;
        return switch (c) {
            '-', '0'...'9' => self.parseNumber(),
            '"' => self.parseString(),
            ':' => self.parseBytes(),
            '?' => self.parseBoolean(),
            '*', 'a'...'z', 'A'...'Z' => self.parseToken(),
            else => error.MalformedMacHeader,
        };
    }

    fn parseNumber(self: *DictionaryParser) Error!Item {
        const start = self.pos;
        if (self.peek() == '-') self.pos += 1;
        const digits_start = self.pos;
        var dot: ?usize = null;
        while (self.peek()) |c| {
            if (std.ascii.isDigit(c)) {
                self.pos += 1;
            } else if (c == '.' and dot == null) {
                if (self.pos - digits_start > 12) return error.MalformedMacHeader;
                dot = self.pos;
                self.pos += 1;
            } else break;
        }
        const text = self.input[start..self.pos];
        if (self.pos == digits_start) return error.MalformedMacHeader;
        if (dot) |dot_pos| {
            const fraction = self.pos - dot_pos - 1;
            if (fraction == 0 or fraction > 3) return error.MalformedMacHeader;
            return .{ .decimal = std.fmt.parseFloat(f64, text) catch return error.MalformedMacHeader };
        }
        if (self.pos - digits_start > 15) return error.MalformedMacHeader;
        return .{ .integer = std.fmt.parseInt(i64, text, 10) catch return error.MalformedMacHeader };
    }

    fn parseString(self: *DictionaryParser) Error!Item {
        self.pos += 1;
        var out: std.ArrayList(u8) = .empty;
        while (self.peek()) |c| {
            self.pos += 1;
            switch (c) {
                '"' => return .{ .string = out.items },
                '\\' => {
                    const escaped = self.peek() orelse return error.MalformedMacHeader;
                    if (escaped != '"' and escaped != '\\') return error.MalformedMacHeader;
                    self.pos += 1;
                    out.append(self.allocator, escaped) catch return error.MalformedMacHeader;
                },
                0x20...0x21, 0x23...0x5b, 0x5d...0x7e => out.append(self.allocator, c) catch return error.MalformedMacHeader,
                else => return error.MalformedMacHeader,
            }
        }
        return error.MalformedMacHeader;
    }

    fn parseBytes(self: *DictionaryParser) Error!Item {
        self.pos += 1;
        const end = std.mem.indexOfScalarPos(u8, self.input, self.pos, ':') orelse return error.MalformedMacHeader;
        const encoded = self.input[self.pos..end];
        self.pos = end + 1;
        for (encoded) |c| {
            if (!(std.ascii.isAlphanumeric(c) or c == '+' or c == '/' or c == '=')) return error.MalformedMacHeader;
        }
        const decoder = std.base64.standard.Decoder;
        const len = decoder.calcSizeForSlice(encoded) catch return error.MalformedMacHeader;
        const decoded = self.allocator.alloc(u8, len) catch return error.MalformedMacHeader;
        decoder.decode(decoded, encoded) catch return error.MalformedMacHeader;
        return .{ .bytes = decoded };
    }

    fn parseBoolean(self: *DictionaryParser) Error!Item {
        self.pos += 1;
        const c = self.peek() orelse return error.MalformedMacHeader;
        self.pos += 1;
        return switch (c) {
            '0' => .{ .boolean = false },
            '1' => .{ .boolean = true },
            else => error.MalformedMacHeader,
        };
    }

    fn parseToken(self: *DictionaryParser) Error!Item {
        const start = self.pos;
        self.pos += 1;
        while (self.peek()) |c| {
            const tchar = std.ascii.isAlphanumeric(c) or std.mem.indexOfScalar(u8, "!#$%&'*+-.^_`|~:/", c) != null;
            if (!tchar) break;
            self.pos += 1;
        }
        return .{ .token = self.input[start..self.pos] };
    }
};

fn hexBytes(comptime hex: []const u8) [hex.len / 2]u8 {
    var out: [hex.len / 2]u8 = undefined;
    _ = std.fmt.hexToBytes(&out, hex) catch unreachable;
    return out;
}

test "KEM matches the RFC 9180 A.3.1 vector" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();
    // pkEm is the syncer's key share and skRm the author's KEM key.
    const pk_em = hexBytes("04a92719c6195d5085104f469a8b9814d5838ff72b60501e2c4466e5e67b325ac98536d7b61a1af4b78e5b7f951c0900be863c403ce65c9bfcb9382657222d18c4");
    const sk_rm = hexBytes("f3ce7fdae57e1a310d87f1ebbde6f328be0a99cdbcadf4d6589cf29de4b8ffd2");
    const point = try P256.fromSec1(&pk_em);
    const key_share = try zat.multicodec.formatDidKey(allocator, .p256, &point.toCompressedSec1());
    const shared = try decapsulate(sk_rm, key_share);
    try std.testing.expectEqualStrings(
        "c0d26aeab536609a572b07695d933b589dcf363ff9d93c93adea537aeabb8cb8",
        &std.fmt.bytesToHex(shared, .lower),
    );
}

test "KEM matches the reference implementation and round trips" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();
    // @atproto/space 0.0.0-spaces-alpha-20261022141228 decapsulate(P256Keypair.import(0x31 * 32), did of 0x32 * 32).
    const author_secret: [32]u8 = @splat(0x31);
    const shared = try decapsulate(author_secret, "did:key:zDnaemu5viSrXFCjxg3KigUcZ57xLDhgQBHMCJTZtatnsZjM6");
    try std.testing.expectEqualStrings(
        "629cf3ba48dc1db7eeff7e9af7d1fe838f62d8152e59a142c7f46791f361bcc6",
        &std.fmt.bytesToHex(shared, .lower),
    );

    const author = try zat.Keypair.fromSecretKey(.p256, author_secret);
    const encapsulated = try encapsulate(std.testing.io, allocator, try author.did(allocator));
    try std.testing.expectEqualSlices(u8, &encapsulated.shared_secret, &try decapsulate(author_secret, encapsulated.key_share));
    const other: [32]u8 = @splat(0x33);
    try std.testing.expect(!std.mem.eql(u8, &encapsulated.shared_secret, &try decapsulate(other, encapsulated.key_share)));
}

test "KEM rejects keys that are not P-256" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const k256 = try zat.Keypair.fromSecretKey(.secp256k1, @splat(0x31));
    try std.testing.expectError(error.InvalidKemKey, decapsulate(@splat(0x31), try k256.did(arena.allocator())));
    try std.testing.expectError(error.InvalidKemKey, decapsulate(@splat(0x31), "did:key:nope"));
}

test "commit and record MACs match the reference implementation" {
    const secret: [32]u8 = @splat(0x42);
    // Commit for the space-reference-repo fixture, MAC'd by @atproto/space computeCommitMac.
    const commit: Commit = .{
        .repo = "did:plc:refuser",
        .space = "at://did:plc:refauthority/space/com.example.forum/general",
        .rev = "3mstarrefrevaa",
        .data = hexBytes("869574498510cb637faf8de78a9557ea1223e0263ef37d70a07deec680e55c50"),
    };
    const mac = computeCommitMac(&secret, commit);
    try std.testing.expectEqualStrings("9f90b6801163a1b25654eab6c1781938a55a96f4355b65cc749dd6a52f3537ee", &std.fmt.bytesToHex(mac, .lower));
    try std.testing.expect(verifyCommitMac(&secret, commit, &mac));
    var tampered = commit;
    tampered.data[0] ^= 1;
    try std.testing.expect(!verifyCommitMac(&secret, tampered, &mac));

    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const cid = try zat.cbor.Cid.fromString(arena.allocator(), "bafyreihtk4vh7vfupnk2ad5q2yceepsqftci5vija3tsrtuuga7shg56eu");
    const record_mac = computeRecordMac(
        &secret,
        "at://did:plc:refauthority/space/com.example.forum/general/did:plc:refuser/com.example.post/aaa",
        cid.raw,
    );
    try std.testing.expectEqualStrings("f775baac9f49eb70ce2836c41d85ecb1e1361bbe05f22d3438ae5d00109f0ad4", &std.fmt.bytesToHex(record_mac, .lower));
}

test "commit MAC length-prefixes each field" {
    const secret: [32]u8 = @splat(7);
    const a: Commit = .{ .repo = "did:example:alice", .space = "s", .rev = "3kbc", .data = @splat(1) };
    const b: Commit = .{ .repo = "did:example:alice3", .space = "s", .rev = "kbc", .data = @splat(1) };
    try std.testing.expect(!std.mem.eql(u8, &computeCommitMac(&secret, a), &computeCommitMac(&secret, b)));
}

test "MAC headers format and parse like the reference implementation" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();
    var mac: [32]u8 = undefined;
    for (&mac, 0..) |*byte, idx| byte.* = @intCast(idx);
    const header = try formatMacHeader(allocator, .commit, &mac);
    try std.testing.expectEqualStrings("type=\"atproto-space-repo\", ver=1, mac=:AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8=:", header);
    try std.testing.expectEqualSlices(u8, &mac, try parseMacHeader(allocator, header, .commit));

    const want = try formatWantMacHeader(allocator, .commit, "did:key:zDnaExample");
    try std.testing.expectEqualStrings("type=\"atproto-space-repo\", ver=1, key_share=\"did:key:zDnaExample\"", want);
    try std.testing.expectEqualStrings("did:key:zDnaExample", try parseWantMacHeader(allocator, want, .commit));
    try std.testing.expectEqualStrings("did:key:zDnaX", try parseWantMacHeader(allocator, "key_share=\"did:key:zDnaX\";p=1,ver=1,   type=\"atproto-space-record\", extra=(1 2)", .record));
}

test "want-mac headers reject missing, malformed, and mismatched values" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const allocator = arena.allocator();
    try std.testing.expectError(error.MissingMacHeader, parseWantMacHeader(allocator, null, .commit));
    try std.testing.expectError(error.MalformedMacHeader, parseWantMacHeader(allocator, "???", .commit));
    try std.testing.expectError(error.MalformedMacHeader, parseWantMacHeader(allocator, "", .commit));
    try std.testing.expectError(error.MalformedMacHeader, parseWantMacHeader(allocator, "type=\"atproto-space-repo\", ver=1,", .commit));
    try std.testing.expectError(error.MissingKeyShare, parseWantMacHeader(allocator, "type=\"atproto-space-repo\", ver=1", .commit));
    try std.testing.expectError(error.MissingKeyShare, parseWantMacHeader(allocator, "type=\"atproto-space-repo\", ver=1, key_share=did", .commit));
    try std.testing.expectError(error.UnsupportedMacType, parseWantMacHeader(allocator, "type=atproto-space-repo, ver=1, key_share=\"x\"", .commit));
    try std.testing.expectError(error.UnsupportedMacType, parseWantMacHeader(allocator, "type=\"atproto-space-repo\", ver=1, key_share=\"x\"", .record));
    try std.testing.expectError(error.UnsupportedMacVersion, parseWantMacHeader(allocator, "type=\"atproto-space-repo\", ver=2, key_share=\"x\"", .commit));
    try std.testing.expectError(error.UnsupportedMacVersion, parseWantMacHeader(allocator, "type=\"atproto-space-repo\", key_share=\"x\"", .commit));
    try std.testing.expectError(error.MalformedMacHeader, parseMacHeader(allocator, "type=\"atproto-space-repo\", ver=1, mac=\"abc\"", .commit));
}