/*
 * mdns-repeater.c — Lightweight IPv4/IPv6 mDNS reflector with health check
 *
 * Listens on all specified interfaces for mDNS packets (port 5353,
 * multicast groups 224.0.0.251 and ff02::fb) and repeats them
 * to every other interface.
 *
 * Writes a heartbeat to /tmp/mdns-health for container health checks.
 *
 * Build:
 *   gcc -static -O2 -o mdns-repeater mdns-repeater.c -lpthread
 *
 * Usage:
 *   mdns-repeater eth0 eth1 eth2 ...
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <pthread.h>
#include <time.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/select.h>
#include <sys/stat.h>
#include <net/if.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define MDNS_PORT      5353
#define MDNS_GROUP_IP  "224.0.0.251"
#define MDNS_GROUP_IP6 "ff02::fb"
#define PACKET_SZ      4096
#define MAX_IFACES     32
#define HEALTH_FILE    "/tmp/mdns-health"

/* ── Per-interface state ────────────────────────────────────── */
typedef struct {
    char         name[IFNAMSIZ];
    int          index;
    int          fd4;        /* IPv4 socket */
    int          fd6;        /* IPv6 socket */
    uint64_t     pkts_in;    /* health: packets received */
    uint64_t     pkts_out;   /* health: packets forwarded */
} iface_t;

static iface_t  ifaces[MAX_IFACES];
static int      n_ifaces = 0;
static volatile int running = 1;

/* ── Signal handler ─────────────────────────────────────────── */
static void handle_signal(int sig) {
    (void)sig;
    running = 0;
}

/* ── Write health heartbeat ─────────────────────────────────── */
static void write_health(void) {
    FILE *f = fopen(HEALTH_FILE, "w");
    if (!f) return;

    time_t now = time(NULL);
    uint64_t total_in = 0, total_out = 0;

    fprintf(f, "ts=%ld\n", (long)now);
    for (int i = 0; i < n_ifaces; i++) {
        fprintf(f, "iface.%s.pkts_in=%lu\n",
                ifaces[i].name, (unsigned long)ifaces[i].pkts_in);
        fprintf(f, "iface.%s.pkts_out=%lu\n",
                ifaces[i].name, (unsigned long)ifaces[i].pkts_out);
        total_in  += ifaces[i].pkts_in;
        total_out += ifaces[i].pkts_out;
    }
    fprintf(f, "total.pkts_in=%lu\n", (unsigned long)total_in);
    fprintf(f, "total.pkts_out=%lu\n", (unsigned long)total_out);
    fprintf(f, "n_ifaces=%d\n", n_ifaces);
    fprintf(f, "alive=1\n");
    fclose(f);
}

/* ── Heartbeat thread: writes health file every 10 seconds ──── */
static void *heartbeat_thread(void *arg) {
    (void)arg;
    write_health(); /* initial write */
    while (running) {
        sleep(10);
        write_health();
    }
    return NULL;
}

/* ── Join IPv4 mDNS group on an fd ──────────────────────────── */
static int join_mcast4(int fd, int ifindex) {
    struct ip_mreqn mreq = {0};
    inet_pton(AF_INET, MDNS_GROUP_IP, &mreq.imr_multiaddr);
    mreq.imr_ifindex = ifindex;
    if (setsockopt(fd, IPPROTO_IP, IP_ADD_MEMBERSHIP, &mreq, sizeof(mreq)) < 0) {
        fprintf(stderr, "IP_ADD_MEMBERSHIP(%s): %s\n", MDNS_GROUP_IP, strerror(errno));
        return -1;
    }
    return 0;
}

/* ── Join IPv6 mDNS group on an fd ──────────────────────────── */
static int join_mcast6(int fd, int ifindex) {
    struct ipv6_mreq mreq = {0};
    inet_pton(AF_INET6, MDNS_GROUP_IP6, &mreq.ipv6mr_multiaddr);
    mreq.ipv6mr_interface = ifindex;
    if (setsockopt(fd, IPPROTO_IPV6, IPV6_JOIN_GROUP, &mreq, sizeof(mreq)) < 0) {
        fprintf(stderr, "IPV6_JOIN_GROUP(%s): %s\n", MDNS_GROUP_IP6, strerror(errno));
        return -1;
    }
    return 0;
}

/* ── Set common socket options for mDNS ─────────────────────── */
static int setup_sockopts(int fd, int arr_idx, int sys_ifindex, int af) {
    int on = 1;
    int off = 0;

    /* Reuse address (allow multiple listeners) */
    if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) < 0)
        perror("SO_REUSEADDR");
#if defined(SO_REUSEPORT)
    if (setsockopt(fd, SOL_SOCKET, SO_REUSEPORT, &on, sizeof(on)) < 0)
        perror("SO_REUSEPORT");
#endif

    /* Bind to interface so packets from wrong interface are dropped */
    if (setsockopt(fd, SOL_SOCKET, SO_BINDTODEVICE, ifaces[arr_idx].name,
                   strlen(ifaces[arr_idx].name)) < 0)
        perror("SO_BINDTODEVICE");

    if (af == AF_INET) {
        if (setsockopt(fd, IPPROTO_IP, IP_MULTICAST_LOOP, &off, sizeof(off)) < 0)
            perror("IP_MULTICAST_LOOP");
        if (setsockopt(fd, IPPROTO_IP, IP_MULTICAST_IF, &sys_ifindex, sizeof(sys_ifindex)) < 0)
            perror("IP_MULTICAST_IF");
    } else {
        if (setsockopt(fd, IPPROTO_IPV6, IPV6_MULTICAST_LOOP, &off, sizeof(off)) < 0)
            perror("IPV6_MULTICAST_LOOP");
        if (setsockopt(fd, IPPROTO_IPV6, IPV6_MULTICAST_IF, &sys_ifindex, sizeof(sys_ifindex)) < 0)
            perror("IPV6_MULTICAST_IF");
    }
    return 0;
}

/* ── Create and bind an mDNS socket ─────────────────────────── */
static int create_mdns_socket(int af, int arr_idx) {
    struct sockaddr_storage addr = {0};
    socklen_t addrlen;
    int sys_ifindex = ifaces[arr_idx].index;
    int fd;

    fd = socket(af, SOCK_DGRAM, 0);
    if (fd < 0) {
        perror("socket");
        return -1;
    }

    setup_sockopts(fd, arr_idx, sys_ifindex, af);

    if (af == AF_INET) {
        struct sockaddr_in *sa = (struct sockaddr_in *)&addr;
        sa->sin_family      = AF_INET;
        sa->sin_port        = htons(MDNS_PORT);
        sa->sin_addr.s_addr = INADDR_ANY;
        addrlen = sizeof(struct sockaddr_in);
    } else {
        struct sockaddr_in6 *sa6 = (struct sockaddr_in6 *)&addr;
        sa6->sin6_family    = AF_INET6;
        sa6->sin6_port      = htons(MDNS_PORT);
        sa6->sin6_addr      = in6addr_any;
        sa6->sin6_scope_id  = sys_ifindex;
        addrlen = sizeof(struct sockaddr_in6);
    }

    if (bind(fd, (struct sockaddr *)&addr, addrlen) < 0) {
        fprintf(stderr, "bind(%s): %s\n", ifaces[arr_idx].name, strerror(errno));
        close(fd);
        return -1;
    }

    if (af == AF_INET) {
        if (join_mcast4(fd, sys_ifindex) < 0) { close(fd); return -1; }
    } else {
        if (join_mcast6(fd, sys_ifindex) < 0) { close(fd); return -1; }
    }

    return fd;
}

/* ── Send a packet to all interfaces except the source ──────── */
static void repeat(const char *buf, int len, int src_idx, int af) {
    struct msghdr   msg     = {0};
    struct iovec    iov[1];
    char            cmsg_buf[CMSG_SPACE(sizeof(struct in_pktinfo))];

    iov[0].iov_base = (void *)buf;
    iov[0].iov_len  = len;
    msg.msg_iov     = iov;
    msg.msg_iovlen  = 1;

    for (int i = 0; i < n_ifaces; i++) {
        if (i == src_idx) continue;

        int fd = (af == AF_INET) ? ifaces[i].fd4 : ifaces[i].fd6;
        if (fd < 0) continue;

        if (af == AF_INET) {
            struct in_pktinfo pktinfo = {0};
            pktinfo.ipi_ifindex = ifaces[i].index;
            inet_pton(AF_INET, MDNS_GROUP_IP, &pktinfo.ipi_spec_dst);

            msg.msg_control    = cmsg_buf;
            msg.msg_controllen = sizeof(cmsg_buf);

            struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
            cmsg->cmsg_level = IPPROTO_IP;
            cmsg->cmsg_type  = IP_PKTINFO;
            cmsg->cmsg_len   = CMSG_LEN(sizeof(struct in_pktinfo));
            memcpy(CMSG_DATA(cmsg), &pktinfo, sizeof(pktinfo));
            msg.msg_controllen = cmsg->cmsg_len;

            msg.msg_name    = NULL;
            msg.msg_namelen = 0;
        } else {
            struct sockaddr_in6 dst = {0};
            dst.sin6_family   = AF_INET6;
            dst.sin6_port     = htons(MDNS_PORT);
            inet_pton(AF_INET6, MDNS_GROUP_IP6, &dst.sin6_addr);
            dst.sin6_scope_id = ifaces[i].index;

            msg.msg_name    = &dst;
            msg.msg_namelen = sizeof(dst);
            msg.msg_control    = NULL;
            msg.msg_controllen = 0;
        }

        ssize_t sent = sendmsg(fd, &msg, 0);
        if (sent < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
            fprintf(stderr, "sendmsg(%s): %s\n", ifaces[i].name, strerror(errno));
        } else {
            ifaces[i].pkts_out++;
        }
    }
}

/* ── Reader thread: read packets from one interface ─────────── */
static void *reader_thread(void *arg) {
    int idx = (int)(intptr_t)arg;
    iface_t *iface = &ifaces[idx];
    char buf[PACKET_SZ];
    int maxfd = (iface->fd4 > iface->fd6) ? iface->fd4 : iface->fd6;

    while (running) {
        fd_set rfds;
        FD_ZERO(&rfds);
        if (iface->fd4 >= 0) FD_SET(iface->fd4, &rfds);
        if (iface->fd6 >= 0) FD_SET(iface->fd6, &rfds);

        struct timeval tv = {1, 0};
        int ret = select(maxfd + 1, &rfds, NULL, NULL, &tv);
        if (ret < 0) {
            if (errno == EINTR) continue;
            perror("select");
            break;
        }
        if (ret == 0) continue;

        if (iface->fd4 >= 0 && FD_ISSET(iface->fd4, &rfds)) {
            int len = recv(iface->fd4, buf, sizeof(buf), 0);
            if (len > 0) {
                iface->pkts_in++;
                repeat(buf, len, idx, AF_INET);
            }
        }
        if (iface->fd6 >= 0 && FD_ISSET(iface->fd6, &rfds)) {
            int len = recv(iface->fd6, buf, sizeof(buf), 0);
            if (len > 0) {
                iface->pkts_in++;
                repeat(buf, len, idx, AF_INET6);
            }
        }
    }
    return NULL;
}

/* ── Print usage ────────────────────────────────────────────── */
static void usage(const char *prog) {
    fprintf(stderr,
        "Usage: %s <interface> [interface ...]\n"
        "\n"
        "Reflect mDNS packets (IPv4 + IPv6) between all listed interfaces.\n"
        "Writes health status to " HEALTH_FILE " every 10s.\n"
        "\n"
        "Example:\n"
        "  %s eth0 eth1 br-lan\n",
        prog, prog);
}

/* ── Main ────────────────────────────────────────────────────── */
int main(int argc, char **argv) {
    if (argc < 2) {
        usage(argv[0]);
        return 1;
    }

    if (argc - 1 > MAX_IFACES) {
        fprintf(stderr, "Too many interfaces (max %d)\n", MAX_IFACES);
        return 1;
    }

    printf("[*] mdns-repeater starting...\n");

    /* Parse interfaces */
    for (int i = 1; i < argc; i++) {
        strncpy(ifaces[n_ifaces].name, argv[i], IFNAMSIZ - 1);
        ifaces[n_ifaces].name[IFNAMSIZ - 1] = '\0';

        ifaces[n_ifaces].index = if_nametoindex(argv[i]);
        if (ifaces[n_ifaces].index == 0) {
            fprintf(stderr, "Unknown interface: %s\n", argv[i]);
            return 1;
        }

        ifaces[n_ifaces].fd4 = create_mdns_socket(AF_INET, n_ifaces);
        if (ifaces[n_ifaces].fd4 < 0)
            fprintf(stderr, "WARNING: Could not create IPv4 mDNS socket on %s\n", argv[i]);

        ifaces[n_ifaces].fd6 = create_mdns_socket(AF_INET6, n_ifaces);
        if (ifaces[n_ifaces].fd6 < 0)
            fprintf(stderr, "WARNING: Could not create IPv6 mDNS socket on %s\n", argv[i]);

        if (ifaces[n_ifaces].fd4 < 0 && ifaces[n_ifaces].fd6 < 0) {
            fprintf(stderr, "FATAL: No sockets created for %s\n", argv[i]);
            return 1;
        }

        ifaces[n_ifaces].pkts_in  = 0;
        ifaces[n_ifaces].pkts_out = 0;

        printf("[+] %-16s (idx=%3d)  fd4=%d  fd6=%d\n",
               ifaces[n_ifaces].name, ifaces[n_ifaces].index,
               ifaces[n_ifaces].fd4, ifaces[n_ifaces].fd6);
        n_ifaces++;
    }

    /* Install signal handlers */
    signal(SIGINT,  handle_signal);
    signal(SIGTERM, handle_signal);
    signal(SIGQUIT, handle_signal);

    /* Start heartbeat thread */
    pthread_t hb_thread;
    pthread_create(&hb_thread, NULL, heartbeat_thread, NULL);

    /* Start one reader thread per interface */
    pthread_t threads[MAX_IFACES];
    for (int i = 0; i < n_ifaces; i++) {
        if (pthread_create(&threads[i], NULL, reader_thread, (void*)(intptr_t)i) != 0) {
            perror("pthread_create");
            return 1;
        }
    }

    printf("[*] Reflecting mDNS between %d interfaces (IPv4 + IPv6)\n", n_ifaces);

    /* Wait for all reader threads */
    for (int i = 0; i < n_ifaces; i++) {
        pthread_join(threads[i], NULL);
    }

    /* Stop heartbeat */
    running = 0;
    pthread_join(hb_thread, NULL);

    /* Cleanup */
    for (int i = 0; i < n_ifaces; i++) {
        if (ifaces[i].fd4 >= 0) close(ifaces[i].fd4);
        if (ifaces[i].fd6 >= 0) close(ifaces[i].fd6);
    }
    unlink(HEALTH_FILE);

    printf("[*] Shutdown complete.\n");
    return 0;
}