// 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 Record Layer
 *
 * Handles record framing, AEAD encrypt/decrypt, and non-blocking I/O.
 * The record layer does NOT call read()/write() from inside the handshake
 * state machine. Instead, the caller drives fill/flush around handshake_step.
 */

#include "tls13.h"
#include <errno.h>
#include <stdio.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

/* ── Nonce construction ───────────────────────────────────────────── */

void litls_record_nonce(const litls_traffic_keys *keys, uint8_t nonce[12]) {
    memcpy(nonce, keys->iv, 12);
    /* XOR the 64-bit sequence number into the last 8 bytes */
    uint64_t seq = keys->seq;
    for (int i = 11; i >= 4; i--) {
        nonce[i] ^= (uint8_t)(seq & 0xFF);
        seq >>= 8;
    }
}

/* ── Non-blocking I/O ─────────────────────────────────────────────── */

litls_io_want litls_record_fill(litls_tls13_ctx *ctx) {
    /* Compact: shift unconsumed data to front */
    if (ctx->rbuf.pos > 0 && ctx->rbuf.len > ctx->rbuf.pos) {
        memmove(ctx->rbuf.data, ctx->rbuf.data + ctx->rbuf.pos, ctx->rbuf.len - ctx->rbuf.pos);
        ctx->rbuf.len -= ctx->rbuf.pos;
        ctx->rbuf.pos = 0;
    } else if (ctx->rbuf.pos > 0 && ctx->rbuf.len == ctx->rbuf.pos) {
        ctx->rbuf.len = 0;
        ctx->rbuf.pos = 0;
    }

    if (ctx->rbuf.len >= sizeof(ctx->rbuf.data))
        return LITLS_DONE;

    ssize_t n = read(ctx->fd, ctx->rbuf.data + ctx->rbuf.len, sizeof(ctx->rbuf.data) - ctx->rbuf.len);
    if (n > 0) {
        ctx->rbuf.len += (size_t)n;
        return LITLS_DONE;
    }
    if (n == 0)
        return LITLS_CLOSED;
    if (errno == EAGAIN || errno == EWOULDBLOCK)
        return LITLS_WANT_READ;

    snprintf(ctx->error, sizeof(ctx->error), "read: %s", strerror(errno));
    return LITLS_ERROR;
}

litls_io_want litls_record_flush(litls_tls13_ctx *ctx) {
    while (ctx->wbuf.pos < ctx->wbuf.len) {
        ssize_t n = send(ctx->fd, ctx->wbuf.data + ctx->wbuf.pos, ctx->wbuf.len - ctx->wbuf.pos, MSG_NOSIGNAL);
        if (n > 0) {
            ctx->wbuf.pos += (size_t)n;
            continue;
        }
        if (n < 0 && (errno == EAGAIN || errno == EWOULDBLOCK))
            return LITLS_WANT_WRITE;
        if (n < 0) {
            snprintf(ctx->error, sizeof(ctx->error), "write: %s", strerror(errno));
            return LITLS_ERROR;
        }
        /* send() returned 0 for non-zero length — treat as error */
        snprintf(ctx->error, sizeof(ctx->error), "write: send returned 0");
        return LITLS_ERROR;
    }
    /* All written, reset buffer */
    ctx->wbuf.len = 0;
    ctx->wbuf.pos = 0;
    return LITLS_DONE;
}

/* ── Record write ─────────────────────────────────────────────────── */

int litls_record_write(litls_tls13_ctx *ctx, uint8_t content_type, const uint8_t *data, size_t len) {
    /* RFC 8446 §5.1: TLSPlaintext.length must not exceed 2^14. Enforce at the
     * record layer so no caller can overflow the stack buffer below, even if a
     * future handshake builder or direct caller bypasses litls_tls13_write()
     * fragmentation. */
    if (len > TLS13_MAX_RECORD) {
        snprintf(ctx->error, sizeof(ctx->error), "record plaintext too large: %zu", len);
        return -1;
    }

    litls_traffic_keys *keys = NULL;

    if (ctx->hs_state == TLS13_HS_DONE)
        keys = &ctx->app_write;
    else if (ctx->encrypting)
        keys = &ctx->hs_write;

    if (keys) {
        /* Encrypted record: outer type is application_data (0x17) */
        /* inner plaintext = data || content_type (1 byte) */
        size_t inner_len = len + 1; /* data + real content type */
        size_t ct_len    = inner_len + TLS13_TAG_LEN;

        if (ctx->wbuf.len + TLS13_RECORD_HDR + ct_len > sizeof(ctx->wbuf.data)) {
            snprintf(ctx->error, sizeof(ctx->error), "write buffer overflow");
            return -1;
        }

        uint8_t *hdr = ctx->wbuf.data + ctx->wbuf.len;
        hdr[0]       = TLS13_CT_APPLICATION_DATA;
        hdr[1]       = 0x03;
        hdr[2]       = 0x03; /* legacy version TLS 1.2 */
        hdr[3]       = (uint8_t)(ct_len >> 8);
        hdr[4]       = (uint8_t)(ct_len);

        /* Build inner plaintext in a temporary buffer */
        uint8_t inner[TLS13_MAX_RECORD + 1];
        memcpy(inner, data, len);
        inner[len] = content_type;

        /* Build nonce */
        uint8_t nonce[12];
        litls_record_nonce(keys, nonce);

        /* AAD is the 5-byte record header */
        uint8_t *ct_out = hdr + TLS13_RECORD_HDR;
        uint8_t tag[16];

        int rc = ctx->suite.encrypt(keys->key, nonce, hdr, TLS13_RECORD_HDR, inner, inner_len, ct_out, tag);
        if (rc != 0) {
            snprintf(ctx->error, sizeof(ctx->error), "AEAD encrypt failed");
            return -1;
        }

        /* Append tag after ciphertext */
        memcpy(ct_out + inner_len, tag, TLS13_TAG_LEN);
        keys->seq++;

        ctx->wbuf.len += TLS13_RECORD_HDR + ct_len;
    } else {
        /* Plaintext record */
        if (ctx->wbuf.len + TLS13_RECORD_HDR + len > sizeof(ctx->wbuf.data)) {
            snprintf(ctx->error, sizeof(ctx->error), "write buffer overflow");
            return -1;
        }

        uint8_t *hdr = ctx->wbuf.data + ctx->wbuf.len;
        hdr[0]       = content_type;
        hdr[1]       = 0x03;
        hdr[2]       = 0x03; /* legacy record version TLS 1.2 (RFC 8446 §5.1) */
        hdr[3]       = (uint8_t)(len >> 8);
        hdr[4]       = (uint8_t)(len);
        memcpy(hdr + TLS13_RECORD_HDR, data, len);

        ctx->wbuf.len += TLS13_RECORD_HDR + len;
    }

    return 0;
}

/* ── Record read ──────────────────────────────────────────────────── */

int litls_record_read(litls_tls13_ctx *ctx, uint8_t *ct, uint8_t **out, size_t *out_len) {
    size_t avail = ctx->rbuf.len - ctx->rbuf.pos;
    uint8_t *buf = ctx->rbuf.data + ctx->rbuf.pos;

    /* Need at least the 5-byte header */
    if (avail < TLS13_RECORD_HDR)
        return 1; /* incomplete */

    uint8_t rec_type = buf[0];
    size_t rec_len   = ((size_t)buf[3] << 8) | buf[4];

    /* Sanity check */
    if (rec_len > TLS13_MAX_CIPHERTEXT + 256) {
        snprintf(ctx->error, sizeof(ctx->error), "record too large: %zu", rec_len);
        return -1;
    }

    /* Need full record */
    if (avail < TLS13_RECORD_HDR + rec_len)
        return 1; /* incomplete */

    uint8_t *rec_data = buf + TLS13_RECORD_HDR;

    /* Silently ignore ChangeCipherSpec (middlebox compat) */
    if (rec_type == TLS13_CT_CHANGE_CIPHER_SPEC) {
        ctx->rbuf.pos += TLS13_RECORD_HDR + rec_len;
        return 1; /* signal "try again" to caller */
    }

    /* Determine if this record should be decrypted */
    litls_traffic_keys *keys = NULL;
    if (rec_type == TLS13_CT_APPLICATION_DATA) {
        if (ctx->hs_state == TLS13_HS_DONE)
            keys = &ctx->app_read;
        else if (ctx->decrypting)
            keys = &ctx->hs_read;
    }

    /* RFC 8446 §5.1: once read keys are installed, every record other than
     * a tolerated ChangeCipherSpec (handled above) MUST arrive as a protected
     * application_data record. An unprotected handshake or alert record at
     * this point is fatal -- accepting it would let an off-path attacker inject
     * a cleartext close_notify (truncation attack) or corrupt the handshake
     * transcript with forged plaintext messages. */
    if ((ctx->decrypting || ctx->hs_state == TLS13_HS_DONE) && rec_type != TLS13_CT_APPLICATION_DATA) {
        snprintf(ctx->error, sizeof(ctx->error), "unprotected record (type %d) after key change", rec_type);
        return -1;
    }

    if (keys && rec_type == TLS13_CT_APPLICATION_DATA) {
        /* Encrypted record */
        if (rec_len < TLS13_TAG_LEN + 1) {
            snprintf(ctx->error, sizeof(ctx->error), "encrypted record too short");
            return -1;
        }

        size_t cipher_len = rec_len - TLS13_TAG_LEN;
        uint8_t *tag      = rec_data + cipher_len;

        uint8_t nonce[12];
        litls_record_nonce(keys, nonce);

        /* AAD = 5-byte record header */
        int rc =
            ctx->suite.decrypt(keys->key, nonce, buf, TLS13_RECORD_HDR, rec_data, cipher_len, tag, ctx->decrypt_buf);
        if (rc != 0) {
            snprintf(ctx->error, sizeof(ctx->error), "AEAD decrypt failed");
            return -1;
        }
        keys->seq++;

        /* Unmask content type: last non-zero byte of plaintext is the real type.
         * Scan all bytes in constant time to avoid leaking plaintext length. */
        size_t pt_len = 0;
        for (size_t i = 0; i < cipher_len; i++) {
            /* Update pt_len whenever we see a non-zero byte */
            size_t mask = (size_t)(-(ctx->decrypt_buf[i] != 0));
            pt_len      = (pt_len & ~mask) | ((i + 1) & mask);
        }

        if (pt_len == 0) {
            snprintf(ctx->error, sizeof(ctx->error), "empty inner plaintext");
            return -1;
        }

        *ct      = ctx->decrypt_buf[pt_len - 1]; /* real content type */
        *out     = ctx->decrypt_buf;
        *out_len = pt_len - 1;
    } else {
        /* Plaintext record */
        *ct      = rec_type;
        *out     = rec_data;
        *out_len = rec_len;
    }

    ctx->rbuf.pos += TLS13_RECORD_HDR + rec_len;

    /* Handle alerts */
    if (*ct == TLS13_CT_ALERT) {
        if (*out_len >= 2) {
            uint8_t level = (*out)[0];
            uint8_t desc  = (*out)[1];
            (void)level;
            if (desc == TLS13_ALERT_CLOSE_NOTIFY) {
                ctx->peer_closed = 1;
                return -2; /* special: closed */
            }
            snprintf(ctx->error, sizeof(ctx->error), "alert: level=%d desc=%d", level, desc);
        } else {
            snprintf(ctx->error, sizeof(ctx->error), "malformed alert record");
        }
        return -1;
    }

    return 0;
}