// 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 Key Schedule (RFC 8446 Section 7.1)
 *
 * Pure computation, no I/O. Implements HKDF-Expand-Label, Derive-Secret,
 * and the full handshake/application key derivation.
 */

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

/*
 * HKDF-Expand-Label(Secret, Label, Context, Length) =
 *   HKDF-Expand(Secret, HkdfLabel, Length)
 *
 * HkdfLabel = struct {
 *   uint16 length;
 *   opaque label<7..255> = "tls13 " + Label;
 *   opaque context<0..255> = Context;
 * };
 */

int litls_tls13_expand_label_256(const uint8_t secret[32], const char *label, const uint8_t *context,
                                 size_t context_len, uint8_t *out, size_t out_len) {
    /* Build HkdfLabel info */
    size_t label_len       = strlen(label);
    size_t prefix_len      = 6; /* "tls13 " */
    size_t total_label_len = prefix_len + label_len;

    /* HkdfLabel: 2 (length) + 1 (label_len_byte) + total_label_len + 1 (context_len_byte) + context_len */
    uint8_t info[512];
    if (4 + total_label_len + context_len > sizeof(info))
        return -1;
    size_t pos = 0;

    info[pos++] = (uint8_t)(out_len >> 8);
    info[pos++] = (uint8_t)(out_len);
    info[pos++] = (uint8_t)(total_label_len);
    memcpy(info + pos, "tls13 ", 6);
    pos += 6;
    memcpy(info + pos, label, label_len);
    pos += label_len;
    info[pos++] = (uint8_t)(context_len);
    if (context_len > 0) {
        memcpy(info + pos, context, context_len);
        pos += context_len;
    }

    return litls_hkdf_expand_sha256(secret, info, pos, out, out_len);
}

int litls_tls13_expand_label_384(const uint8_t secret[48], const char *label, const uint8_t *context,
                                 size_t context_len, uint8_t *out, size_t out_len) {
    size_t label_len       = strlen(label);
    size_t prefix_len      = 6;
    size_t total_label_len = prefix_len + label_len;

    uint8_t info[512];
    if (4 + total_label_len + context_len > sizeof(info))
        return -1;
    size_t pos = 0;

    info[pos++] = (uint8_t)(out_len >> 8);
    info[pos++] = (uint8_t)(out_len);
    info[pos++] = (uint8_t)(total_label_len);
    memcpy(info + pos, "tls13 ", 6);
    pos += 6;
    memcpy(info + pos, label, label_len);
    pos += label_len;
    info[pos++] = (uint8_t)(context_len);
    if (context_len > 0) {
        memcpy(info + pos, context, context_len);
        pos += context_len;
    }

    return litls_hkdf_expand_sha384(secret, info, pos, out, out_len);
}

int litls_tls13_derive_secret_256(const uint8_t secret[32], const char *label, const uint8_t transcript_hash[32],
                                  uint8_t out[32]) {
    return litls_tls13_expand_label_256(secret, label, transcript_hash, 32, out, 32);
}

int litls_tls13_derive_secret_384(const uint8_t secret[48], const char *label, const uint8_t transcript_hash[48],
                                  uint8_t out[48]) {
    return litls_tls13_expand_label_384(secret, label, transcript_hash, 48, out, 48);
}

/* Extract traffic key and IV from a traffic secret */
void litls_tls13_traffic_keys_256(const uint8_t traffic_secret[32], uint8_t *key, size_t key_len, uint8_t iv[12]) {
    litls_tls13_expand_label_256(traffic_secret, "key", NULL, 0, key, key_len);
    litls_tls13_expand_label_256(traffic_secret, "iv", NULL, 0, iv, 12);
}

void litls_tls13_traffic_keys_384(const uint8_t traffic_secret[48], uint8_t *key, size_t key_len, uint8_t iv[12]) {
    litls_tls13_expand_label_384(traffic_secret, "key", NULL, 0, key, key_len);
    litls_tls13_expand_label_384(traffic_secret, "iv", NULL, 0, iv, 12);
}

/*
 * Derive handshake keys after ServerHello.
 *
 * Input: ctx->shared_secret (from X25519), transcript hash up to ServerHello
 * Output: ctx->hs_read, ctx->hs_write keys, ctx->server_hs_secret, ctx->client_hs_secret
 */
int litls_tls13_derive_handshake_keys(litls_tls13_ctx *ctx) {
    int hash_len = ctx->suite.hash_len;
    int key_len  = ctx->suite.key_len;
    uint8_t transcript_hash[48];
    litls_transcript_hash(&ctx->transcript, transcript_hash);

    /* Zero IKM for early secret */
    uint8_t zeros[48];
    memset(zeros, 0, (size_t)hash_len);

    if (hash_len == 32) {
        /* Early Secret = HKDF-Extract(0, 0) */
        uint8_t early_secret[32];
        litls_hkdf_extract_sha256(NULL, 0, zeros, 32, early_secret);

        /* Derive-Secret(early_secret, "derived", Hash("")) */
        uint8_t empty_hash[32];
        litls_sha256((const uint8_t *)"", 0, empty_hash);
        uint8_t derived[32];
        litls_tls13_derive_secret_256(early_secret, "derived", empty_hash, derived);

        /* Handshake Secret = HKDF-Extract(derived, shared_secret) */
        uint8_t handshake_secret[32];
        litls_hkdf_extract_sha256(derived, 32, ctx->shared_secret, 32, handshake_secret);

        /* client_handshake_traffic_secret */
        litls_tls13_derive_secret_256(handshake_secret, "c hs traffic", transcript_hash, ctx->client_hs_secret);

        /* server_handshake_traffic_secret */
        litls_tls13_derive_secret_256(handshake_secret, "s hs traffic", transcript_hash, ctx->server_hs_secret);

        /* Derive traffic keys: swap read/write based on role */
        const uint8_t *read_secret  = ctx->is_server ? ctx->client_hs_secret : ctx->server_hs_secret;
        const uint8_t *write_secret = ctx->is_server ? ctx->server_hs_secret : ctx->client_hs_secret;
        litls_tls13_traffic_keys_256(read_secret, ctx->hs_read.key, (size_t)key_len, ctx->hs_read.iv);
        litls_tls13_traffic_keys_256(write_secret, ctx->hs_write.key, (size_t)key_len, ctx->hs_write.iv);
        ctx->hs_read.seq  = 0;
        ctx->hs_write.seq = 0;

        /* Store for later: Derive-Secret(handshake_secret, "derived", Hash("")) for master secret */
        litls_tls13_derive_secret_256(handshake_secret, "derived", empty_hash, derived);
        /* master_secret = HKDF-Extract(derived, 0) */
        litls_hkdf_extract_sha256(derived, 32, zeros, 32, ctx->master_secret);

        litls_secure_zero(early_secret, sizeof(early_secret));
        litls_secure_zero(derived, sizeof(derived));
        litls_secure_zero(handshake_secret, sizeof(handshake_secret));
    } else {
        /* SHA-384 path */
        uint8_t early_secret[48];
        litls_hkdf_extract_sha384(NULL, 0, zeros, 48, early_secret);

        uint8_t empty_hash[48];
        litls_sha384((const uint8_t *)"", 0, empty_hash);
        uint8_t derived[48];
        litls_tls13_derive_secret_384(early_secret, "derived", empty_hash, derived);

        uint8_t handshake_secret[48];
        litls_hkdf_extract_sha384(derived, 48, ctx->shared_secret, 32, handshake_secret);

        litls_tls13_derive_secret_384(handshake_secret, "c hs traffic", transcript_hash, ctx->client_hs_secret);
        litls_tls13_derive_secret_384(handshake_secret, "s hs traffic", transcript_hash, ctx->server_hs_secret);

        const uint8_t *read_secret  = ctx->is_server ? ctx->client_hs_secret : ctx->server_hs_secret;
        const uint8_t *write_secret = ctx->is_server ? ctx->server_hs_secret : ctx->client_hs_secret;
        litls_tls13_traffic_keys_384(read_secret, ctx->hs_read.key, (size_t)key_len, ctx->hs_read.iv);
        litls_tls13_traffic_keys_384(write_secret, ctx->hs_write.key, (size_t)key_len, ctx->hs_write.iv);
        ctx->hs_read.seq  = 0;
        ctx->hs_write.seq = 0;

        litls_tls13_derive_secret_384(handshake_secret, "derived", empty_hash, derived);
        litls_hkdf_extract_sha384(derived, 48, zeros, 48, ctx->master_secret);

        litls_secure_zero(early_secret, sizeof(early_secret));
        litls_secure_zero(derived, sizeof(derived));
        litls_secure_zero(handshake_secret, sizeof(handshake_secret));
    }

    return 0;
}

/*
 * Derive application keys after server Finished.
 *
 * Input: ctx->master_secret (computed during handshake key derivation),
 *        transcript hash through server Finished
 * Output: ctx->app_read, ctx->app_write keys
 */
int litls_tls13_derive_app_keys(litls_tls13_ctx *ctx) {
    int hash_len = ctx->suite.hash_len;
    int key_len  = ctx->suite.key_len;
    uint8_t transcript_hash[48];
    litls_transcript_hash(&ctx->transcript, transcript_hash);

    if (hash_len == 32) {
        uint8_t client_app_secret[32], server_app_secret[32];
        litls_tls13_derive_secret_256(ctx->master_secret, "c ap traffic", transcript_hash, client_app_secret);
        litls_tls13_derive_secret_256(ctx->master_secret, "s ap traffic", transcript_hash, server_app_secret);

        const uint8_t *read_secret  = ctx->is_server ? client_app_secret : server_app_secret;
        const uint8_t *write_secret = ctx->is_server ? server_app_secret : client_app_secret;
        litls_tls13_traffic_keys_256(read_secret, ctx->app_read.key, (size_t)key_len, ctx->app_read.iv);
        litls_tls13_traffic_keys_256(write_secret, ctx->app_write.key, (size_t)key_len, ctx->app_write.iv);

        litls_secure_zero(client_app_secret, sizeof(client_app_secret));
        litls_secure_zero(server_app_secret, sizeof(server_app_secret));
    } else {
        uint8_t client_app_secret[48], server_app_secret[48];
        litls_tls13_derive_secret_384(ctx->master_secret, "c ap traffic", transcript_hash, client_app_secret);
        litls_tls13_derive_secret_384(ctx->master_secret, "s ap traffic", transcript_hash, server_app_secret);

        const uint8_t *read_secret  = ctx->is_server ? client_app_secret : server_app_secret;
        const uint8_t *write_secret = ctx->is_server ? server_app_secret : client_app_secret;
        litls_tls13_traffic_keys_384(read_secret, ctx->app_read.key, (size_t)key_len, ctx->app_read.iv);
        litls_tls13_traffic_keys_384(write_secret, ctx->app_write.key, (size_t)key_len, ctx->app_write.iv);

        litls_secure_zero(client_app_secret, sizeof(client_app_secret));
        litls_secure_zero(server_app_secret, sizeof(server_app_secret));
    }

    ctx->app_read.seq  = 0;
    ctx->app_write.seq = 0;
    return 0;
}