/*
 * fwupd.c — a small firmware-update receiver for a fictional sensor board.
 *
 * Original code written for h4rithd.com. It is compiled for arm64 and its real
 * disassembly is the background of the site. fw_write_chunk() contains a
 * deliberate signedness bug (CWE-195): a 32-bit length read off the wire is
 * stored in a signed int, so a "negative" length passes the bounds check and
 * reaches memcpy() sign-extended to a huge size_t.
 *
 *   clang -arch arm64 -O2 -fno-inline -o fwupd fwupd.c            (vulnerable)
 *   clang -arch arm64 -O2 -fno-inline -DHARDENED -o fwupd fwupd.c (fixed)
 *
 * Do not ship this. It exists to be taken apart.
 */
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>

#define NOINLINE __attribute__((noinline))

#define FW_MAGIC        0x48344657u /* "H4FW" */
#define FW_CHUNK_MAX    0x400
#define FW_SLOT_SIZE    0x10000
#define FW_MAX_SESSIONS 8
#define RING_SIZE       64

enum {
    OP_HELLO   = 0x01,
    OP_BEGIN   = 0x02,
    OP_CHUNK   = 0x03,
    OP_VERIFY  = 0x04,
    OP_COMMIT  = 0x05,
    OP_ABORT   = 0x06,
    OP_STATUS  = 0x07,
    OP_SETKEY  = 0x08,
};

enum {
    ERR_OK = 0,
    ERR_SHORT = -1,
    ERR_MAGIC = -2,
    ERR_STATE = -3,
    ERR_RANGE = -4,
    ERR_CRC = -5,
    ERR_AUTH = -6,
    ERR_OP = -7,
};

struct fw_session {
    uint32_t id;
    uint32_t state;
    uint32_t crc;
    uint32_t expected_crc;
    uint64_t offset;
    uint64_t image_size;
    uint8_t  key[16];
};

struct fw_ctx {
    uint8_t  chunk[FW_CHUNK_MAX];
    struct fw_session *session;
    uint8_t  (*slot)[FW_SLOT_SIZE];
    uint32_t chunks_seen;
    uint32_t flags;
};

struct ring {
    uint32_t head;
    uint32_t tail;
    uint8_t  data[RING_SIZE];
};

static struct fw_session g_sessions[FW_MAX_SESSIONS];
static uint8_t g_slot[FW_SLOT_SIZE];
static struct ring g_log;

/* ---- byte order helpers ------------------------------------------------ */

NOINLINE uint16_t read_u16_le(const uint8_t *p) { return (uint16_t)(p[0] | p[1] << 8); }
NOINLINE uint16_t read_u16_be(const uint8_t *p) { return (uint16_t)(p[0] << 8 | p[1]); }

NOINLINE uint32_t read_u32_be(const uint8_t *p)
{
    return (uint32_t)p[0] << 24 | (uint32_t)p[1] << 16 | (uint32_t)p[2] << 8 | p[3];
}

NOINLINE uint32_t read_u32_le(const uint8_t *p)
{
    return (uint32_t)p[3] << 24 | (uint32_t)p[2] << 16 | (uint32_t)p[1] << 8 | p[0];
}

NOINLINE void write_u32_be(uint8_t *p, uint32_t v)
{
    p[0] = (uint8_t)(v >> 24);
    p[1] = (uint8_t)(v >> 16);
    p[2] = (uint8_t)(v >> 8);
    p[3] = (uint8_t)v;
}

/* ---- checksums ---------------------------------------------------------- */

NOINLINE uint32_t crc32_update(uint32_t crc, const uint8_t *p, size_t n)
{
    crc = ~crc;
    while (n--) {
        crc ^= *p++;
        for (int k = 0; k < 8; k++)
            crc = (crc >> 1) ^ (0xEDB88320u & (0u - (crc & 1u)));
    }
    return ~crc;
}

NOINLINE uint32_t adler32(const uint8_t *p, size_t n)
{
    uint32_t a = 1, b = 0;
    while (n) {
        size_t block = n < 5552 ? n : 5552;
        n -= block;
        while (block--) {
            a += *p++;
            b += a;
        }
        a %= 65521u;
        b %= 65521u;
    }
    return b << 16 | a;
}

NOINLINE uint32_t fnv1a32(const uint8_t *p, size_t n)
{
    uint32_t h = 0x811C9DC5u;
    for (size_t i = 0; i < n; i++) {
        h ^= p[i];
        h *= 0x01000193u;
    }
    return h;
}

/* ---- decoding ----------------------------------------------------------- */

NOINLINE int varint_decode(const uint8_t *p, size_t n, uint64_t *out)
{
    uint64_t v = 0;
    for (size_t i = 0; i < n && i < 10; i++) {
        v |= (uint64_t)(p[i] & 0x7F) << (7 * i);
        if (!(p[i] & 0x80)) {
            *out = v;
            return (int)i + 1;
        }
    }
    return ERR_SHORT;
}

NOINLINE int hex_nibble(int c)
{
    if (c >= '0' && c <= '9') return c - '0';
    if (c >= 'a' && c <= 'f') return c - 'a' + 10;
    if (c >= 'A' && c <= 'F') return c - 'A' + 10;
    return -1;
}

NOINLINE int hex_decode(uint8_t *dst, size_t cap, const char *src, size_t n)
{
    if (n & 1 || n / 2 > cap) return ERR_RANGE;
    for (size_t i = 0; i < n; i += 2) {
        int hi = hex_nibble(src[i]), lo = hex_nibble(src[i + 1]);
        if (hi < 0 || lo < 0) return ERR_RANGE;
        dst[i / 2] = (uint8_t)(hi << 4 | lo);
    }
    return (int)(n / 2);
}

NOINLINE int utf8_valid(const uint8_t *p, size_t n)
{
    size_t i = 0;
    while (i < n) {
        uint8_t c = p[i];
        size_t need = c < 0x80 ? 0 : (c >> 5) == 0x6 ? 1 : (c >> 4) == 0xE ? 2 : (c >> 3) == 0x1E ? 3 : 4;
        if (need == 4 || i + need >= n + (need == 0)) return need == 0;
        for (size_t k = 1; k <= need; k++)
            if ((p[i + k] & 0xC0) != 0x80) return 0;
        i += need + 1;
    }
    return 1;
}

NOINLINE int tlv_next(const uint8_t *p, size_t n, uint8_t *type, const uint8_t **val, uint16_t *len)
{
    if (n < 3) return ERR_SHORT;
    *type = p[0];
    *len = read_u16_be(p + 1);
    if ((size_t)*len + 3 > n) return ERR_RANGE;
    *val = p + 3;
    return *len + 3;
}

/* ---- comparisons -------------------------------------------------------- */

NOINLINE int ct_equal(const uint8_t *a, const uint8_t *b, size_t n)
{
    uint8_t d = 0;
    for (size_t i = 0; i < n; i++) d |= a[i] ^ b[i];
    return d == 0;
}

NOINLINE int checked_add_u64(uint64_t a, uint64_t b, uint64_t *out)
{
    return !__builtin_add_overflow(a, b, out);
}

NOINLINE size_t str_copy(char *dst, const char *src, size_t cap)
{
    size_t i = 0;
    if (!cap) return 0;
    for (; i + 1 < cap && src[i]; i++) dst[i] = src[i];
    dst[i] = '\0';
    return i;
}

/* ---- log ring ----------------------------------------------------------- */

NOINLINE int ring_push(struct ring *r, uint8_t b)
{
    uint32_t next = (r->head + 1) % RING_SIZE;
    if (next == r->tail) return 0;
    r->data[r->head] = b;
    r->head = next;
    return 1;
}

NOINLINE int ring_pop(struct ring *r, uint8_t *b)
{
    if (r->head == r->tail) return 0;
    *b = r->data[r->tail];
    r->tail = (r->tail + 1) % RING_SIZE;
    return 1;
}

/* ---- sessions ----------------------------------------------------------- */

NOINLINE struct fw_session *session_find(uint32_t id)
{
    for (int i = 0; i < FW_MAX_SESSIONS; i++)
        if (g_sessions[i].id == id && g_sessions[i].state) return &g_sessions[i];
    return NULL;
}

NOINLINE struct fw_session *session_open(uint32_t id, uint64_t image_size)
{
    if (!id || image_size > FW_SLOT_SIZE) return NULL;
    for (int i = 0; i < FW_MAX_SESSIONS; i++) {
        if (!g_sessions[i].state) {
            g_sessions[i].id = id;
            g_sessions[i].state = 1;
            g_sessions[i].offset = 0;
            g_sessions[i].crc = 0;
            g_sessions[i].image_size = image_size;
            return &g_sessions[i];
        }
    }
    return NULL;
}

NOINLINE void session_close(struct fw_session *s)
{
    memset(s, 0, sizeof *s);
}

/* ---- the interesting part ---------------------------------------------- */

/*
 * CHUNK packet: [op:1][session:4][len:4][payload:len]
 *
 * len arrives big-endian from the wire. Stored as int32_t, a value such as
 * 0xFFFFFFF0 becomes -16: it is not "greater than" FW_CHUNK_MAX, so the check
 * passes, and memcpy receives it sign-extended to 0xFFFFFFFFFFFFFFF0.
 */
NOINLINE int fw_write_chunk(struct fw_ctx *ctx, const uint8_t *pkt, size_t n)
{
#ifdef HARDENED
    uint32_t len = read_u32_be(pkt + 5);
#else
    int32_t len = (int32_t)read_u32_be(pkt + 5);
#endif
    if (len > FW_CHUNK_MAX)
        return ERR_RANGE;

    memcpy(ctx->chunk, pkt + 9, len);

    struct fw_session *s = ctx->session;
    s->crc = crc32_update(s->crc, ctx->chunk, (size_t)len);
    s->offset += (size_t)len;
    ctx->chunks_seen++;
    (void)n;
    return ERR_OK;
}

NOINLINE int fw_verify(struct fw_ctx *ctx)
{
    struct fw_session *s = ctx->session;
    if (!s || s->state != 1) return ERR_STATE;
    if (s->offset != s->image_size) return ERR_RANGE;
    return s->crc == s->expected_crc ? ERR_OK : ERR_CRC;
}

NOINLINE int fw_commit(struct fw_ctx *ctx)
{
    int rc = fw_verify(ctx);
    if (rc) return rc;
    ctx->session->state = 2;
    ring_push(&g_log, OP_COMMIT);
    return ERR_OK;
}

NOINLINE int fw_set_key(struct fw_ctx *ctx, const uint8_t *pkt, size_t n)
{
    uint8_t type;
    const uint8_t *val;
    uint16_t len;
    if (!ctx->session) return ERR_STATE;
    int used = tlv_next(pkt + 5, n - 5, &type, &val, &len);
    if (used < 0) return used;
    if (type != 0x4B || len != sizeof ctx->session->key) return ERR_AUTH;
    memcpy(ctx->session->key, val, sizeof ctx->session->key);
    return ERR_OK;
}

NOINLINE int fw_status(const struct fw_ctx *ctx, uint8_t *out, size_t cap)
{
    if (cap < 12) return ERR_SHORT;
    const struct fw_session *s = ctx->session;
    write_u32_be(out, s ? s->id : 0);
    write_u32_be(out + 4, s ? (uint32_t)s->offset : 0);
    write_u32_be(out + 8, ctx->chunks_seen);
    return 12;
}

NOINLINE int handle_packet(struct fw_ctx *ctx, const uint8_t *pkt, size_t n)
{
    if (n < 9) return ERR_SHORT;
    uint32_t sid = read_u32_be(pkt + 1);

    switch (pkt[0]) {
    case OP_HELLO:
        return read_u32_le(pkt + 1) == FW_MAGIC ? ERR_OK : ERR_MAGIC;
    case OP_BEGIN:
        ctx->session = session_open(sid, read_u32_be(pkt + 5));
        return ctx->session ? ERR_OK : ERR_STATE;
    case OP_CHUNK:
        ctx->session = session_find(sid);
        return ctx->session ? fw_write_chunk(ctx, pkt, n) : ERR_STATE;
    case OP_VERIFY:
        return fw_verify(ctx);
    case OP_COMMIT:
        return fw_commit(ctx);
    case OP_ABORT:
        if (ctx->session) session_close(ctx->session);
        ctx->session = NULL;
        return ERR_OK;
    case OP_STATUS:
        return fw_status(ctx, ctx->chunk, sizeof ctx->chunk);
    case OP_SETKEY:
        return fw_set_key(ctx, pkt, n);
    default:
        return ERR_OP;
    }
}

int main(void)
{
    static uint8_t pkt[0x2000];
    struct fw_ctx ctx = { .slot = &g_slot };
    uint8_t len_be[2];

    while (read(0, len_be, 2) == 2) {
        size_t n = read_u16_be(len_be);
        if (n > sizeof pkt) return 1;
        if (read(0, pkt, n) != (ssize_t)n) return 1;
        int rc = handle_packet(&ctx, pkt, n);
        uint8_t r = (uint8_t)rc;
        ring_push(&g_log, pkt[0]);
        write(1, &r, 1);
    }
    return 0;
}