/*
 * DF-0753 — FIXED-code-path variant of the stale-pointer harness.
 *
 * This file transcribes the SAME mpls_forward/mpls_output/mpls_push/
 * mpls_swap/mpls_pop/m_prepend/m_pullup code as mpls_stale_harness.c,
 * but with the FIX applied: mpls_output/mpls_swap/mpls_pop take
 * `struct mbuf **mp` and write the new head back through *mp.
 *
 * Expected result: ALL scenarios (including A/A2/B that double-freed
 * in the unpatched harness) now complete with ZERO double-frees and
 * ZERO use-after-frees — the caller always sees the correct head.
 *
 * Build:  cc -O2 -Wall -o mpls_stale_harness_fixed mpls_stale_harness_fixed.c
 * Run:    ./mpls_stale_harness_fixed
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <assert.h>

#define MHLEN        84
#define MLEN         100
#define M_PKTHDR     0x02
#define M_EXT        0x04
#define M_MPLSLABELED 0x4000
#define ETHER_HDR_LEN 14
#define MPLS_SHIM_LEN 4

struct mbuf {
    uint32_t magic;
    uint32_t state;
    struct mbuf *m_next;
    int  m_flags;
    int  m_len;
    char *m_data;
    int  pkthdr_len;
    char m_pktdat[MLEN];
    int  refcnt;
};

#define LIVE_MAGIC  0x4d42464d
#define DEAD_MAGIC  0xdeadc0de

static int alloc_count = 0;
static int free_count  = 0;
static int double_free_detected = 0;
static int uaf_detected = 0;

static struct mbuf *mbuf_alloc(int flags)
{
    struct mbuf *m = calloc(1, sizeof(struct mbuf));
    assert(m);
    m->magic = LIVE_MAGIC;
    m->state = 1;
    m->m_flags = flags;
    m->m_data = m->m_pktdat;
    m->m_next = NULL;
    m->refcnt = 1;
    alloc_count++;
    return m;
}

static struct mbuf *m_gethdr(int how, int type)
{
    (void)how; (void)type;
    return mbuf_alloc(M_PKTHDR);
}

static struct mbuf *check_live(struct mbuf *m, const char *where)
{
    if (m == NULL) return NULL;
    if (m->state == 0 || m->magic == DEAD_MAGIC) {
        printf("  [UAF] deref of freed mbuf %p at %s\n", (void*)m, where);
        uaf_detected++;
        return NULL;
    }
    return m;
}

static void m_free_one(struct mbuf *m)
{
    if (m == NULL) return;
    if (m->state == 0 || m->magic == DEAD_MAGIC) {
        printf("  [DOUBLE-FREE] m_freem on already-freed mbuf %p\n", (void*)m);
        double_free_detected++;
        return;
    }
    m->state = 0;
    m->magic = DEAD_MAGIC;
    free_count++;
}

static void m_freem(struct mbuf *m)
{
    while (m) {
        struct mbuf *next = m->m_next;
        m_free_one(m);
        m = next;
    }
}

static void M_MOVE_PKTHDR(struct mbuf *to, struct mbuf *from)
{
    assert(to->m_flags & M_PKTHDR);
    assert(from->m_flags & M_PKTHDR);
    to->m_flags |= (from->m_flags & (M_PKTHDR|M_MPLSLABELED));
    to->pkthdr_len = from->pkthdr_len;
}

static int M_LEADINGSPACE(struct mbuf *m)
{
    return (int)(m->m_data - m->m_pktdat);
}

#define M_PREPEND(mp, plen) do {                          \
    if (M_LEADINGSPACE(*(mp)) >= (plen)) {                \
        (*(mp))->m_data -= (plen);                        \
        (*(mp))->m_len   += (plen);                        \
    } else {                                              \
        *(mp) = m_prepend_harness(*(mp), (plen));         \
    }                                                     \
    if (*(mp) && (*(mp))->m_flags & M_PKTHDR)             \
        (*(mp))->pkthdr_len += (plen);                    \
} while(0)

static struct mbuf *m_prepend_harness(struct mbuf *m, int len)
{
    struct mbuf *mn;
    if (m->m_flags & M_PKTHDR)
        mn = m_gethdr(0, 0);
    else
        mn = mbuf_alloc(0);
    if (mn == NULL) { m_freem(m); return NULL; }
    if (m->m_flags & M_PKTHDR)
        M_MOVE_PKTHDR(mn, m);
    mn->m_next = m;
    mn->m_len  = len;
    mn->m_data = mn->m_pktdat;
    return mn;
}

static struct mbuf *m_pullup_harness(struct mbuf *n, int len)
{
    if (n->m_len >= len)
        return n;
    {
        struct mbuf *m;
        if (n->m_flags & M_PKTHDR)
            m = m_gethdr(0, 0);
        else
            m = mbuf_alloc(0);
        if (m == NULL) { m_freem(n); return NULL; }
        m->m_len = 0;
        if (n->m_flags & M_PKTHDR)
            M_MOVE_PKTHDR(m, n);
        {
            int want = len;
            struct mbuf *src = n;
            while (want > 0 && src) {
                struct mbuf *next = src->m_next;
                int cnt = src->m_len < want ? src->m_len : want;
                memcpy(m->m_pktdat + m->m_len, src->m_data, cnt);
                m->m_len += cnt;
                want     -= cnt;
                m_free_one(src);
                src = next;
            }
        }
        m->m_data = m->m_pktdat;
        m->m_next = NULL;
        return m;
    }
}

struct mpls { uint32_t mpls_shim; };
#define MPLS_LABEL(s)   ((s >> 12) & 0xfffff)
#define MPLS_STACK(s)   ((s >> 8) & 1)
#define MPLS_TTL(s)     (s & 0xff)
#define MPLS_SET_LABEL(b,l) (b |= ((l & 0xfffff) << 12))
#define MPLS_SET_STACK(b,s) (b |= ((s & 1) << 8))
#define MPLS_SET_TTL(b,t)   (b |= (t & 0xff))

struct sockaddr_mpls { int smpls_op; uint32_t smpls_label; };
#define MPLSLOP_PUSH 1
#define MPLSLOP_SWAP 2
#define MPLSLOP_POP  3

struct rtentry {
    struct sockaddr_mpls shim[3];
    int nshim;
    int from_mpls;
};
#define AF_MPLS 35

/* ===== FIXED versions: all take struct mbuf **mp ===== */

static int mpls_push(struct mbuf **m, uint32_t label, int s, int ttl)
{
    uint32_t buf = 0;
    M_PREPEND(m, MPLS_SHIM_LEN);
    if (*m == NULL) return -1;
    MPLS_SET_LABEL(buf, label);
    MPLS_SET_STACK(buf, s);
    MPLS_SET_TTL(buf, ttl);
    {
        struct mpls *p = (struct mpls*)(*m)->m_data;
        p->mpls_shim = buf;
    }
    (*m)->m_flags |= M_MPLSLABELED;
    return 0;
}

/* FIXED: takes struct mbuf **mp, writes back new head */
static int mpls_swap_FIXED(struct mbuf **mp, uint32_t label)
{
    struct mbuf *m = *mp;
    if (m->m_len < MPLS_SHIM_LEN) {
        m = m_pullup_harness(m, MPLS_SHIM_LEN);
        if (m == NULL) { *mp = NULL; return -1; }
        *mp = m;   /* *** FIX: propagate new head *** */
    }
    {
        struct mpls *p = (struct mpls*)m->m_data;
        uint32_t buf = p->mpls_shim;
        int ttl = MPLS_TTL(buf);
        if (--ttl <= 0) return -2;
        buf = 0;
        MPLS_SET_LABEL(buf, label);
        MPLS_SET_TTL(buf, ttl);
        p->mpls_shim = buf;
    }
    return 0;
}

static int mpls_pop_FIXED(struct mbuf **mp, int *sbit)
{
    struct mbuf *m = *mp;
    if (m->m_len < MPLS_SHIM_LEN) {
        m = m_pullup_harness(m, MPLS_SHIM_LEN);
        if (m == NULL) { *mp = NULL; return -1; }
        *mp = m;   /* *** FIX *** */
    }
    {
        struct mpls *p = (struct mpls*)m->m_data;
        uint32_t buf = p->mpls_shim;
        *sbit = MPLS_STACK(buf);
    }
    m->m_data += MPLS_SHIM_LEN;
    m->m_len  -= MPLS_SHIM_LEN;
    return 0;
}

/* FIXED: takes struct mbuf **mp, writes *mp = m at the end */
static int mpls_output_FIXED(struct mbuf **mp, struct rtentry *rt)
{
    struct mbuf *m = *mp;   /* load caller's head */
    int i, stackempty;
    int ttl = 255;
    int error = 0;
    stackempty = rt->from_mpls ? 0 : 1;

    for (i = 0; i < rt->nshim; i++) {
        struct sockaddr_mpls *s = &rt->shim[i];
        switch (s->smpls_op) {
        case MPLSLOP_PUSH:
            error = mpls_push(&m, s->smpls_label,
                            (stackempty && i == 0) ? 1 : 0, ttl);
            if (error) goto out;
            stackempty = 0;
            break;
        case MPLSLOP_SWAP:
            if (stackempty) { error = -3; goto out; }
            error = mpls_swap_FIXED(&m, s->smpls_label);  /* *** FIX: &m *** */
            if (error) goto out;
            break;
        case MPLSLOP_POP:
            if (stackempty) { error = -3; goto out; }
            { int sb; error = mpls_pop_FIXED(&m, &sb); }  /* *** FIX: &m *** */
            if (error) goto out;
            break;
        }
    }
out:
    *mp = m;   /* *** FIX: propagate new head to caller *** */
    return error;
}

static int if_output_calls = 0;
static int fake_if_output(struct mbuf *m)
{
    if_output_calls++;
    if (check_live(m, "if_output(m)")) {
        printf("  if_output: received LIVE mbuf %p (m_len=%d flags=0x%x)\n",
               (void*)m, m->m_len, m->m_flags);
        if (m->m_flags & M_MPLSLABELED)
            printf("         (correct head with M_MPLSLABELED — FIX works)\n");
    }
    m_freem(m);
    return 0;
}

static int fake_if_output_error(struct mbuf *m)
{
    /* Faithful to ifq_dispatch (sys/net/if.c:3344): on enqueue failure
     * it returns the error but does NOT free m -- the caller owns it.
     * So we just report and return error; caller's m_freem is the
     * single (correct) free. */
    if_output_calls++;
    check_live(m, "if_output_error(m)");
    /* do NOT free -- caller owns m on error */
    return -1;
}

/* FIXED: mpls_forward passes &m and uses the returned head */
static void mpls_forward_FIXED(struct mbuf *m, struct rtentry *rt,
                               int (*ifp_if_output)(struct mbuf *))
{
    int error;
    printf("  mpls_forward: m=%p (head, m_flags=0x%x)\n",
           (void*)m, m->m_flags);
    error = mpls_output_FIXED(&m, rt);              /* *** FIX: &m *** */
    if (error) { printf("  mpls_output returned %d\n", error); goto bad; }
    error = ifp_if_output(m);                        /* now uses CORRECT head */
    if (error) { printf("  if_output returned %d\n", error); goto bad; }
    printf("  mpls_forward: forwarded OK\n");
    return;
bad:
    printf("  mpls_forward bad: m_freem(%p)\n", (void*)m);
    check_live(m, "m_freem at bad:");
    m_freem(m);
}

static struct mbuf *make_rx_frame(int payload_len, int leading_space)
{
    struct mbuf *m = m_gethdr(0, 0);
    m->m_data = m->m_pktdat + leading_space;
    memset(m->m_pktdat, 0xAA, sizeof(m->m_pktdat));
    {
        struct mpls *p = (struct mpls*)m->m_data;
        MPLS_SET_LABEL(p->mpls_shim, 100);
        MPLS_SET_STACK(p->mpls_shim, 1);
        MPLS_SET_TTL(p->mpls_shim, 64);
    }
    m->m_len = MPLS_SHIM_LEN + payload_len;
    m->pkthdr_len = m->m_len;
    return m;
}

static void reset_accounting(void)
{
    alloc_count = free_count = 0;
    double_free_detected = uaf_detected = 0;
    if_output_calls = 0;
}

static void report(const char *label)
{
    printf("\n=== %s ===\n", label);
    printf("  allocs=%d frees=%d  double_free=%d  uaf=%d\n",
           alloc_count, free_count, double_free_detected, uaf_detected);
    if (double_free_detected)
        printf("  *** FAIL: DOUBLE-FREE still present ***\n");
    else
        printf("  PASS: no double-free (fix eliminates stale pointer)\n");
}

int main(void)
{
    struct rtentry rt_push;
    int i;

    printf("DF-0753 FIXED-code-path harness\n");
    printf("mpls_output/mpls_swap/mpls_pop now take struct mbuf **mp\n");
    printf("and propagate the new head through *mp.\n\n");

    /* Scenario A: PUSH no-headroom + if_output success */
    printf("--------------------------------------------------------\n");
    printf("Scenario A (FIXED): PUSH leading_space=2, if_output success\n");
    printf("--------------------------------------------------------\n");
    {
        struct rtentry rt = {0};
        rt.from_mpls = 1; rt.nshim = 1;
        rt.shim[0].smpls_op = MPLSLOP_PUSH; rt.shim[0].smpls_label = 999;
        reset_accounting();
        struct mbuf *m = make_rx_frame(20, 2);
        mpls_forward_FIXED(m, &rt, fake_if_output);
    }
    report("A FIXED: PUSH no-headroom, success");

    /* Scenario A2: PUSH no-headroom + if_output error */
    printf("\n--------------------------------------------------------\n");
    printf("Scenario A2 (FIXED): PUSH no-headroom, if_output ERROR\n");
    printf("--------------------------------------------------------\n");
    {
        struct rtentry rt = {0};
        rt.from_mpls = 1; rt.nshim = 1;
        rt.shim[0].smpls_op = MPLSLOP_PUSH; rt.shim[0].smpls_label = 999;
        reset_accounting();
        struct mbuf *m = make_rx_frame(20, 2);
        mpls_forward_FIXED(m, &rt, fake_if_output_error);
    }
    report("A2 FIXED: PUSH no-headroom, if_output error");

    /* Scenario B: SWAP fragmented */
    printf("\n--------------------------------------------------------\n");
    printf("Scenario B (FIXED): SWAP m_len=2 -> m_pullup\n");
    printf("--------------------------------------------------------\n");
    {
        struct rtentry rt = {0};
        rt.from_mpls = 1; rt.nshim = 1;
        rt.shim[0].smpls_op = MPLSLOP_SWAP; rt.shim[0].smpls_label = 200;
        reset_accounting();
        struct mbuf *m = m_gethdr(0, 0);
        struct mbuf *m2 = mbuf_alloc(0);
        m->m_flags |= M_MPLSLABELED;
        m->m_len = 2;
        memcpy(m->m_pktdat, "\x00\x00", 2);
        m->m_data = m->m_pktdat;
        m2->m_len = 32;
        memset(m2->m_pktdat, 0xBB, 32);
        m2->m_data = m2->m_pktdat;
        m->m_next = m2;
        m->pkthdr_len = 34;
        mpls_forward_FIXED(m, &rt, fake_if_output);
    }
    report("B FIXED: SWAP fragmented");

    printf("\n=========================================================\n");
    printf("FIXED harness complete.\n");
    if (double_free_detected == 0 && uaf_detected == 0)
        printf("RESULT: ALL scenarios PASS — fix eliminates stale pointer.\n");
    else
        printf("RESULT: FAIL — stale pointer still present.\n");
    return 0;
}