/*
 * DF-0754 — FIXED-code-path variant of the harness.
 *
 * Transcribes the SAME mpls_output/mpls_push/mpls_swap/mpls_pop/
 * mpls_output_process/ip_output-dispatch code as harness.c, but with the
 * FIX applied (identical to the DF-0753 root-cause fix):
 *   - mpls_output()       takes `struct mbuf **mp`   (mpls_output.c:50)
 *   - mpls_swap()         takes `struct mbuf **mp`   (mpls_output.c:171)
 *   - mpls_pop()          takes `struct mbuf **mp`   (mpls_output.c:196)
 *   - mpls_output_process() takes `struct mbuf **mp` (mpls_output.c:134)
 *   - ip_output dispatch  passes &m                  (ip_output.c:695,739)
 * and the new head is written back through *mp at every rebind, plus an
 * `out:` label that always propagates the (possibly new) head before return.
 *
 * Expected: ALL DF-0754 scenarios (M1a/M1b/M1c/M2a/M2b) that double-freed /
 * leaked / sent a stale mbuf in the unpatched harness now complete with
 * ZERO double-frees, ZERO UAFs, ZERO leaks, and the driver receives the
 * CORRECT (new) head mbuf.
 *
 * Build:  cc -O2 -Wall -o harness_fixed harness_fixed.c
 * Run:    ./harness_fixed
 */

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

/* Userspace stand-ins for kernel-only defines used by the transcribed code. */
#define ENOBUFS    105
#define ETIMEDOUT  110
#define ENOTSUP    91
typedef unsigned char boolean_t;
#ifndef TRUE
#define TRUE 1
#endif
#ifndef FALSE
#define FALSE 0
#endif

#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;
    int  label;
};

#define LIVE_MAGIC  0x4d42464d
#define DEAD_MAGIC  0xdeadc0de

static int alloc_count = 0, free_count = 0;
static int double_free_detected = 0, uaf_detected = 0;
static int live_allocs = 0;
static int fail_after = -1, alloc_budget = 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++; live_allocs++;
    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 (label=%d) at %s\n",
               (void*)m, m->label, 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 (label=%d)\n",
               (void*)m, m->label);
        double_free_detected++;
        return;
    }
    m->state = 0; m->magic = DEAD_MAGIC;
    free_count++; live_allocs--;
}
static void m_freem(struct mbuf *m)
{ while (m) { struct mbuf *n = m->m_next; m_free_one(m); m = n; } }

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);                        \
        (*(mp))->pkthdr_len += (plen);                     \
    } else {                                              \
        *(mp) = m_prepend_harness(*(mp), (plen));         \
        if (*(mp)) (*(mp))->pkthdr_len += (plen);         \
    }                                                     \
} while(0)

static struct mbuf *m_prepend_harness(struct mbuf *m, int len)
{
    struct mbuf *mn;
    if (fail_after >= 0 && alloc_budget >= fail_after) {
        printf("  [inject] m_prepend: alloc OOM -> m_freem(m=%p label=%d) + return NULL\n",
               (void*)m, m->label);
        m_freem(m);
        return NULL;
    }
    alloc_budget++;
    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;
    if (fail_after >= 0 && alloc_budget >= fail_after) {
        printf("  [inject] m_pullup: alloc OOM -> m_freem(n=%p label=%d) + return NULL\n",
               (void*)n, n->label);
        m_freem(n);
        return NULL;
    }
    alloc_budget++;
    {
        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_TTL(s)     (s & 0xff)
#define MPLS_STACK(s)   ((s >> 8) & 1)
#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, from_mpls, rt_flags;
};
#define RTF_MPLSOPS 0x100

static int label_counter = 100;

/* ================== THE FIXED FUNCTIONS ================== */

/* mpls_push — unchanged (already took struct mbuf **). */
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 (ENOBUFS);
    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;
    (*m)->label = ++label_counter;
    printf("  mpls_push: new head m=%p label=%d\n", (void*)(*m), (*m)->label);
    return 0;
}

/* mpls_swap — FIXED: takes struct mbuf **mp, writes *mp on rebind. */
static int mpls_swap(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 (ENOBUFS); }
        *mp = m;                       /* *** THE 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 (ETIMEDOUT);
        buf = 0;
        MPLS_SET_LABEL(buf, label);
        MPLS_SET_TTL(buf, ttl);
        p->mpls_shim = buf;
    }
    return 0;
}

/* mpls_pop — FIXED: takes struct mbuf **mp, writes *mp on rebind. */
static int mpls_pop(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 (ENOBUFS); }
        *mp = m;                       /* *** THE 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;
}

/* mpls_output — FIXED: takes struct mbuf **mp; out: label always writes *mp. */
static int mpls_output(struct mbuf **mp, struct rtentry *rt)
{
    int i, error = 0;
    int stackempty;
    int ttl = 255;
    struct mbuf *m = *mp;              /* *** THE FIX: local copy *** */

    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;       /* *** FIX: goto out (not return) *** */
            stackempty = 0;
            m->m_flags |= M_MPLSLABELED;
            break;
        case MPLSLOP_SWAP:
            if (stackempty) { error = ENOTSUP; goto out; }
            error = mpls_swap(&m, s->smpls_label);   /* *** FIX: &m *** */
            if (error) goto out;
            break;
        case MPLSLOP_POP:
            if (stackempty) { error = ENOTSUP; goto out; }
            { int sb; error = mpls_pop(&m, &sb); }   /* *** FIX: &m *** */
            if (error) goto out;
            break;
        default:
            error = ENOTSUP; goto out;
        }
    }
out:
    *mp = m;                            /* *** THE FIX: propagate head *** */
    return (error);
}

/* mpls_output_process — FIXED: takes struct mbuf **mp. */
static int mpls_output_process(struct mbuf **mp, struct rtentry *rt)
{
    int error;
    if (!(rt->rt_flags & RTF_MPLSOPS))
        return TRUE;
    error = mpls_output(mp, rt);        /* *** FIX: mp (not m) *** */
    if (error) {
        printf("  mpls_output_process: error=%d -> m_freem(*mp=%p label=%d)\n",
               error, (void*)(*mp), *mp ? (*mp)->label : -1);
        m_freem(*mp);                   /* *** FIX: *mp (current head) *** */
        *mp = NULL;
        return FALSE;
    }
    return TRUE;
}

static int if_output_deref_count = 0;
static int if_output_received_label = -1;
static int fake_if_output(struct mbuf *m)
{
    if_output_deref_count++;
    printf("  if_output: received m=%p label=%d\n", (void*)m, m ? m->label : -1);
    if (check_live(m, "if_output(m) deref")) {
        if_output_received_label = m->label;
        printf("  if_output: m LIVE, m_len=%d pkthdr.len=%d M_PKTHDR=%d "
               "M_MPLSLABELED=%d\n",
               m->m_len, m->pkthdr_len,
               (m->m_flags & M_PKTHDR)?1:0,
               (m->m_flags & M_MPLSLABELED)?1:0);
        m_freem(m);
    }
    return 0;
}

/* ip_output dispatch — FIXED: passes &m. */
static void ip_output_mpls_dispatch(struct mbuf *m, struct rtentry *rt)
{
    int cont;
    printf("  ip_output: m=%p label=%d (head)\n", (void*)m, m->label);
    cont = mpls_output_process(&m, rt);               /* *** FIX: &m *** */
    if (!cont) {
        printf("  ip_output: mpls_output_process=FALSE -> goto done\n");
        return;
    }
    printf("  ip_output: mpls_output_process=TRUE -> ifp->if_output(m=%p)\n",
           (void*)m);
    (void)fake_if_output(m);
}

static struct mbuf *make_ip_pkt(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));
    m->m_data[0] = 0x45; m->m_data[8] = 64;
    m->m_len = 20; m->pkthdr_len = 20;
    m->label = ++label_counter;
    return m;
}
static struct mbuf *make_mpls_frag_chain(void)
{
    struct mbuf *m = m_gethdr(0, 0);
    struct mbuf *m2 = mbuf_alloc(0);
    uint32_t shim = 0;
    MPLS_SET_LABEL(shim, 100);
    MPLS_SET_STACK(shim, 1);
    MPLS_SET_TTL(shim, 64);
    m->m_flags |= M_MPLSLABELED;
    m->m_len = 2;
    memcpy(m->m_data, &shim, 2);
    memcpy(m2->m_pktdat, ((char*)&shim) + 2, 2);
    m2->m_len = 32;
    memset(m2->m_pktdat + 2, 0xBB, 30);
    m2->m_data = m2->m_pktdat;
    m2->label = ++label_counter;
    m->m_next = m2;
    m->pkthdr_len = 34;
    m->label = ++label_counter;
    return m;
}

static void reset_accounting(void)
{
    alloc_count = free_count = 0;
    double_free_detected = uaf_detected = 0;
    live_allocs = 0;
    if_output_deref_count = 0;
    if_output_received_label = -1;
    fail_after = -1; alloc_budget = 0;
}

static void report(const char *label, const char *expect)
{
    printf("\n  --- %s ---\n", label);
    printf("  allocs=%d frees=%d  outstanding(live)=%d  double_free=%d  uaf=%d\n",
           alloc_count, free_count, live_allocs, double_free_detected, uaf_detected);
    printf("  if_output derefs=%d  last_received_label=%d\n",
           if_output_deref_count, if_output_received_label);
    printf("  EXPECT (fixed): %s\n", expect);
    if (double_free_detected == 0 && uaf_detected == 0 && live_allocs == 0)
        printf("  *** PASS — fix eliminates the manifestation ***\n");
    else
        printf("  *** FAIL — manifestation still present ***\n");
}

int main(void)
{
    struct rtentry rt_push, rt_swap, rt_push_then_swap;

    printf("================================================================\n");
    printf("DF-0754 FIXED harness — mpls_output(struct mbuf **mp)\n");
    printf("All M1/M2 scenarios should PASS (zero double-free/UAF/leak).\n");
    printf("================================================================\n\n");

    printf("----------------------------------------------------------------\n");
    printf("M1a (fixed): PUSH + m_prepend OOM\n");
    printf("----------------------------------------------------------------\n");
    reset_accounting();
    memset(&rt_push, 0, sizeof(rt_push));
    rt_push.from_mpls = 1; rt_push.rt_flags = RTF_MPLSOPS;
    rt_push.nshim = 1;
    rt_push.shim[0].smpls_op = MPLSLOP_PUSH; rt_push.shim[0].smpls_label = 999;
    { struct mbuf *m = make_ip_pkt(0); fail_after = 0;
      ip_output_mpls_dispatch(m, &rt_push); }
    report("M1a fixed", "double_free=0, uaf=0, live_allocs=0 (no stale deref)");

    printf("\n----------------------------------------------------------------\n");
    printf("M1b (fixed): SWAP + m_pullup OOM\n");
    printf("----------------------------------------------------------------\n");
    reset_accounting();
    memset(&rt_swap, 0, sizeof(rt_swap));
    rt_swap.from_mpls = 1; rt_swap.rt_flags = RTF_MPLSOPS;
    rt_swap.nshim = 1;
    rt_swap.shim[0].smpls_op = MPLSLOP_SWAP; rt_swap.shim[0].smpls_label = 200;
    { struct mbuf *m = make_mpls_frag_chain(); fail_after = 0;
      ip_output_mpls_dispatch(m, &rt_swap); }
    report("M1b fixed", "double_free=0, uaf=0, live_allocs=0");

    printf("\n----------------------------------------------------------------\n");
    printf("M1c (fixed): PUSH ok then SWAP TTL-expired\n");
    printf("----------------------------------------------------------------\n");
    reset_accounting();
    memset(&rt_push_then_swap, 0, sizeof(rt_push_then_swap));
    rt_push_then_swap.from_mpls = 1; rt_push_then_swap.rt_flags = RTF_MPLSOPS;
    rt_push_then_swap.nshim = 2;
    rt_push_then_swap.shim[0].smpls_op = MPLSLOP_PUSH; rt_push_then_swap.shim[0].smpls_label = 999;
    rt_push_then_swap.shim[1].smpls_op = MPLSLOP_SWAP; rt_push_then_swap.shim[1].smpls_label = 200;
    { struct mbuf *m = make_ip_pkt(0); fail_after = -1;
      ip_output_mpls_dispatch(m, &rt_push_then_swap); }
    report("M1c fixed", "double_free=0, live_allocs=0 (new head freed, not leaked)");

    printf("\n----------------------------------------------------------------\n");
    printf("M2a (fixed): PUSH realloc SUCCESS\n");
    printf("----------------------------------------------------------------\n");
    reset_accounting();
    { struct mbuf *m = make_ip_pkt(0); fail_after = -1;
      ip_output_mpls_dispatch(m, &rt_push); }
    report("M2a fixed", "if_output_received_label == pushed head label (CORRECT), live_allocs=0");

    printf("\n----------------------------------------------------------------\n");
    printf("M2b (fixed): SWAP m_pullup realloc SUCCESS\n");
    printf("----------------------------------------------------------------\n");
    reset_accounting();
    { struct mbuf *m = make_mpls_frag_chain(); fail_after = -1;
      ip_output_mpls_dispatch(m, &rt_swap); }
    report("M2b fixed", "uaf=0, live_allocs=0 (driver gets the new head)");

    printf("\n----------------------------------------------------------------\n");
    printf("CONTROL (fixed): PUSH leading_space=14\n");
    printf("----------------------------------------------------------------\n");
    reset_accounting();
    { struct mbuf *m = make_ip_pkt(ETHER_HDR_LEN); fail_after = -1;
      ip_output_mpls_dispatch(m, &rt_push); }
    report("CONTROL fixed", "clean (unchanged behavior)");

    printf("\n================================================================\n");
    printf("DF-0754 FIXED harness complete.\n");
    printf("All manifestations eliminated by mpls_output(struct mbuf **mp).\n");
    return 0;
}