ccm.zig

const std = @import("std");
const crypto = std.crypto;
const mem = std.mem;
const debug = std.debug;
const modes = crypto.core.modes;
const AuthenticationError = crypto.errors.AuthenticationError;
const cbc_mac = @import("cbc_mac.zig");

/// CCM (Counter with CBC-MAC) authenticated encryption mode
/// RFC 3610: https://www.rfc-editor.org/rfc/rfc3610
/// NIST SP 800-38C: https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-38C.pdf
///
/// CCM is an AEAD (Authenticated Encryption with Associated Data) mode that combines
/// CTR mode encryption with CBC-MAC authentication.
///
/// Common instances with AES-128 and standard parameters (tag=16, nonce=12)
pub const Aes128Ccm = Ccm(crypto.core.aes.Aes128, 16, 12);
pub const Aes256Ccm = Ccm(crypto.core.aes.Aes256, 16, 12);

/// Generic CCM mode implementation
///
/// Parameters:
/// - BlockCipher: Block cipher type (must have 16-byte blocks)
/// - tag_length: Authentication tag length in bytes (4, 6, 8, 10, 12, 14, or 16)
/// - nonce_length: Nonce length in bytes (7 to 13)
pub fn Ccm(comptime BlockCipher: type, comptime tag_length: usize, comptime nonce_length: usize) type {
    const block_length = BlockCipher.block.block_length;
    debug.assert(block_length == 16); // CCM requires 16-byte blocks

    // Validate tag_length
    comptime {
        if (tag_length < 4 or tag_length > 16 or tag_length % 2 != 0) {
            @compileError("CCM tag_length must be 4, 6, 8, 10, 12, 14, or 16 bytes");
        }
    }

    // Validate nonce_length
    comptime {
        if (nonce_length < 7 or nonce_length > 13) {
            @compileError("CCM nonce_length must be between 7 and 13 bytes");
        }
    }

    const L = 15 - nonce_length; // Counter size in bytes (1 to 8)
    const BlockCipherCtx = @typeInfo(@TypeOf(BlockCipher.initEnc)).@"fn".return_type.?;

    return struct {
        pub const key_length = BlockCipher.key_bits / 8;

        /// Encrypt a message with associated data
        ///
        /// Parameters:
        /// - c: Ciphertext output buffer (must be same length as m)
        /// - tag: Authentication tag output buffer
        /// - m: Plaintext message
        /// - ad: Associated data (authenticated but not encrypted)
        /// - npub: Nonce (public, must be unique for each message with same key)
        /// - key: Encryption key
        pub fn encrypt(c: []u8, tag: *[tag_length]u8, m: []const u8, ad: []const u8, npub: [nonce_length]u8, key: [key_length]u8) void {
            debug.assert(c.len == m.len);

            // Validate message length fits in L bytes
            const max_msg_len: u64 = if (L >= 8) std.math.maxInt(u64) else (@as(u64, 1) << @as(u6, @intCast(L * 8))) - 1;
            debug.assert(m.len <= max_msg_len);

            const cipher_ctx = BlockCipher.initEnc(key);

            // Compute CBC-MAC using the reusable CBC-MAC module
            var mac_result: [block_length]u8 = undefined;
            computeCbcMac(&mac_result, &key, m, ad, npub);

            // Construct counter block for tag encryption (counter = 0)
            var ctr_block: [block_length]u8 = undefined;
            formatCtrBlock(&ctr_block, npub, 0);

            // Encrypt the MAC tag
            var s0: [block_length]u8 = undefined;
            cipher_ctx.encrypt(&s0, &ctr_block);
            for (tag, mac_result[0..tag_length], s0[0..tag_length]) |*t, mac_byte, s_byte| {
                t.* = mac_byte ^ s_byte;
            }

            // Encrypt the plaintext using CTR mode (starting from counter = 1)
            formatCtrBlock(&ctr_block, npub, 1);
            const counter_offset = 1 + nonce_length;
            const counter_size = L;
            modes.ctrSlice(BlockCipherCtx, cipher_ctx, c, m, ctr_block, .big, counter_offset, counter_size);
        }

        /// Decrypt a message with associated data
        ///
        /// Parameters:
        /// - m: Plaintext output buffer (must be same length as c)
        /// - c: Ciphertext
        /// - tag: Authentication tag
        /// - ad: Associated data
        /// - npub: Nonce
        /// - key: Encryption key
        ///
        /// Returns: AuthenticationError if tag verification fails
        pub fn decrypt(m: []u8, c: []const u8, tag: [tag_length]u8, ad: []const u8, npub: [nonce_length]u8, key: [key_length]u8) AuthenticationError!void {
            debug.assert(m.len == c.len);

            const cipher_ctx = BlockCipher.initEnc(key);

            // Decrypt the ciphertext using CTR mode (starting from counter = 1)
            var ctr_block: [block_length]u8 = undefined;
            formatCtrBlock(&ctr_block, npub, 1);
            const counter_offset = 1 + nonce_length;
            const counter_size = L;
            modes.ctrSlice(BlockCipherCtx, cipher_ctx, m, c, ctr_block, .big, counter_offset, counter_size);

            // Compute CBC-MAC over decrypted plaintext using the reusable CBC-MAC module
            var mac_result: [block_length]u8 = undefined;
            computeCbcMac(&mac_result, &key, m, ad, npub);

            // Decrypt the received tag
            formatCtrBlock(&ctr_block, npub, 0);
            var s0: [block_length]u8 = undefined;
            cipher_ctx.encrypt(&s0, &ctr_block);

            // Reconstruct the expected MAC
            var expected_mac: [tag_length]u8 = undefined;
            for (&expected_mac, mac_result[0..tag_length], s0[0..tag_length]) |*e, mac_byte, s_byte| {
                e.* = mac_byte ^ s_byte;
            }

            // Constant-time tag comparison
            const valid = crypto.timing_safe.eql([tag_length]u8, expected_mac, tag);
            if (!valid) {
                crypto.secureZero(u8, m);
                return error.AuthenticationFailed;
            }
        }

        /// Format the counter block for CTR mode
        /// Counter block format: [flags | nonce | counter]
        /// flags = L - 1
        fn formatCtrBlock(block: *[block_length]u8, npub: [nonce_length]u8, counter: u64) void {
            @memset(block, 0);
            block[0] = L - 1; // flags
            @memcpy(block[1..][0..nonce_length], &npub);
            // Counter goes in the last L bytes
            const CounterInt = std.meta.Int(.unsigned, L * 8);
            mem.writeInt(CounterInt, block[1 + nonce_length ..][0..L], @as(CounterInt, @intCast(counter)), .big);
        }

        /// Compute CBC-MAC over the message and associated data
        ///
        /// Note: CCM uses plain CBC-MAC, NOT CMAC (despite stdlib having CMAC available).
        /// CBC-MAC: Simple MAC = Encrypt(prev_mac XOR block), no subkey derivation
        /// CMAC: Derives K1/K2 subkeys and has special final block handling
        /// These produce DIFFERENT outputs. RFC 3610 requires CBC-MAC for CCM.
        ///
        /// This function uses the reusable CbcMac module from cbc_mac.zig.
        fn computeCbcMac(mac: *[block_length]u8, key: *const [key_length]u8, m: []const u8, ad: []const u8, npub: [nonce_length]u8) void {
            const CbcMac = cbc_mac.CbcMac(BlockCipher);
            var ctx = CbcMac.init(key);

            // Process B_0 block
            var b0: [block_length]u8 = undefined;
            formatB0Block(&b0, m.len, ad.len, npub);
            ctx.update(&b0);

            // Process associated data if present
            // RFC 3610: AD is (encoded_length || ad) padded to block boundary
            if (ad.len > 0) {
                // Encode and add associated data length
                var ad_len_encoding: [10]u8 = undefined;
                const ad_len_size = encodeAdLength(&ad_len_encoding, ad.len);

                // Process AD with padding to block boundary
                ctx.update(ad_len_encoding[0..ad_len_size]);
                ctx.update(ad);

                // Add zero padding to reach block boundary
                const total_ad_size = ad_len_size + ad.len;
                const remainder = total_ad_size % block_length;
                if (remainder > 0) {
                    const padding = [_]u8{0} ** block_length;
                    ctx.update(padding[0 .. block_length - remainder]);
                }
            }

            // Process plaintext message
            ctx.update(m);

            // Finalize MAC
            ctx.final(mac);
        }

        /// Format the B_0 block for CBC-MAC
        /// B_0 format: [flags | nonce | message_length]
        /// flags = 64*Adata + 8*M' + L'
        /// where: Adata = (ad.len > 0), M' = (tag_length - 2)/2, L' = L - 1
        fn formatB0Block(block: *[block_length]u8, msg_len: usize, ad_len: usize, npub: [nonce_length]u8) void {
            @memset(block, 0);

            const Adata: u8 = if (ad_len > 0) 1 else 0;
            const M_prime: u8 = @intCast((tag_length - 2) / 2);
            const L_prime: u8 = L - 1;

            block[0] = (Adata << 6) | (M_prime << 3) | L_prime;
            @memcpy(block[1..][0..nonce_length], &npub);

            // Encode message length in last L bytes
            const LengthInt = std.meta.Int(.unsigned, L * 8);
            mem.writeInt(LengthInt, block[1 + nonce_length ..][0..L], @as(LengthInt, @intCast(msg_len)), .big);
        }

        /// Encode associated data length according to CCM specification
        /// Returns the number of bytes written
        fn encodeAdLength(buf: *[10]u8, ad_len: usize) usize {
            if (ad_len < 65280) { // 2^16 - 2^8
                // Encode as 2 bytes
                mem.writeInt(u16, buf[0..2], @as(u16, @intCast(ad_len)), .big);
                return 2;
            } else if (ad_len <= std.math.maxInt(u32)) {
                // Encode as 0xff || 0xfe || 4 bytes
                buf[0] = 0xff;
                buf[1] = 0xfe;
                mem.writeInt(u32, buf[2..6], @as(u32, @intCast(ad_len)), .big);
                return 6;
            } else {
                // Encode as 0xff || 0xff || 8 bytes
                buf[0] = 0xff;
                buf[1] = 0xff;
                mem.writeInt(u64, buf[2..10], @as(u64, @intCast(ad_len)), .big);
                return 10;
            }
        }
    };
}

const testing = std.testing;

test "Aes128Ccm - RFC 3610 Packet Vector #1" {
    // MAC: 8 octets, Message: 23 octets, Nonce: 13 octets
    const Ccm813 = Ccm(crypto.core.aes.Aes128, 8, 13);

    const key = [_]u8{ 0xC0, 0xC1, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xCB, 0xCC, 0xCD, 0xCE, 0xCF };
    const nonce = [_]u8{ 0x00, 0x00, 0x00, 0x03, 0x02, 0x01, 0x00, 0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5 };
    const ad = [_]u8{ 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07 };
    const plaintext = [_]u8{ 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E };

    const expected_ciphertext = [_]u8{ 0x58, 0x8C, 0x97, 0x9A, 0x61, 0xC6, 0x63, 0xD2, 0xF0, 0x66, 0xD0, 0xC2, 0xC0, 0xF9, 0x89, 0x80, 0x6D, 0x5F, 0x6B, 0x61, 0xDA, 0xC3, 0x84 };
    const expected_tag = [_]u8{ 0x17, 0xE8, 0xD1, 0x2C, 0xFD, 0xF9, 0x26, 0xE0 };

    var ciphertext: [plaintext.len]u8 = undefined;
    var tag: [8]u8 = undefined;

    Ccm813.encrypt(&ciphertext, &tag, &plaintext, &ad, nonce, key);

    try testing.expectEqualSlices(u8, &expected_ciphertext, &ciphertext);
    try testing.expectEqualSlices(u8, &expected_tag, &tag);

    // Test decryption
    var decrypted: [plaintext.len]u8 = undefined;
    try Ccm813.decrypt(&decrypted, &ciphertext, tag, &ad, nonce, key);
    try testing.expectEqualSlices(u8, &plaintext, &decrypted);
}

test "Aes128Ccm - Standard parameters (tag=16, nonce=12)" {
    const key: [Aes128Ccm.key_length]u8 = [_]u8{0x42} ** Aes128Ccm.key_length;
    const nonce: [12]u8 = [_]u8{0x01} ** 12;
    const ad = "Additional data";
    const plaintext = "Hello, CCM mode!";

    var ciphertext: [plaintext.len]u8 = undefined;
    var tag: [16]u8 = undefined;

    Aes128Ccm.encrypt(&ciphertext, &tag, plaintext, ad, nonce, key);

    // Test decryption
    var decrypted: [plaintext.len]u8 = undefined;
    try Aes128Ccm.decrypt(&decrypted, &ciphertext, tag, ad, nonce, key);
    try testing.expectEqualSlices(u8, plaintext, &decrypted);

    // Test that wrong tag fails
    tag[0] ^= 1;
    try testing.expectError(error.AuthenticationFailed, Aes128Ccm.decrypt(&decrypted, &ciphertext, tag, ad, nonce, key));
}

test "Aes128Ccm - Empty message" {
    const key: [Aes128Ccm.key_length]u8 = [_]u8{0x42} ** Aes128Ccm.key_length;
    const nonce: [12]u8 = [_]u8{0x01} ** 12;
    const ad = "Additional data";
    const plaintext = "";

    var ciphertext: [plaintext.len]u8 = undefined;
    var tag: [16]u8 = undefined;

    Aes128Ccm.encrypt(&ciphertext, &tag, plaintext, ad, nonce, key);

    var decrypted: [plaintext.len]u8 = undefined;
    try Aes128Ccm.decrypt(&decrypted, &ciphertext, tag, ad, nonce, key);
    try testing.expectEqualSlices(u8, plaintext, &decrypted);
}

test "Aes128Ccm - No associated data" {
    const key: [Aes128Ccm.key_length]u8 = [_]u8{0x42} ** Aes128Ccm.key_length;
    const nonce: [12]u8 = [_]u8{0x01} ** 12;
    const ad = "";
    const plaintext = "Hello, CCM mode!";

    var ciphertext: [plaintext.len]u8 = undefined;
    var tag: [16]u8 = undefined;

    Aes128Ccm.encrypt(&ciphertext, &tag, plaintext, ad, nonce, key);

    var decrypted: [plaintext.len]u8 = undefined;
    try Aes128Ccm.decrypt(&decrypted, &ciphertext, tag, ad, nonce, key);
    try testing.expectEqualSlices(u8, plaintext, &decrypted);
}

test "Aes256Ccm - Basic test" {
    const key: [Aes256Ccm.key_length]u8 = [_]u8{0x42} ** Aes256Ccm.key_length;
    const nonce: [12]u8 = [_]u8{0x01} ** 12;
    const ad = "Additional data";
    const plaintext = "Hello, AES-256-CCM!";

    var ciphertext: [plaintext.len]u8 = undefined;
    var tag: [16]u8 = undefined;

    Aes256Ccm.encrypt(&ciphertext, &tag, plaintext, ad, nonce, key);

    var decrypted: [plaintext.len]u8 = undefined;
    try Aes256Ccm.decrypt(&decrypted, &ciphertext, tag, ad, nonce, key);
    try testing.expectEqualSlices(u8, plaintext, &decrypted);
}