/*
 * mini_httpd.c — Vulnerable embedded-style HTTP daemon
 *
 * Practice target simulating a router/IoT web interface.
 * Single-threaded, persistent connections (HTTP/1.1 keep-alive).
 * Contains multiple heap vulnerabilities exploitable on
 * glibc <= 2.27 (pre-tcache-hardening) or uClibc.
 *
 * Compile: make
 * Run:     ./mini_httpd           (default port 8080)
 *          ./mini_httpd 9090      (custom port)
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <errno.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

/* ------------------------------------------------------------------ */
/*  Configuration                                                      */
/* ------------------------------------------------------------------ */

#define DEFAULT_PORT    8080
#define BACKLOG         5
#define MAX_REQUEST     8192
#define MAX_HEADERS     32
#define MAX_SESSIONS    8
#define MAX_PARAMS      16
#define SESS_ID_LEN     8
#define DEV_ID_SZ       32
#define PARAM_VAL_SZ    64
#define FW_DESC_SZ      64

/* ------------------------------------------------------------------ */
/*  Types                                                              */
/* ------------------------------------------------------------------ */

typedef struct {
    char name[64];
    char value[512];
} header_t;

typedef void (*log_fn_t)(const char *);

typedef struct {
    char     id[SESS_ID_LEN + 1];
    char     username[32];
    int      priv_level;          /* 0=guest  1=user  2=admin */
    log_fn_t log_action;          /* called on status check    */
    int      active;
} session_t;

typedef struct {
    char      method[16];
    char      path[512];
    header_t  headers[MAX_HEADERS];
    int       n_headers;
    int       content_length;
    char     *body;               /* heap-allocated if CL > 0 */
} request_t;

/* ------------------------------------------------------------------ */
/*  Globals                                                            */
/* ------------------------------------------------------------------ */

static session_t *g_sessions[MAX_SESSIONS];
static char      *g_device_id  = NULL;
static int        g_running    = 1;
static int        g_req_count  = 0;

/* ------------------------------------------------------------------ */
/*  Logging callbacks (targets for UAF)                                */
/* ------------------------------------------------------------------ */

static void log_user(const char *msg)
{
    fprintf(stderr, "[USER]  %s\n", msg);
}

static void log_admin(const char *msg)
{
    fprintf(stderr, "[ADMIN] %s\n", msg);
}

/* ------------------------------------------------------------------ */
/*  Utility helpers                                                    */
/* ------------------------------------------------------------------ */

static void send_response(int fd, int code, const char *body)
{
    const char *reason;
    switch (code) {
    case 200: reason = "OK";                  break;
    case 400: reason = "Bad Request";         break;
    case 401: reason = "Unauthorized";        break;
    case 404: reason = "Not Found";           break;
    case 500: reason = "Internal Server Error"; break;
    default:  reason = "Unknown";             break;
    }

    char hdr[512];
    int  blen = body ? (int)strlen(body) : 0;
    int  hlen = snprintf(hdr, sizeof(hdr),
        "HTTP/1.1 %d %s\r\n"
        "Content-Type: text/html\r\n"
        "Content-Length: %d\r\n"
        "Connection: keep-alive\r\n"
        "\r\n", code, reason, blen);

    write(fd, hdr, hlen);
    if (body && blen > 0)
        write(fd, body, blen);
}

static const char *get_header(const request_t *req, const char *name)
{
    for (int i = 0; i < req->n_headers; i++)
        if (strcasecmp(req->headers[i].name, name) == 0)
            return req->headers[i].value;
    return NULL;
}

/* Read a line (up to \r\n) from fd into buf. Returns bytes read, 0 on
   disconnect, -1 on error.  Intentionally byte-at-a-time for
   simplicity — this is a CTF target, not production code. */
static int read_line(int fd, char *buf, int max)
{
    int i = 0;
    char c;
    while (i < max - 1) {
        int n = read(fd, &c, 1);
        if (n <= 0) return n == 0 ? 0 : -1;
        if (c == '\r') {
            read(fd, &c, 1);   /* consume \n */
            break;
        }
        buf[i++] = c;
    }
    buf[i] = '\0';
    return i;
}

/* ------------------------------------------------------------------ */
/*  Request parser                                                     */
/* ------------------------------------------------------------------ */

static int parse_request(int fd, request_t *req)
{
    memset(req, 0, sizeof(*req));
    req->body = NULL;

    /* --- request line --- */
    char line[MAX_REQUEST];
    int n = read_line(fd, line, sizeof(line));
    if (n <= 0) return -1;

    char *sp1 = strchr(line, ' ');
    if (!sp1) return -1;
    *sp1 = '\0';
    strncpy(req->method, line, sizeof(req->method) - 1);

    char *sp2 = strchr(sp1 + 1, ' ');
    if (sp2) *sp2 = '\0';
    strncpy(req->path, sp1 + 1, sizeof(req->path) - 1);

    /* --- headers --- */
    while (req->n_headers < MAX_HEADERS) {
        n = read_line(fd, line, sizeof(line));
        if (n <= 0) break;
        if (line[0] == '\0') break;          /* blank line = end */

        char *colon = strchr(line, ':');
        if (!colon) continue;
        *colon = '\0';
        char *val = colon + 1;
        while (*val == ' ') val++;

        strncpy(req->headers[req->n_headers].name,  line, 63);
        strncpy(req->headers[req->n_headers].value, val,  511);
        req->n_headers++;
    }

    /* --- Content-Length & body --- */
    const char *cl = get_header(req, "Content-Length");
    if (cl) req->content_length = atoi(cl);

    if (req->content_length > 0 && req->content_length < MAX_REQUEST) {
        req->body = malloc(req->content_length + 1);
        if (!req->body) return -1;
        int total = 0;
        while (total < req->content_length) {
            n = read(fd, req->body + total,
                     req->content_length - total);
            if (n <= 0) break;
            total += n;
        }
        req->body[total] = '\0';
    }

    return 0;
}

static void cleanup_request(request_t *req)
{
    if (req->body) {
        free(req->body);
        req->body = NULL;
    }
}

/* ------------------------------------------------------------------ */
/*  Session management                                                 */
/* ------------------------------------------------------------------ */

static session_t *find_session(const char *id)
{
    if (!id) return NULL;
    for (int i = 0; i < MAX_SESSIONS; i++)
        if (g_sessions[i] && strcmp(g_sessions[i]->id, id) == 0)
            return g_sessions[i];
    return NULL;
}

static session_t *create_session(const char *user, int priv)
{
    for (int i = 0; i < MAX_SESSIONS; i++) {
        if (!g_sessions[i]) {
            session_t *s = malloc(sizeof(session_t));
            if (!s) return NULL;
            snprintf(s->id, sizeof(s->id), "%08x", rand());
            strncpy(s->username, user, sizeof(s->username) - 1);
            s->username[sizeof(s->username) - 1] = '\0';
            s->priv_level = priv;
            s->log_action = (priv >= 2) ? log_admin : log_user;
            s->active     = 1;
            g_sessions[i] = s;
            return s;
        }
    }
    return NULL;
}

/* ------------------------------------------------------------------ */
/*  Route: GET /                                                       */
/* ------------------------------------------------------------------ */

static void handle_index(int fd, request_t *req)
{
    (void)req;
    const char *page =
        "<html><body>"
        "<h2>MiniHTTPD — IoT Gateway v1.3.7</h2>"
        "<p>Endpoints:</p>"
        "<ul>"
        "<li>POST /api/login       — body: user=NAME&pass=PASS</li>"
        "<li>GET  /api/logout      — header: X-Session-Id</li>"
        "<li>GET  /api/status      — header: X-Session-Id</li>"
        "<li>POST /api/config      — body: key=val&key2=val2</li>"
        "<li>POST /api/firmware    — header: X-Firmware-Size, X-Firmware-Desc</li>"
        "<li>POST /api/diag        — body: PING|UPTIME|ECHO msg</li>"
        "<li>GET  /api/log         — header: X-Log-Size (opt)</li>"
        "<li>GET  /api/debug       — leak heap / GOT addresses (cheat)</li>"
        "</ul>"
        "</body></html>";
    send_response(fd, 200, page);
}

/* ------------------------------------------------------------------ */
/*  Route: GET /api/debug   — address leak helper                      */
/* ------------------------------------------------------------------ */

static void handle_debug(int fd, request_t *req)
{
    (void)req;
    char buf[1024];
    snprintf(buf, sizeof(buf),
        "<pre>\n"
        "=== Debug Info ===\n"
        "g_device_id ptr : %p\n"
        "g_sessions[0]   : %p\n"
        "heap sample     : %p\n"
        "log_user()      : %p\n"
        "log_admin()     : %p\n"
        "system@plt      : %p\n"
        "free@got        : %p\n"
        "puts@got        : %p\n"
        "sizeof(session) : %zu\n"
        "</pre>\n",
        (void *)g_device_id,
        (void *)g_sessions[0],
        (void *)malloc(1),        /* quick heap probe */
        (void *)log_user,
        (void *)log_admin,
        (void *)system,
        (void *)&free,
        (void *)&puts,
        sizeof(session_t));
    send_response(fd, 200, buf);
}

/* ------------------------------------------------------------------ */
/*  Route: GET /api/info   — device info (processes X-Device-Id)       */
/* ------------------------------------------------------------------ */

static void handle_info(int fd, request_t *req)
{
    const char *dev = get_header(req, "X-Device-Id");

    if (dev) {
        /* Update stored device identifier.
           Router vendors typically keep a short device tag for
           telemetry; we allocate a small fixed buffer for it. */
        if (g_device_id)
            free(g_device_id);

        g_device_id = malloc(DEV_ID_SZ);
        if (!g_device_id) {
            send_response(fd, 500, "alloc failed");
            return;
        }

        /* Copy the device-id string from the header */
        strcpy(g_device_id, dev);
    }

    char resp[256];
    snprintf(resp, sizeof(resp),
        "{\"device_id\":\"%s\",\"uptime\":%d}",
        g_device_id ? g_device_id : "unknown",
        g_req_count);
    send_response(fd, 200, resp);
}

/* ------------------------------------------------------------------ */
/*  Route: POST /api/login                                             */
/* ------------------------------------------------------------------ */

static void handle_login(int fd, request_t *req)
{
    if (!req->body) {
        send_response(fd, 400, "Missing credentials");
        return;
    }

    char user[32] = {0}, pass[32] = {0};

    /* Minimal body parser: user=xxx&pass=yyy */
    char *u = strstr(req->body, "user=");
    char *p = strstr(req->body, "pass=");
    if (u) {
        u += 5;
        char *end = strchr(u, '&');
        if (end) *end = '\0';
        strncpy(user, u, sizeof(user) - 1);
        if (end) *end = '&';
    }
    if (p) {
        p += 5;
        char *end = strchr(p, '&');
        if (end) *end = '\0';
        strncpy(pass, p, sizeof(pass) - 1);
        if (end) *end = '&';
    }

    /* Hardcoded creds — typical embedded device */
    int priv = -1;
    if (strcmp(user, "admin") == 0 && strcmp(pass, "admin") == 0)
        priv = 2;
    else if (strcmp(user, "user") == 0 && strcmp(pass, "user") == 0)
        priv = 1;

    if (priv < 0) {
        send_response(fd, 401, "Bad credentials");
        return;
    }

    session_t *s = create_session(user, priv);
    if (!s) {
        send_response(fd, 500, "Session table full");
        return;
    }

    char resp[128];
    snprintf(resp, sizeof(resp),
        "{\"session\":\"%s\",\"priv\":%d}", s->id, s->priv_level);
    send_response(fd, 200, resp);
}

/* ------------------------------------------------------------------ */
/*  Route: GET /api/logout                                             */
/* ------------------------------------------------------------------ */

static void handle_logout(int fd, request_t *req)
{
    const char *sid = get_header(req, "X-Session-Id");
    session_t  *s   = find_session(sid);

    if (!s) {
        send_response(fd, 400, "No such session");
        return;
    }

    s->active = 0;
    free(s);
    /* Developer note: session removed from active use.
       The global table entry is not cleared here because
       the audit log reader iterates g_sessions[] later. */

    send_response(fd, 200, "{\"status\":\"logged_out\"}");
}

/* ------------------------------------------------------------------ */
/*  Route: GET /api/status                                             */
/* ------------------------------------------------------------------ */

static void handle_status(int fd, request_t *req)
{
    const char *sid = get_header(req, "X-Session-Id");
    session_t  *s   = find_session(sid);

    if (!s) {
        send_response(fd, 401, "Invalid session");
        return;
    }

    /* Log the status check via the session's logging callback */
    char logmsg[128];
    snprintf(logmsg, sizeof(logmsg), "status check by %s", s->username);
    s->log_action(logmsg);

    char resp[256];
    snprintf(resp, sizeof(resp),
        "{\"user\":\"%s\",\"priv\":%d,\"active\":%d}",
        s->username, s->priv_level, s->active);
    send_response(fd, 200, resp);
}

/* ------------------------------------------------------------------ */
/*  Route: POST /api/config                                            */
/* ------------------------------------------------------------------ */

static void handle_config(int fd, request_t *req)
{
    if (!req->body) {
        send_response(fd, 400, "No config data");
        return;
    }

    /* Parse URL-encoded config: key=value&key2=value2&... */
    char *saveptr = NULL;
    char *work    = strdup(req->body);
    if (!work) {
        send_response(fd, 500, "alloc");
        return;
    }

    int   count = 0;
    char *pairs[MAX_PARAMS];
    char *values[MAX_PARAMS];

    char *tok = strtok_r(work, "&", &saveptr);
    while (tok && count < MAX_PARAMS) {
        char *eq = strchr(tok, '=');
        if (!eq) { tok = strtok_r(NULL, "&", &saveptr); continue; }
        *eq = '\0';

        pairs[count] = strdup(tok);             /* key   */
        values[count] = malloc(PARAM_VAL_SZ);   /* fixed-size value buf */
        if (!values[count]) break;

        /* Store the value — trusting that config values are short */
        strcpy(values[count], eq + 1);

        fprintf(stderr, "[CONFIG] %s = %s\n", pairs[count], values[count]);
        count++;
        tok = strtok_r(NULL, "&", &saveptr);
    }

    /* Cleanup */
    for (int i = 0; i < count; i++) {
        free(pairs[i]);
        free(values[i]);
    }
    free(work);

    char resp[64];
    snprintf(resp, sizeof(resp), "{\"saved\":%d}", count);
    send_response(fd, 200, resp);
}

/* ------------------------------------------------------------------ */
/*  Route: POST /api/firmware                                          */
/* ------------------------------------------------------------------ */

static void handle_firmware(int fd, request_t *req)
{
    /* X-Firmware-Size declares the total firmware image size.
       The body of this request carries the first chunk. */
    const char *sz_str   = get_header(req, "X-Firmware-Size");
    const char *desc_str = get_header(req, "X-Firmware-Desc");

    if (!sz_str) {
        send_response(fd, 400, "Missing X-Firmware-Size");
        return;
    }

    size_t fw_size = strtoul(sz_str, NULL, 10);
    if (fw_size == 0) {
        send_response(fd, 400, "Invalid firmware size");
        return;
    }

    /* Allocate receive buffer for the declared firmware size */
    char *fw_buf = malloc(fw_size);
    if (!fw_buf) {
        send_response(fd, 500, "Firmware buffer alloc failed");
        return;
    }

    /* Copy in whatever body data we received (may be partial) */
    if (req->body && req->content_length > 0) {
        size_t copy_len = (size_t)req->content_length;
        if (copy_len > fw_size) copy_len = fw_size;
        memcpy(fw_buf, req->body, copy_len);
    }

    /* Validate firmware magic bytes */
    if (!req->body || req->content_length < 4 ||
        memcmp(fw_buf, "FWUP", 4) != 0) {
        send_response(fd, 400, "Bad firmware magic");
        free(fw_buf);
        return;
    }

    /* Store firmware metadata */
    char *meta = malloc(FW_DESC_SZ);
    if (meta && desc_str) {
        strcpy(meta, desc_str);
    } else if (meta) {
        strcpy(meta, "no description");
    }

    fprintf(stderr, "[FW] Accepted %zu bytes, desc: %s\n", fw_size,
            meta ? meta : "(null)");

    char resp[128];
    snprintf(resp, sizeof(resp),
        "{\"status\":\"ok\",\"size\":%zu,\"meta\":\"%s\"}",
        fw_size, meta ? meta : "");
    send_response(fd, 200, resp);

    free(meta);
    free(fw_buf);
}

/* ------------------------------------------------------------------ */
/*  Route: GET /api/log   — activity log                               */
/* ------------------------------------------------------------------ */

/* Internal log-cursor used for paginated log output.
   Tracks position, formatter callback, and buffer ref. */
typedef struct {
    int          seq;
    int          total;
    log_fn_t     emit;          /* formatter callback   */
    char         tag[24];
    void        *buf_ref;       /* back-pointer to data */
} log_cursor_t;

#define LOG_BUF_SZ  128

static void handle_log(int fd, request_t *req)
{
    /* X-Log-Size: how many bytes the client wants back.
       Intended for scripted log consumers that know their
       terminal width or parse buffer size. */
    const char *sz_str = get_header(req, "X-Log-Size");
    int resp_sz = sz_str ? atoi(sz_str) : LOG_BUF_SZ;
    if (resp_sz <= 0 || resp_sz > 4096) resp_sz = LOG_BUF_SZ;

    /* Allocate the text buffer */
    char *logbuf = malloc(LOG_BUF_SZ);
    if (!logbuf) {
        send_response(fd, 500, "alloc");
        return;
    }

    /* Allocate cursor context (lives right after logbuf on the heap) */
    log_cursor_t *cur = malloc(sizeof(log_cursor_t));
    if (!cur) {
        free(logbuf);
        send_response(fd, 500, "alloc");
        return;
    }

    cur->seq     = g_req_count;
    cur->total   = g_req_count;
    cur->emit    = log_admin;        /* function pointer       */
    strncpy(cur->tag, "activity", sizeof(cur->tag));
    cur->buf_ref = logbuf;           /* heap address reference */

    /* Format the log text into the fixed buffer */
    int written = 0;
    written += snprintf(logbuf + written, LOG_BUF_SZ - written,
        "=== Activity Log ===\n"
        "Requests: %d\n"
        "Device: %s\n",
        g_req_count,
        g_device_id ? g_device_id : "(unset)");

    for (int i = 0; i < MAX_SESSIONS && written < LOG_BUF_SZ; i++) {
        if (g_sessions[i]) {
            written += snprintf(logbuf + written, LOG_BUF_SZ - written,
                "Session[%d]: %s (priv=%d)\n",
                i, g_sessions[i]->id, g_sessions[i]->priv_level);
        }
    }

    /* Send the log data.  The text buffer is LOG_BUF_SZ bytes,
       but we honour the client's requested resp_sz for the
       Content-Length so scripted consumers get a fixed-size read. */
    char hdr[256];
    int hlen = snprintf(hdr, sizeof(hdr),
        "HTTP/1.1 200 OK\r\n"
        "Content-Type: application/octet-stream\r\n"
        "Content-Length: %d\r\n"
        "Connection: keep-alive\r\n"
        "\r\n", resp_sz);

    write(fd, hdr, hlen);
    write(fd, logbuf, resp_sz);

    free(cur);
    free(logbuf);
}

/* ------------------------------------------------------------------ */
/*  Route: POST /api/diag                                              */
/* ------------------------------------------------------------------ */

static void handle_diag(int fd, request_t *req)
{
    if (!req->body || req->content_length <= 0) {
        send_response(fd, 400, "No diagnostic data");
        return;
    }

    char *result = malloc(256);
    if (!result) {
        send_response(fd, 500, "alloc");
        return;
    }

    if (strncmp(req->body, "PING", 4) == 0) {
        snprintf(result, 256, "{\"pong\":\"%s\"}", req->body + 4);
    } else if (strncmp(req->body, "UPTIME", 6) == 0) {
        snprintf(result, 256, "{\"requests\":%d}", g_req_count);
    } else if (strncmp(req->body, "ECHO ", 5) == 0) {
        /* Echo back the message for connectivity testing */
        snprintf(result, 256, req->body + 5);
    } else {
        /* Unrecognized diagnostic command — reject and clean up */
        free(result);
        free(req->body);       /* release input buffer early */
        send_response(fd, 400, "{\"error\":\"unknown command\"}");
        return;
    }

    send_response(fd, 200, result);
    free(result);
}

/* ------------------------------------------------------------------ */
/*  Router                                                             */
/* ------------------------------------------------------------------ */

static void dispatch(int fd, request_t *req)
{
    g_req_count++;

    if (strcmp(req->method, "GET") == 0) {
        if (strcmp(req->path, "/") == 0)
            handle_index(fd, req);
        else if (strcmp(req->path, "/api/info") == 0)
            handle_info(fd, req);
        else if (strcmp(req->path, "/api/debug") == 0)
            handle_debug(fd, req);
        else if (strcmp(req->path, "/api/logout") == 0)
            handle_logout(fd, req);
        else if (strcmp(req->path, "/api/status") == 0)
            handle_status(fd, req);
        else if (strcmp(req->path, "/api/log") == 0)
            handle_log(fd, req);
        else
            send_response(fd, 404, "Not found");

    } else if (strcmp(req->method, "POST") == 0) {
        if (strcmp(req->path, "/api/login") == 0)
            handle_login(fd, req);
        else if (strcmp(req->path, "/api/config") == 0)
            handle_config(fd, req);
        else if (strcmp(req->path, "/api/firmware") == 0)
            handle_firmware(fd, req);
        else if (strcmp(req->path, "/api/diag") == 0)
            handle_diag(fd, req);
        else
            send_response(fd, 404, "Not found");
    } else {
        send_response(fd, 400, "Unsupported method");
    }
}

/* ------------------------------------------------------------------ */
/*  Connection handler — loops for keep-alive                          */
/* ------------------------------------------------------------------ */

static void handle_connection(int client_fd)
{
    fprintf(stderr, "[CONN] new connection fd=%d\n", client_fd);

    while (g_running) {
        request_t req;
        if (parse_request(client_fd, &req) < 0)
            break;

        dispatch(client_fd, &req);

        /* Check for Connection: close */
        const char *conn = get_header(&req, "Connection");
        int do_close = (conn && strcasecmp(conn, "close") == 0);

        cleanup_request(&req);

        if (do_close)
            break;
    }

    close(client_fd);
    fprintf(stderr, "[CONN] closed fd=%d\n", client_fd);
}

/* ------------------------------------------------------------------ */
/*  Main                                                               */
/* ------------------------------------------------------------------ */

static void sigint_handler(int sig)
{
    (void)sig;
    g_running = 0;
}

int main(int argc, char *argv[])
{
    int port = (argc > 1) ? atoi(argv[1]) : DEFAULT_PORT;

    signal(SIGINT,  sigint_handler);
    signal(SIGPIPE, SIG_IGN);
    srand(0xdeadbeef);        /* deterministic sessions for CTF */

    int sfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sfd < 0) { perror("socket"); return 1; }

    int opt = 1;
    setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_port   = htons(port),
        .sin_addr.s_addr = INADDR_ANY
    };

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

    fprintf(stderr, "[*] mini_httpd listening on :%d\n", port);
    //fprintf(stderr, "[*] Heap base (approx): %p\n", malloc(1));

    while (g_running) {
        struct sockaddr_in peer;
        socklen_t plen = sizeof(peer);
        int cfd = accept(sfd, (struct sockaddr *)&peer, &plen);
        if (cfd < 0) {
            if (errno == EINTR) continue;
            perror("accept");
            break;
        }

        /* Single-threaded: handle one connection at a time.
           Typical for resource-constrained embedded HTTP daemons. */
        handle_connection(cfd);
    }

    close(sfd);
    fprintf(stderr, "[*] shutting down\n");
    return 0;
}
