#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/ioctl.h>
#include <net/if.h>
#include <netinet/ether.h>
#include <arpa/inet.h>
#include <linux/if_packet.h>
#include <sys/mman.h>
#include <unistd.h>
#include <assert.h>

#define MTU                   1500
#define FRAME_MEM_SIZE        ( 4 * 1024 * 1024 )
#define BLOCK_SIZE            ( 4 * 1024 )

typedef uint8_t u8;
typedef uint16_t u16;
typedef uint32_t u32;

struct ring {
    int sock;
    int type;
    int version;
    void ** frames;
    size_t frame_num;
    size_t frame_count;
    void * mmap_base;
    size_t mmap_len;
    size_t frame_length;
    struct tpacket_req tpacket_req;
};

struct llc_header {
    u8 dsap;
    u8 ssap;
    u16 control;
};

static inline void get_own_mac(const char * if_name, u8 mac[6]);
static inline void setup_ring(struct ring * ring, int sock, int type);
static inline void send_broadcast_packet(struct ring * tx_ring, size_t payload_size, u8 sender_mac[6]);
static inline u32 round_up_to_pow_2(u32 x);

int main (int argc, char **argv) {

    if (argc < 3) {
        printf("usage: %s <interface name> <payload size>\n", argv[0]);
        exit(EXIT_FAILURE);
    }

    const char * if_name = argv[1];
    u8 own_mac[6];
    get_own_mac(argv[1], own_mac);

    char * endptr = NULL;
    size_t payload_size = strtol(argv[2], &endptr, 0);
    /* Ensure a number was parsed */
    assert(endptr != argv[2]);

    /* Open a raw socket */
    int sock = socket(PF_PACKET, SOCK_RAW, htons(ETH_P_ALL));
    assert(sock != -1);
    /* Use TPACKET_V2 */
    int ver = TPACKET_V2;
    (void) setsockopt(sock, SOL_PACKET, PACKET_VERSION, &ver, sizeof(ver));
    struct sockaddr_ll sa = {
        .sll_family = PF_PACKET,
        .sll_protocol = htons(ETH_P_ALL),
        .sll_ifindex = if_nametoindex(if_name)
    };
    /* Bind the socket to a network device */
    (void) bind(sock, (const struct sockaddr *) &sa, sizeof(sa));

    /* Set up the rings in the kernel */
    struct ring tx_ring;
    struct ring rx_ring;
    setup_ring(&tx_ring, sock, PACKET_TX_RING);
    setup_ring(&rx_ring, sock, PACKET_RX_RING);

    /* Map both rings into userspace in one mmap() call */
    size_t mmap_len = rx_ring.mmap_len + tx_ring.mmap_len;
    void * mmap_base = mmap(NULL, mmap_len, PROT_READ | PROT_WRITE, \
        MAP_SHARED | MAP_LOCKED | MAP_POPULATE, sock, 0);
    assert(mmap_base != MAP_FAILED);

    /* RX ring is first in the mapped memory, followed by the TX ring */
    rx_ring.mmap_base = mmap_base;
    tx_ring.mmap_base = mmap_base + rx_ring.mmap_len;

    /* Set up the pointers */
    for (size_t i = 0; i < rx_ring.frame_count; i++)
        rx_ring.frames[i] = rx_ring.mmap_base + (i * rx_ring.frame_length);
    for (size_t i = 0; i < tx_ring.frame_count; i++)
        tx_ring.frames[i] = tx_ring.mmap_base + (i * tx_ring.frame_length);

    send_broadcast_packet(&tx_ring, payload_size, own_mac);

    return 0;
}

static inline void get_own_mac(const char * if_name, u8 mac[6]) {

    /* Open any socket as a channel to the kernel */
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    assert(sock != -1);
    /* Get the HW address */
    struct ifreq ifr;
    (void) strncpy(ifr.ifr_name, if_name, sizeof(ifr.ifr_name) - 1);
    ifr.ifr_name[sizeof(ifr.ifr_name) - 1] = '\0';
    (void) ioctl(sock, SIOCGIFHWADDR, &ifr);
    (void) close(sock);
    (void) memcpy(mac, ifr.ifr_hwaddr.sa_data, 6);
}

static inline void setup_ring(struct ring * ring, int sock, int type) {

    ring->sock = sock;
    ring->type = type;
    ring->version = TPACKET_V2;
    u32 frame_size = round_up_to_pow_2(MTU + TPACKET_HDRLEN + TPACKET_ALIGNMENT);
    u32 block_size = BLOCK_SIZE;
    if (frame_size > block_size)
        block_size = frame_size;
    u32 block_count = FRAME_MEM_SIZE / block_size;
    u32 frame_count = (block_size / frame_size) * block_count;
    u32 ring_size = frame_count * sizeof(void *);
    /* Allocate an array of pointers to the frames in shared buffer */
    ring->frames = malloc(ring_size);
    assert(ring->frames);

    ring->tpacket_req.tp_block_size = block_size;
    ring->tpacket_req.tp_block_nr = block_count;
    ring->tpacket_req.tp_frame_size = frame_size;
    ring->tpacket_req.tp_frame_nr = frame_count;

    ring->mmap_len = block_size * block_count;
    ring->frame_count = frame_count;
    ring->frame_length = frame_size;
    ring->frame_num = 0;

    (void) setsockopt(sock, SOL_PACKET, type, &ring->tpacket_req, sizeof(ring->tpacket_req));
}

static inline void send_broadcast_packet(struct ring * tx_ring, size_t payload_size, u8 sender_mac[6]) {

    struct tpacket2_hdr * tp_header = tx_ring->frames[tx_ring->frame_num];
    /* Only three lower bits encode the TX status (assert there is space in the ring) */
    u32 tp_status = tp_header->tp_status & 0x7;
    assert(tp_status == TP_STATUS_AVAILABLE);
    u8 * buf = (u8 *)(void *)tp_header + TPACKET2_HDRLEN - \
        sizeof(struct sockaddr_ll);

    /* Fill in the MAC addresses */
    struct ether_header * eth = (struct ether_header *) buf;
    (void) memset(eth->ether_dhost, 0xFF, 6);
    (void) memcpy(eth->ether_shost, sender_mac, 6);
    eth->ether_type = htons(sizeof(struct llc_header) + payload_size);
    /* Fill in the LLC header */
    struct llc_header * llc = (struct llc_header *)(eth + 1);
    llc->dsap = 0;
    llc->ssap = 0;
    llc->control = 0;
    /* Poison the payload */
    u8 * payload = (u8 *)(llc + 1);
    (void) memset(payload, 0xAA, payload_size);

    size_t total_len = sizeof(*eth) + sizeof(*llc) + payload_size;
    tp_header->tp_len = total_len;
    /* Relinquish control of the frame to the kernel */
    tp_header->tp_status = TP_STATUS_SEND_REQUEST;

    /* Ping the kernel to send the packet */
    ssize_t ret = send(tx_ring->sock, NULL, 0, MSG_DONTWAIT);
    assert(ret == total_len);
    tx_ring->frame_num = (tx_ring->frame_num + 1) % tx_ring->frame_count;
}

static inline u32 round_up_to_pow_2(u32 x) {

    u32 i;
    for (i = 1; i < x; i <<= 1) {
        ;
    }
    return i;
}