/*
 * demo_server.c — deliberately vulnerable key-value store server
 *                  for CS496/596 in-class fuzzing demonstration
 *
 * This server implements a different protocol and different vulnerability
 * classes than vuln_server.c.  Use it for live class demos without
 * spoiling the homework target.
 *
 * PROTOCOL (TCP port 4141, line-oriented, CRLF terminated)
 * ─────────────────────────────────────────────────────────
 *   SET <key> <value>   →  +OK\r\n
 *   GET <key>           →  +<value>\r\n  |  -ERR not found\r\n
 *   DEL <key>           →  +OK\r\n       |  -ERR not found\r\n
 *   KEYS                →  +<key1>,<key2>,...\r\n
 *   LEN <key>           →  +<length>\r\n  |  -ERR not found\r\n
 *   APPEND <key> <data> →  +OK\r\n
 *   RESET               →  +OK\r\n  (clear the store)
 *   INFO                →  +VulnStore/1.0 <uptime>s <count> keys\r\n
 *   QUIT                →  +BYE\r\n
 *
 * INTENTIONAL VULNERABILITIES
 * ────────────────────────────
 *   1. SET handler — stack buffer overflow
 *      value_buf[64] on the stack; strcpy'd from the line buffer.
 *      Sending SET k <64+ bytes> overwrites the return address of cmd_set().
 *
 *   2. LEN handler — signed/unsigned integer confusion → out-of-bounds read
 *      atoi() returns int; compared against size_t.  A negative result
 *      (empty key, non-numeric key stored as value) wraps on the comparison
 *      and the code reads past the end of the values[] array.
 *      More subtle than a simple overflow — easy to miss in code review.
 *
 *   3. APPEND handler — integer truncation → heap overflow
 *      existing_len + new_len is computed as int, truncating to 0 when the
 *      sum overflows INT_MAX.  The realloc() call uses the truncated size,
 *      allocating far less memory than the subsequent memcpy writes.
 *
 * COMPILE:
 *   gcc -o demo_server demo_server.c -fno-stack-protector -no-pie -g
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>
#include <arpa/inet.h>
#include <sys/socket.h>

#define PORT       4141
#define BACKLOG    5
#define MAX_KEYS   32
#define KEY_LEN    64
#define VAL_LEN    256

/* ── in-memory key-value store ──────────────────────────────────────────── */
static char  keys[MAX_KEYS][KEY_LEN];
static char *vals[MAX_KEYS];       /* heap-allocated value strings */
static int   nkeys = 0;
static time_t start_time;

static int find_key(const char *k)
{
    for (int i = 0; i < nkeys; i++)
        if (strcmp(keys[i], k) == 0) return i;
    return -1;
}

/* ── vulnerability 1: stack buffer overflow ─────────────────────────────── */
static void cmd_set(int sock, const char *key, const char *value)
{
    char value_buf[64];      /* intentionally small */

    int idx = find_key(key);
    if (idx < 0) {
        if (nkeys >= MAX_KEYS) {
            dprintf(sock, "-ERR store full\r\n");
            return;
        }
        idx = nkeys++;
        strncpy(keys[idx], key, KEY_LEN - 1);
        keys[idx][KEY_LEN - 1] = '\0';
        vals[idx] = NULL;
    }

    strcpy(value_buf, value);    /* ← CWE-121: value_buf is only 64 bytes */

    free(vals[idx]);
    vals[idx] = strdup(value_buf);
    dprintf(sock, "+OK\r\n");
}

/* ── vulnerability 2: signed/unsigned confusion → OOB read ──────────────── */
static void cmd_len(int sock, const char *key)
{
    int idx = find_key(key);
    if (idx < 0) {
        dprintf(sock, "-ERR not found\r\n");
        return;
    }

    /*
     * atoi() returns a signed int.  If the stored value happens to be a
     * numeric string (e.g., "-1"), assigning it to a size_t causes the
     * negative value to wrap to SIZE_MAX.  The subsequent use of `len`
     * as an array index or a length in snprintf can read far outside
     * the valid range.
     *
     * In this code the "vulnerability" is latent: it manifests when a
     * caller does `GET k` and uses the returned length to size a buffer,
     * then finds the actual string is longer.
     */
    int len = atoi(vals[idx]);   /* ← CWE-195: signed-to-unsigned confusion */
    if ((size_t)len > VAL_LEN) {
        /* len was negative → cast to size_t is huge → branch taken */
        dprintf(sock, "-ERR internal length error\r\n");
        return;
    }
    dprintf(sock, "+%d\r\n", (int)strlen(vals[idx]));
}

/* ── vulnerability 3: integer truncation → heap overflow ────────────────── */
static void cmd_append(int sock, const char *key, const char *data)
{
    int idx = find_key(key);
    if (idx < 0) {
        /* key doesn't exist — create it */
        if (nkeys >= MAX_KEYS) {
            dprintf(sock, "-ERR store full\r\n");
            return;
        }
        idx = nkeys++;
        strncpy(keys[idx], key, KEY_LEN - 1);
        keys[idx][KEY_LEN - 1] = '\0';
        vals[idx] = strdup(data);
        dprintf(sock, "+OK\r\n");
        return;
    }

    int existing_len = (int)strlen(vals[idx]);
    int new_len      = (int)strlen(data);

    /*
     * If data is very long, existing_len + new_len + 1 can overflow int,
     * wrapping to a small positive number.  realloc() allocates that small
     * size; the subsequent strcat writes far beyond the allocation.
     *
     * CWE-190 (integer overflow) → CWE-122 (heap buffer overflow).
     * Triggerable by sending APPEND k <2GB string> — boofuzz won't send
     * 2GB in one test case, but sending repeated APPEND calls against
     * the same key accumulates the total length.
     */
    int total = existing_len + new_len + 1;  /* ← can overflow if new_len huge */

    char *newbuf = realloc(vals[idx], (size_t)total);
    if (!newbuf) {
        dprintf(sock, "-ERR out of memory\r\n");
        return;
    }
    vals[idx] = newbuf;
    strcat(vals[idx], data);    /* ← writes beyond allocation if total wrapped */
    dprintf(sock, "+OK\r\n");
}

/* ── safe handlers ──────────────────────────────────────────────────────── */
static void cmd_get(int sock, const char *key)
{
    int idx = find_key(key);
    if (idx < 0) dprintf(sock, "-ERR not found\r\n");
    else          dprintf(sock, "+%s\r\n", vals[idx] ? vals[idx] : "");
}

static void cmd_del(int sock, const char *key)
{
    int idx = find_key(key);
    if (idx < 0) { dprintf(sock, "-ERR not found\r\n"); return; }
    free(vals[idx]);
    /* swap with last entry to keep array compact */
    if (idx < nkeys - 1) {
        strncpy(keys[idx], keys[nkeys - 1], KEY_LEN);
        vals[idx] = vals[nkeys - 1];
    }
    nkeys--;
    dprintf(sock, "+OK\r\n");
}

static void cmd_keys(int sock)
{
    if (nkeys == 0) { dprintf(sock, "+\r\n"); return; }
    dprintf(sock, "+");
    for (int i = 0; i < nkeys; i++) {
        dprintf(sock, "%s%s", keys[i], i < nkeys - 1 ? "," : "");
    }
    dprintf(sock, "\r\n");
}

static void cmd_reset(int sock)
{
    for (int i = 0; i < nkeys; i++) { free(vals[i]); vals[i] = NULL; }
    nkeys = 0;
    dprintf(sock, "+OK\r\n");
}

static void cmd_info(int sock)
{
    time_t uptime = time(NULL) - start_time;
    dprintf(sock, "+VulnStore/1.0 %lds %d keys\r\n", (long)uptime, nkeys);
}

/* ── dispatch ───────────────────────────────────────────────────────────── */
static void handle_client(int sock)
{
    char line[512];
    ssize_t n;

    dprintf(sock, "+VulnStore 1.0 (demo target — port 4141)\r\n");

    while ((n = recv(sock, line, sizeof(line) - 1, 0)) > 0) {
        line[n] = '\0';
        char *end = strpbrk(line, "\r\n");
        if (end) *end = '\0';

        /* tokenise: cmd [arg1 [arg2]] */
        char *cmd  = strtok(line,  " ");
        char *arg1 = strtok(NULL,  " ");
        char *arg2 = strtok(NULL,  "");   /* rest of line */
        if (!cmd || !*cmd) continue;

        if      (strcasecmp(cmd, "SET")    == 0 && arg1 && arg2)
            cmd_set(sock, arg1, arg2);
        else if (strcasecmp(cmd, "GET")    == 0 && arg1)
            cmd_get(sock, arg1);
        else if (strcasecmp(cmd, "DEL")    == 0 && arg1)
            cmd_del(sock, arg1);
        else if (strcasecmp(cmd, "LEN")    == 0 && arg1)
            cmd_len(sock, arg1);
        else if (strcasecmp(cmd, "APPEND") == 0 && arg1 && arg2)
            cmd_append(sock, arg1, arg2);
        else if (strcasecmp(cmd, "KEYS")   == 0)
            cmd_keys(sock);
        else if (strcasecmp(cmd, "RESET")  == 0)
            cmd_reset(sock);
        else if (strcasecmp(cmd, "INFO")   == 0)
            cmd_info(sock);
        else if (strcasecmp(cmd, "QUIT")   == 0) {
            dprintf(sock, "+BYE\r\n");
            break;
        } else {
            dprintf(sock, "-ERR unknown command: %.50s\r\n", cmd);
        }
    }
    close(sock);
}

int main(void)
{
    int srv, cli, opt = 1;
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));

    start_time = time(NULL);

    if ((srv = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
        perror("socket"); return 1;
    }
    setsockopt(srv, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    addr.sin_family      = AF_INET;
    addr.sin_addr.s_addr = INADDR_ANY;
    addr.sin_port        = htons(PORT);

    if (bind(srv, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        perror("bind"); return 1;
    }
    listen(srv, BACKLOG);

    printf("VulnStore listening on port %d\n", PORT);
    printf("Vulnerabilities: SET(stack overflow) LEN(int confusion) APPEND(int truncation)\n");

    for (;;) {
        cli = accept(srv, NULL, NULL);
        if (cli >= 0) handle_client(cli);
    }
}