// SPDX-FileCopyrightText: © 2022—2026 Vladimir Zorin <vladimir@deviant.guru>
// SPDX-License-Identifier: LicenseRef-OWL-1.0-or-later
// Licensed under OWL v1.0+. See LICENSE.
/*
 * LITLS -- TLS 1.3 Handshake Messages
 *
 * ClientHello builder, ServerHello/EE/Certificate/CV/Finished parsers,
 * transcript hash management, and handshake reassembly.
 */

#include "tls13.h"
#include <stdio.h>
#include <string.h>

/* HelloRetryRequest sentinel random (SHA-256 of "HelloRetryRequest") */
const uint8_t TLS13_HRR_RANDOM[32] = {0xCF, 0x21, 0xAD, 0x74, 0xE5, 0x9A, 0x61, 0x11, 0xBE, 0x1D, 0x8C,
                                      0x02, 0x1E, 0x65, 0xB8, 0x91, 0xC2, 0xA2, 0x11, 0x16, 0x7A, 0xBB,
                                      0x8C, 0x5E, 0x07, 0x9E, 0x09, 0xE2, 0xC8, 0xA8, 0x33, 0x9C};

/* ── Helper: put bytes ────────────────────────────────────────────── */

static inline void put16(uint8_t *p, uint16_t v) {
    p[0] = (uint8_t)(v >> 8);
    p[1] = (uint8_t)(v);
}

static inline void put24(uint8_t *p, uint32_t v) {
    p[0] = (uint8_t)(v >> 16);
    p[1] = (uint8_t)(v >> 8);
    p[2] = (uint8_t)(v);
}

static inline uint16_t get16(const uint8_t *p) {
    return (uint16_t)((p[0] << 8) | p[1]);
}

static inline uint32_t get24(const uint8_t *p) {
    return ((uint32_t)p[0] << 16) | ((uint32_t)p[1] << 8) | p[2];
}

/* ── Transcript hash ──────────────────────────────────────────────── */

void litls_transcript_init(litls_transcript *t) {
    litls_sha256_init(&t->sha256);
    litls_sha384_init(&t->sha384);
    t->use_384 = -1; /* both active until suite is selected */
}

void litls_transcript_update(litls_transcript *t, const uint8_t *data, size_t len) {
    if (t->use_384 <= 0)
        litls_sha256_update(&t->sha256, data, len);
    if (t->use_384 != 0)
        litls_sha384_update(&t->sha384, data, len);
}

void litls_transcript_hash(const litls_transcript *t, uint8_t *out) {
    if (t->use_384 == 1) {
        /* Copy context to avoid modifying the running hash */
        litls_sha384_ctx tmp = t->sha384;
        litls_sha384_final(&tmp, out);
    } else {
        litls_sha256_ctx tmp = t->sha256;
        litls_sha256_final(&tmp, out);
    }
}

int litls_transcript_hash_len(const litls_transcript *t) {
    return (t->use_384 == 1) ? 48 : 32;
}

/* ── ClientHello builder ──────────────────────────────────────────── */

int litls_build_client_hello(litls_tls13_ctx *ctx, uint8_t *out, size_t cap, size_t *out_len) {
    uint8_t *p   = out;
    uint8_t *end = out + cap;

    /* We build the handshake message body first, then prepend the header */

    /* Skip 4-byte handshake header for now */
    uint8_t *body = p + 4;
    uint8_t *b    = body;

    /* client_version = TLS 1.2 (legacy) */
    if (b + 2 > end)
        return -1;
    put16(b, TLS13_VERSION_12);
    b += 2;

    /* random (32 bytes) */
    if (b + 32 > end)
        return -1;
    litls_random_bytes(b, 32);
    b += 32;

    /* legacy_session_id (32 bytes for middlebox compat) */
    if (b + 33 > end)
        return -1;
    b[0] = 32; /* length */
    litls_random_bytes(b + 1, 32);
    b += 33;

    /* cipher_suites (3 suites = 6 bytes + 2 length) */
    if (b + 8 > end)
        return -1;
    put16(b, 6); /* length */
    b += 2;
    put16(b, TLS13_CHACHA20_POLY1305_SHA256);
    b += 2;
    put16(b, TLS13_AES_128_GCM_SHA256);
    b += 2;
    put16(b, TLS13_AES_256_GCM_SHA384);
    b += 2;

    /* legacy_compression_methods = {null} */
    if (b + 2 > end)
        return -1;
    b[0] = 1; /* length */
    b[1] = 0; /* null */
    b += 2;

    /* ── Extensions ── */
    uint8_t *ext_len_ptr = b;
    b += 2; /* will fill in extensions length later */

    /* supported_versions (TLS 1.3 only) */
    if (b + 7 > end)
        return -1;
    put16(b, TLS13_EXT_SUPPORTED_VERSIONS);
    b += 2;
    put16(b, 3); /* extension data length */
    b += 2;
    b[0] = 2; /* list length */
    b += 1;
    put16(b, TLS13_VERSION_13);
    b += 2;

    /* supported_groups (x25519 only) */
    if (b + 8 > end)
        return -1;
    put16(b, TLS13_EXT_SUPPORTED_GROUPS);
    b += 2;
    put16(b, 4); /* extension data length */
    b += 2;
    put16(b, 2); /* list length */
    b += 2;
    put16(b, TLS13_GROUP_X25519);
    b += 2;

    /* key_share (x25519 public key) */
    if (b + 42 > end)
        return -1;
    put16(b, TLS13_EXT_KEY_SHARE);
    b += 2;
    put16(b, 38); /* extension data length */
    b += 2;
    put16(b, 36); /* client shares length */
    b += 2;
    put16(b, TLS13_GROUP_X25519);
    b += 2;
    put16(b, 32); /* key exchange length */
    b += 2;
    memcpy(b, ctx->x25519_pub, 32);
    b += 32;

    /* signature_algorithms (RFC 8446 §4.2.3 forbids RSASSA-PKCS1-v1_5 for TLS
     * 1.3 handshake signatures, so rsa_pkcs1_sha256 is not advertised here.
     * RSA PKCS#1 v1.5 in *certificate chain* signatures is still accepted by
     * `litls_x509_verify_chain`.) */
    if (b + 12 > end)
        return -1;
    put16(b, TLS13_EXT_SIGNATURE_ALGORITHMS);
    b += 2;
    put16(b, 8); /* extension data length */
    b += 2;
    put16(b, 6); /* list length */
    b += 2;
    put16(b, TLS13_SIG_ECDSA_SECP256R1_SHA256);
    b += 2;
    put16(b, TLS13_SIG_RSA_PSS_RSAE_SHA256);
    b += 2;
    put16(b, TLS13_SIG_ED25519);
    b += 2;

    /* server_name (SNI) */
    if (ctx->sni[0] != '\0') {
        size_t sni_len      = strlen(ctx->sni);
        size_t ext_data_len = sni_len + 5; /* list_len(2) + type(1) + name_len(2) + name */
        if (b + 4 + ext_data_len > end)
            return -1;
        put16(b, TLS13_EXT_SERVER_NAME);
        b += 2;
        put16(b, (uint16_t)ext_data_len);
        b += 2;
        put16(b, (uint16_t)(sni_len + 3)); /* server name list length */
        b += 2;
        b[0] = 0; /* host_name type */
        b += 1;
        put16(b, (uint16_t)sni_len);
        b += 2;
        memcpy(b, ctx->sni, sni_len);
        b += sni_len;
    }

    /* Cookie extension (for HelloRetryRequest) */
    if (ctx->hrr_cookie_len > 0) {
        size_t ext_data_len = 2 + ctx->hrr_cookie_len;
        if (b + 4 + ext_data_len > end)
            return -1;
        put16(b, TLS13_EXT_COOKIE);
        b += 2;
        put16(b, (uint16_t)ext_data_len);
        b += 2;
        put16(b, (uint16_t)ctx->hrr_cookie_len);
        b += 2;
        memcpy(b, ctx->hrr_cookie, ctx->hrr_cookie_len);
        b += ctx->hrr_cookie_len;
    }

    /* Fill in extensions length */
    size_t ext_total = (size_t)(b - ext_len_ptr - 2);
    put16(ext_len_ptr, (uint16_t)ext_total);

    /* Fill in handshake header */
    size_t body_len = (size_t)(b - body);
    p[0]            = TLS13_HT_CLIENT_HELLO;
    put24(p + 1, (uint32_t)body_len);

    *out_len = (size_t)(b - out);
    return 0;
}

/* ── Shared helpers (used by both client and server) ──────────────── */

static inline uint32_t _hs_get24(const uint8_t *p) {
    return ((uint32_t)p[0] << 16) | ((uint32_t)p[1] << 8) | p[2];
}

int litls_get_hs_message(litls_tls13_ctx *ctx, uint8_t *msg_type, const uint8_t **msg, size_t *msg_len) {
    /* Check if hs_buf already has a complete message */
    if (ctx->hs_buf_len >= 4) {
        uint32_t total = _hs_get24(ctx->hs_buf + 1);
        if (ctx->hs_buf_len >= 4 + total) {
            *msg_type = ctx->hs_buf[0];
            *msg      = ctx->hs_buf + 4;
            *msg_len  = total;
            return 1;
        }
    }

    /* Need more data from a record */
    uint8_t ct;
    uint8_t *data;
    size_t data_len;
    int rc = litls_record_read(ctx, &ct, &data, &data_len);
    if (rc == 1)
        return 0;
    if (rc < 0)
        return rc;

    if (ct != TLS13_CT_HANDSHAKE) {
        snprintf(ctx->error, sizeof(ctx->error), "expected handshake, got type %d", ct);
        return -1;
    }

    return litls_hs_reassemble(ctx, data, data_len, msg_type, msg, msg_len);
}

void litls_consume_hs_message(litls_tls13_ctx *ctx, size_t msg_len) {
    size_t consumed = 4 + msg_len;
    if (ctx->hs_buf_len > consumed) {
        memmove(ctx->hs_buf, ctx->hs_buf + consumed, ctx->hs_buf_len - consumed);
        ctx->hs_buf_len -= consumed;
    } else {
        ctx->hs_buf_len = 0;
    }
}

int litls_init_cipher_suite(litls_tls13_ctx *ctx, uint16_t suite_id) {
    switch (suite_id) {
    case TLS13_CHACHA20_POLY1305_SHA256:
        ctx->suite.id           = suite_id;
        ctx->suite.hash_len     = 32;
        ctx->suite.key_len      = 32;
        ctx->suite.iv_len       = 12;
        ctx->suite.encrypt      = (litls_aead_encrypt_fn)litls_chacha20_poly1305_encrypt;
        ctx->suite.decrypt      = (litls_aead_decrypt_fn)litls_chacha20_poly1305_decrypt;
        ctx->transcript.use_384 = 0;
        return 0;
    case TLS13_AES_128_GCM_SHA256:
        ctx->suite.id           = suite_id;
        ctx->suite.hash_len     = 32;
        ctx->suite.key_len      = 16;
        ctx->suite.iv_len       = 12;
        ctx->suite.encrypt      = (litls_aead_encrypt_fn)litls_aes128_gcm_encrypt;
        ctx->suite.decrypt      = (litls_aead_decrypt_fn)litls_aes128_gcm_decrypt;
        ctx->transcript.use_384 = 0;
        return 0;
    case TLS13_AES_256_GCM_SHA384:
        ctx->suite.id           = suite_id;
        ctx->suite.hash_len     = 48;
        ctx->suite.key_len      = 32;
        ctx->suite.iv_len       = 12;
        ctx->suite.encrypt      = (litls_aead_encrypt_fn)litls_aes256_gcm_encrypt;
        ctx->suite.decrypt      = (litls_aead_decrypt_fn)litls_aes256_gcm_decrypt;
        ctx->transcript.use_384 = 1;
        return 0;
    default:
        snprintf(ctx->error, sizeof(ctx->error), "unsupported cipher suite 0x%04X", suite_id);
        return -1;
    }
}

int litls_build_finished(litls_tls13_ctx *ctx, const uint8_t *hs_secret, uint8_t *out, size_t *out_len) {
    int hash_len = ctx->suite.hash_len;
    uint8_t transcript_hash[48];
    litls_transcript_hash(&ctx->transcript, transcript_hash);

    uint8_t finished_key[48];
    uint8_t verify_data[48];

    if (hash_len == 32) {
        litls_tls13_expand_label_256(hs_secret, "finished", NULL, 0, finished_key, 32);
        litls_hmac_sha256(finished_key, 32, transcript_hash, 32, verify_data);
    } else {
        litls_tls13_expand_label_384(hs_secret, "finished", NULL, 0, finished_key, 48);
        litls_hmac_sha384(finished_key, 48, transcript_hash, 48, verify_data);
    }

    out[0] = TLS13_HT_FINISHED;
    out[1] = 0;
    out[2] = 0;
    out[3] = (uint8_t)hash_len;
    memcpy(out + 4, verify_data, (size_t)hash_len);

    *out_len = 4 + (size_t)hash_len;

    litls_secure_zero(finished_key, sizeof(finished_key));
    litls_secure_zero(verify_data, sizeof(verify_data));
    litls_secure_zero(transcript_hash, sizeof(transcript_hash));
    return 0;
}

/* ── ServerHello parser ───────────────────────────────────────────── */

int litls_parse_server_hello(litls_tls13_ctx *ctx, const uint8_t *data, size_t len) {
    if (len < 2 + 32 + 1)
        goto err;

    const uint8_t *p = data;
    /* server_version (legacy) */
    p += 2;

    /* random */
    const uint8_t *server_random = p;
    p += 32;

    /* Check for HelloRetryRequest */
    if (memcmp(server_random, TLS13_HRR_RANDOM, 32) == 0) {
        /* This is an HRR, not a real ServerHello. RFC 8446 §4.1.4: a client
         * that receives a second HRR in the same handshake MUST abort. */
        if (ctx->hrr_done) {
            snprintf(ctx->error, sizeof(ctx->error), "second HelloRetryRequest");
            return -1;
        }
        ctx->hrr_seen = 1;
    }

    /* legacy_session_id_echo */
    if (p >= data + len)
        goto err;
    uint8_t sess_id_len = *p++;
    if (p + sess_id_len > data + len)
        goto err;
    p += sess_id_len;

    /* cipher_suite */
    if (p + 2 > data + len)
        goto err;
    uint16_t suite_id = get16(p);
    p += 2;

    if (ctx->hrr_seen) {
        /* Remember the suite the server commits to in the HRR. */
        ctx->hrr_suite = suite_id;
    } else if (ctx->hrr_done) {
        /* RFC 8446 §4.1.4: the ServerHello after an HRR MUST NOT change the
         * cipher suite the server committed to in the HRR. A different suite
         * (different hash) would also corrupt the transcript hash. */
        if (suite_id != ctx->hrr_suite) {
            snprintf(ctx->error, sizeof(ctx->error), "cipher suite changed after HelloRetryRequest");
            return -1;
        }
    }

    if (litls_init_cipher_suite(ctx, suite_id) != 0)
        return -1;

    /* legacy_compression_method (must be 0) */
    if (p + 1 > data + len)
        goto err;
    if (*p++ != 0)
        goto err;

    /* Extensions */
    if (p + 2 > data + len)
        goto err;
    uint16_t ext_len = get16(p);
    p += 2;

    const uint8_t *ext_end = p + ext_len;
    if (ext_end > data + len)
        goto err;

    int got_version   = 0;
    int got_key_share = 0;

    while (p + 4 <= ext_end) {
        uint16_t etype = get16(p);
        p += 2;
        uint16_t elen = get16(p);
        p += 2;

        if (p + elen > ext_end)
            goto err;

        switch (etype) {
        case TLS13_EXT_SUPPORTED_VERSIONS:
            if (elen != 2)
                goto err;
            if (get16(p) != TLS13_VERSION_13)
                goto err;
            got_version = 1;
            break;

        case TLS13_EXT_KEY_SHARE:
            if (ctx->hrr_seen) {
                /* HRR key_share contains only the selected group */
                if (elen < 2)
                    goto err;
                uint16_t group = get16(p);
                if (group != TLS13_GROUP_X25519) {
                    snprintf(ctx->error, sizeof(ctx->error), "HRR requested unsupported group 0x%04X", group);
                    return -1;
                }
            } else {
                /* ServerHello key_share: group(2) + key_len(2) + key(32) */
                if (elen < 36)
                    goto err;
                uint16_t group   = get16(p);
                uint16_t key_len = get16(p + 2);
                if (group != TLS13_GROUP_X25519 || key_len != 32) {
                    snprintf(ctx->error, sizeof(ctx->error), "unsupported key share group/length");
                    return -1;
                }
                /* Compute shared secret */
                if (litls_x25519(ctx->shared_secret, ctx->x25519_priv, p + 4) != 0) {
                    snprintf(ctx->error, sizeof(ctx->error), "X25519 failed");
                    return -1;
                }
            }
            got_key_share = 1;
            break;

        case TLS13_EXT_COOKIE:
            if (ctx->hrr_seen && elen >= 2) {
                uint16_t cookie_len = get16(p);
                if (cookie_len > sizeof(ctx->hrr_cookie) || (size_t)(cookie_len + 2) > elen)
                    goto err;
                memcpy(ctx->hrr_cookie, p + 2, cookie_len);
                ctx->hrr_cookie_len = cookie_len;
            }
            break;

        default:
            break; /* ignore unknown extensions */
        }

        p += elen;
    }

    if (!got_version) {
        snprintf(ctx->error, sizeof(ctx->error), "no supported_versions in ServerHello");
        return -1;
    }
    if (!got_key_share) {
        snprintf(ctx->error, sizeof(ctx->error), "no key_share in ServerHello");
        return -1;
    }

    return 0;

err:
    snprintf(ctx->error, sizeof(ctx->error), "malformed ServerHello");
    return -1;
}

/* ── EncryptedExtensions parser ───────────────────────────────────── */

int litls_parse_encrypted_extensions(litls_tls13_ctx *ctx, const uint8_t *data, size_t len) {
    /* Just validate minimum structure. We don't use ALPN or early data. */
    if (len < 2)
        goto err;
    uint16_t ext_len = get16(data);
    if ((size_t)(ext_len + 2) > len)
        goto err;
    (void)ctx;
    return 0;

err:
    snprintf(ctx->error, sizeof(ctx->error), "malformed EncryptedExtensions");
    return -1;
}

/* ── CertificateRequest parser ────────────────────────────────────── */

int litls_parse_certificate_request(litls_tls13_ctx *ctx, const uint8_t *data, size_t len) {
    if (len < 3)
        goto err;

    const uint8_t *p = data;

    /* certificate_request_context <0..255> */
    uint8_t ctx_len = *p++;
    if (p + ctx_len > data + len)
        goto err;
    if (ctx_len > 0)
        memcpy(ctx->cert_request_ctx, p, ctx_len);
    ctx->cert_request_ctx_len = ctx_len;
    p += ctx_len;

    /* extensions <2..2^16-1> -- skip for V1 */
    if (p + 2 > data + len)
        goto err;
    uint16_t ext_len = get16(p);
    p += 2;
    if (p + ext_len > data + len)
        goto err;

    ctx->cert_requested = 1;
    return 0;

err:
    snprintf(ctx->error, sizeof(ctx->error), "malformed CertificateRequest");
    return -1;
}

/* ── Certificate parser ───────────────────────────────────────────── */

int litls_parse_certificate(litls_tls13_ctx *ctx, const uint8_t *data, size_t len, const uint8_t **chain_ders,
                            size_t *chain_lens, int *chain_count, int max_chain) {
    if (len < 4)
        goto err;

    const uint8_t *p = data;

    /* certificate_request_context (must be empty for server certs) */
    uint8_t ctx_len = *p++;
    if (ctx_len != 0)
        goto err;

    /* certificate_list length (3 bytes) */
    if (p + 3 > data + len)
        goto err;
    uint32_t list_len = get24(p);
    p += 3;

    if (p + list_len > data + len)
        goto err;

    const uint8_t *list_end = p + list_len;
    int count               = 0;

    while (p + 3 <= list_end && count < max_chain) {
        uint32_t cert_len = get24(p);
        p += 3;
        if (p + cert_len > list_end)
            goto err;

        chain_ders[count] = p;
        chain_lens[count] = cert_len;
        count++;
        p += cert_len;

        /* Per-certificate extensions (2 bytes length + data) */
        if (p + 2 > list_end)
            goto err;
        uint16_t ext_len = get16(p);
        p += 2;
        if (p + ext_len > list_end)
            goto err;
        p += ext_len;
    }

    if (count == 0)
        goto err;

    *chain_count = count;
    return 0;

err:
    snprintf(ctx->error, sizeof(ctx->error), "malformed Certificate message");
    return -1;
}

/* ── CertificateVerify parser ─────────────────────────────────────── */

int litls_parse_certificate_verify(litls_tls13_ctx *ctx, const uint8_t *data, size_t len, uint16_t *sig_algo,
                                   const uint8_t **sig_out, size_t *sig_len_out) {
    if (len < 4)
        goto err;

    *sig_algo        = get16(data);
    uint16_t sig_len = get16(data + 2);

    if (4 + sig_len > len)
        goto err;

    *sig_out     = data + 4;
    *sig_len_out = sig_len;
    return 0;

err:
    snprintf(ctx->error, sizeof(ctx->error), "malformed CertificateVerify");
    return -1;
}

/* ── Finished parser ──────────────────────────────────────────────── */

int litls_parse_finished(litls_tls13_ctx *ctx, const uint8_t *data, size_t len, const uint8_t **verify_data,
                         size_t *verify_len) {
    int hash_len = ctx->suite.hash_len;
    /* RFC 8446 §4.4.4: Finished verify_data is exactly Hash.length bytes. */
    if (len != (size_t)hash_len) {
        snprintf(ctx->error, sizeof(ctx->error), "Finished length mismatch (got %zu, expected %d)", len, hash_len);
        return -1;
    }
    *verify_data = data;
    *verify_len  = (size_t)hash_len;
    return 0;
}

/* ── Handshake reassembly ─────────────────────────────────────────── */

int litls_hs_reassemble(litls_tls13_ctx *ctx, const uint8_t *data, size_t data_len, uint8_t *msg_type,
                        const uint8_t **msg, size_t *msg_len) {
    /* Append incoming data to reassembly buffer */
    if (ctx->hs_buf_len + data_len > TLS13_HS_BUF_SIZE) {
        snprintf(ctx->error, sizeof(ctx->error), "handshake message too large");
        return -1;
    }
    memcpy(ctx->hs_buf + ctx->hs_buf_len, data, data_len);
    ctx->hs_buf_len += data_len;

    /* Check if we have a complete handshake message */
    if (ctx->hs_buf_len < 4)
        return 0; /* need more data */

    uint8_t type       = ctx->hs_buf[0];
    uint32_t total_len = get24(ctx->hs_buf + 1);

    if (ctx->hs_buf_len < 4 + total_len)
        return 0; /* need more data */

    /* Complete message available */
    *msg_type = type;
    *msg      = ctx->hs_buf + 4;
    *msg_len  = total_len;

    return 1; /* message ready */
}

/* ── ClientHello parser (server-side) ─────────────────────────────── */

int litls_parse_client_hello(litls_tls13_ctx *ctx, const uint8_t *data, size_t len) {
    if (len < 2 + 32 + 1)
        goto err;

    const uint8_t *p = data;

    /* client_version (legacy 0x0303) */
    p += 2;

    /* random (32 bytes) -- skip */
    p += 32;

    /* legacy_session_id */
    if (p >= data + len)
        goto err;
    uint8_t sess_id_len = *p++;
    if (sess_id_len > 32 || p + sess_id_len > data + len)
        goto err;
    ctx->session_id_len = sess_id_len;
    if (sess_id_len > 0)
        memcpy(ctx->session_id, p, sess_id_len);
    p += sess_id_len;

    /* cipher_suites */
    if (p + 2 > data + len)
        goto err;
    uint16_t cs_len = get16(p);
    p += 2;
    if (p + cs_len > data + len || cs_len < 2 || (cs_len & 1))
        goto err;

    /* Server preference order: ChaCha20 > AES-128-GCM > AES-256-GCM */
    static const uint16_t preferred[] = {TLS13_CHACHA20_POLY1305_SHA256, TLS13_AES_128_GCM_SHA256,
                                         TLS13_AES_256_GCM_SHA384};
    uint16_t selected_suite           = 0;
    int best_prio                     = 999;

    for (size_t i = 0; i < cs_len; i += 2) {
        uint16_t cs = get16(p + i);
        for (int j = 0; j < 3; j++) {
            if (cs == preferred[j] && j < best_prio) {
                best_prio      = j;
                selected_suite = cs;
            }
        }
    }
    p += cs_len;

    if (!selected_suite) {
        snprintf(ctx->error, sizeof(ctx->error), "no supported cipher suite");
        return -1;
    }

    /* legacy_compression_methods */
    if (p >= data + len)
        goto err;
    uint8_t comp_len = *p++;
    if (p + comp_len > data + len)
        goto err;
    p += comp_len;

    /* Extensions */
    if (p + 2 > data + len)
        goto err;
    uint16_t ext_len = get16(p);
    p += 2;
    const uint8_t *ext_end = p + ext_len;
    if (ext_end > data + len)
        goto err;

    int got_version          = 0;
    int got_key_share        = 0;
    ctx->peer_sig_algo_count = 0;
    ctx->peer_sent_sig_algos = 0;

    while (p + 4 <= ext_end) {
        uint16_t etype = get16(p);
        p += 2;
        uint16_t elen = get16(p);
        p += 2;
        if (p + elen > ext_end)
            goto err;

        switch (etype) {
        case TLS13_EXT_SUPPORTED_VERSIONS: {
            /* List of versions: 1-byte list_len, then 2-byte versions.
             * RFC 8446 §4.2.1: versions<2..254>, vector of 2-byte
             * ProtocolVersion -- list_len must be even and the
             * extension body must be exactly list_len + 1. */
            if (elen < 1)
                goto err;
            uint8_t list_len = p[0];
            if (list_len < 2 || (list_len & 1) || (size_t)list_len + 1 != elen)
                goto err;
            int found_13 = 0;
            for (int i = 0; i < list_len; i += 2) {
                if (get16(p + 1 + i) == TLS13_VERSION_13) {
                    found_13 = 1;
                    break;
                }
            }
            if (!found_13) {
                snprintf(ctx->error, sizeof(ctx->error), "client does not support TLS 1.3");
                return -1;
            }
            got_version = 1;
            break;
        }

        case TLS13_EXT_KEY_SHARE: {
            /* client_shares_len(2) + entries */
            if (elen < 2)
                goto err;
            uint16_t shares_len = get16(p);
            const uint8_t *sp   = p + 2;
            const uint8_t *send = p + 2 + shares_len;
            if (send > p + elen)
                goto err;
            while (sp + 4 <= send) {
                uint16_t group   = get16(sp);
                uint16_t key_len = get16(sp + 2);
                sp += 4;
                if (sp + key_len > send)
                    break;
                if (group == TLS13_GROUP_X25519 && key_len == 32) {
                    /* Compute shared secret using our private key + client's public */
                    if (litls_x25519(ctx->shared_secret, ctx->x25519_priv, sp) != 0) {
                        snprintf(ctx->error, sizeof(ctx->error), "X25519 failed");
                        return -1;
                    }
                    got_key_share = 1;
                }
                sp += key_len;
            }
            break;
        }

        case TLS13_EXT_SERVER_NAME: {
            /* server_name_list: list_len(2) + entries */
            if (elen < 2)
                break;
            uint16_t list_len = get16(p);
            const uint8_t *sp = p + 2;
            if ((size_t)(list_len + 2) > elen)
                break;
            while (sp + 3 <= p + 2 + list_len) {
                uint8_t name_type = *sp++;
                uint16_t name_len = get16(sp);
                sp += 2;
                if (sp + name_len > p + 2 + list_len)
                    break;
                if (name_type == 0 && name_len < TLS13_MAX_SNI_LEN) {
                    memcpy(ctx->sni, sp, name_len);
                    ctx->sni[name_len] = '\0';
                }
                sp += name_len;
            }
            break;
        }

        case TLS13_EXT_SIGNATURE_ALGORITHMS: {
            ctx->peer_sent_sig_algos = 1;
            if (elen < 2)
                break;
            uint16_t list_len = get16(p);
            int count         = 0;
            for (size_t i = 0; i < list_len && i + 1 < elen - 2 && count < 8; i += 2) {
                ctx->peer_sig_algos[count++] = get16(p + 2 + i);
            }
            ctx->peer_sig_algo_count = count;
            break;
        }

        default:
            break;
        }

        p += elen;
    }

    if (!got_version) {
        snprintf(ctx->error, sizeof(ctx->error), "no supported_versions in ClientHello");
        return -1;
    }
    if (!got_key_share) {
        snprintf(ctx->error, sizeof(ctx->error), "no X25519 key_share in ClientHello");
        return -1;
    }

    /* Initialize the selected cipher suite */
    if (litls_init_cipher_suite(ctx, selected_suite) != 0)
        return -1;

    return 0;

err:
    snprintf(ctx->error, sizeof(ctx->error), "malformed ClientHello");
    return -1;
}

/* ── ServerHello builder ──────────────────────────────────────────── */

int litls_build_server_hello(litls_tls13_ctx *ctx, uint8_t *out, size_t cap, size_t *out_len) {
    uint8_t *p   = out;
    uint8_t *end = out + cap;

    /* Skip 4-byte handshake header */
    uint8_t *body = p + 4;
    uint8_t *b    = body;

    /* server_version = TLS 1.2 (legacy) */
    if (b + 2 > end)
        return -1;
    put16(b, TLS13_VERSION_12);
    b += 2;

    /* random (32 bytes) */
    if (b + 32 > end)
        return -1;
    litls_random_bytes(b, 32);
    b += 32;

    /* legacy_session_id_echo */
    if (b + 1 + ctx->session_id_len > end)
        return -1;
    *b++ = ctx->session_id_len;
    if (ctx->session_id_len > 0) {
        memcpy(b, ctx->session_id, ctx->session_id_len);
        b += ctx->session_id_len;
    }

    /* cipher_suite */
    if (b + 2 > end)
        return -1;
    put16(b, ctx->suite.id);
    b += 2;

    /* legacy_compression_method = null */
    if (b + 1 > end)
        return -1;
    *b++ = 0;

    /* Extensions */
    uint8_t *ext_len_ptr = b;
    b += 2;

    /* supported_versions */
    if (b + 6 > end)
        return -1;
    put16(b, TLS13_EXT_SUPPORTED_VERSIONS);
    b += 2;
    put16(b, 2);
    b += 2;
    put16(b, TLS13_VERSION_13);
    b += 2;

    /* key_share (X25519) */
    if (b + 40 > end)
        return -1;
    put16(b, TLS13_EXT_KEY_SHARE);
    b += 2;
    put16(b, 36); /* extension data length: group(2) + key_len(2) + key(32) */
    b += 2;
    put16(b, TLS13_GROUP_X25519);
    b += 2;
    put16(b, 32);
    b += 2;
    memcpy(b, ctx->x25519_pub, 32);
    b += 32;

    /* Fill in extensions length */
    put16(ext_len_ptr, (uint16_t)(b - ext_len_ptr - 2));

    /* Fill in handshake header */
    size_t body_len = (size_t)(b - body);
    p[0]            = TLS13_HT_SERVER_HELLO;
    put24(p + 1, (uint32_t)body_len);

    *out_len = (size_t)(b - out);
    return 0;
}

/* ── EncryptedExtensions builder ──────────────────────────────────── */

int litls_build_encrypted_extensions(uint8_t *out, size_t cap, size_t *out_len) {
    if (cap < 6)
        return -1;
    out[0] = TLS13_HT_ENCRYPTED_EXTENSIONS;
    put24(out + 1, 2); /* body length = 2 */
    put16(out + 4, 0); /* empty extensions */
    *out_len = 6;
    return 0;
}

/* ── Certificate builder ──────────────────────────────────────────── */

int litls_build_certificate(const litls_server_identity *id, const uint8_t *cert_req_ctx, size_t cert_req_ctx_len,
                            uint8_t *out, size_t cap, size_t *out_len) {
    /* Calculate total size:
     * header(4) + ctx_len(1) + ctx + cert_list_len(3) +
     * for each cert: cert_data_len(3) + cert_data + ext_len(2) */
    size_t cert_list_size = 0;
    if (id) {
        for (int i = 0; i < id->cert_count; i++)
            cert_list_size += 3 + id->cert_der_lens[i] + 2;
    }

    size_t body_len = 1 + cert_req_ctx_len + 3 + cert_list_size;
    size_t total    = 4 + body_len;
    if (total > cap)
        return -1;

    uint8_t *p = out;
    /* Handshake header */
    *p++       = TLS13_HT_CERTIFICATE;
    put24(p, (uint32_t)body_len);
    p += 3;

    /* certificate_request_context */
    *p++ = (uint8_t)cert_req_ctx_len;
    if (cert_req_ctx_len > 0) {
        memcpy(p, cert_req_ctx, cert_req_ctx_len);
        p += cert_req_ctx_len;
    }

    /* certificate_list */
    put24(p, (uint32_t)cert_list_size);
    p += 3;

    if (id) {
        for (int i = 0; i < id->cert_count; i++) {
            put24(p, (uint32_t)id->cert_der_lens[i]);
            p += 3;
            memcpy(p, id->cert_chain_der + id->cert_der_offsets[i], id->cert_der_lens[i]);
            p += id->cert_der_lens[i];
            put16(p, 0); /* no per-cert extensions */
            p += 2;
        }
    }

    *out_len = total;
    return 0;
}

/* ── CertificateVerify builder ────────────────────────────────────── */

int litls_build_certificate_verify(litls_tls13_ctx *ctx, const litls_server_identity *id, int is_server, uint8_t *out,
                                   size_t cap, size_t *out_len) {
    /* Build content to sign: 64×0x20 + label + 0x00 + transcript_hash */
    uint8_t content[200];
    size_t content_len = 0;

    memset(content, 0x20, 64);
    content_len             = 64;
    static const char srv[] = "TLS 1.3, server CertificateVerify";
    static const char cli[] = "TLS 1.3, client CertificateVerify";
    const char *label       = is_server ? srv : cli;
    size_t label_len        = 33; /* both labels are 33 bytes */
    memcpy(content + content_len, label, label_len);
    content_len += label_len;
    content[content_len++] = 0x00;

    uint8_t transcript_hash[48];
    litls_transcript_hash(&ctx->transcript, transcript_hash);
    int hash_len = ctx->suite.hash_len;
    memcpy(content + content_len, transcript_hash, (size_t)hash_len);
    content_len += (size_t)hash_len;

    /* Select sig_algo by intersecting identity key type with peer's supported
     * algos. RFC 8446 §4.4.2.2: the server MUST sign with an algorithm the
     * client advertised; there is no "best effort" fallback. */
    uint16_t sig_algo = 0;
    uint16_t wanted   = 0;
    if (id->key_type == LITLS_KEY_ECDSA_P256)
        wanted = TLS13_SIG_ECDSA_SECP256R1_SHA256;
    else if (id->key_type == LITLS_KEY_ED25519)
        wanted = TLS13_SIG_ED25519;

    if (!wanted) {
        snprintf(ctx->error, sizeof(ctx->error), "unsupported identity key type for signing");
        return -1;
    }

    for (int i = 0; i < ctx->peer_sig_algo_count; i++) {
        if (ctx->peer_sig_algos[i] == wanted) {
            sig_algo = wanted;
            break;
        }
    }
    if (!sig_algo) {
        snprintf(ctx->error, sizeof(ctx->error), "no signature algorithm in common with peer");
        return -1;
    }

    /* Sign */
    uint8_t sig_buf[128];
    size_t sig_len = 0;

    if (id->key_type == LITLS_KEY_ECDSA_P256) {
        uint8_t hash[32];
        litls_sha256(content, content_len, hash);
        uint8_t raw_sig[64];
        size_t raw_sig_len;
        if (litls_p256_ecdsa_sign(hash, 32, id->private_key, raw_sig, &raw_sig_len) != 0) {
            snprintf(ctx->error, sizeof(ctx->error), "ECDSA sign failed");
            return -1;
        }
        /* Convert raw r||s to DER for TLS 1.3 */
        if (litls_ecdsa_sig_raw_to_der(raw_sig, raw_sig_len, sig_buf, sizeof(sig_buf), &sig_len) != 0) {
            snprintf(ctx->error, sizeof(ctx->error), "ECDSA DER conversion failed");
            return -1;
        }
    } else if (id->key_type == LITLS_KEY_ED25519) {
        /* Ed25519 signs the content directly */
        if (litls_ed25519_sign(sig_buf, content, content_len, id->public_key, id->private_key) != 0) {
            snprintf(ctx->error, sizeof(ctx->error), "Ed25519 sign failed");
            return -1;
        }
        sig_len = 64;
    } else {
        snprintf(ctx->error, sizeof(ctx->error), "unsupported key type for signing");
        return -1;
    }

    /* Build CertificateVerify: header(4) + sig_algo(2) + sig_len(2) + sig */
    size_t body_len = 2 + 2 + sig_len;
    size_t total    = 4 + body_len;
    if (total > cap)
        return -1;

    uint8_t *p = out;
    *p++       = TLS13_HT_CERTIFICATE_VERIFY;
    put24(p, (uint32_t)body_len);
    p += 3;
    put16(p, sig_algo);
    p += 2;
    put16(p, (uint16_t)sig_len);
    p += 2;
    memcpy(p, sig_buf, sig_len);

    *out_len = total;
    return 0;
}