#ifndef XDP_PARSER_H
#define XDP_PARSER_H 1

#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/udp.h>
#include <stdlib.h>

#ifndef VLAN_MAX_DEPTH
#define VLAN_MAX_DEPTH 2
#endif

#define VLAN_VID_MASK 0x0fff /* VLAN Identifier */

struct hdr_cursor {
  union {
    void *pos;
    __u8 *pos_u8;
  };
};

struct vlan_hdr {
  __be16 h_vlan_TCI;
  __be16 h_vlan_encapsulated_proto;
};

struct collect_vlans {
  __u16 id[VLAN_MAX_DEPTH];
};

static __always_inline int proto_is_vlan(__u16 h_proto) {
  return !!(h_proto == __constant_htons(ETH_P_8021Q) ||
            h_proto == __constant_htons(ETH_P_8021AD));
}

static __always_inline int parse_ethhdr_vlan(struct hdr_cursor *nh,
                                             void *data_end,
                                             struct ethhdr **ethhdr,
                                             struct collect_vlans *vlans) {
  struct ethhdr *eth = nh->pos;
  int hdrsize = sizeof(*eth);
  struct vlan_hdr *vlh;
  __u16 h_proto;
  int i;

  if (nh->pos + hdrsize > data_end)
    return -1;

  nh->pos += hdrsize;
  *ethhdr = eth;
  vlh = nh->pos;
  h_proto = eth->h_proto;

#pragma unroll
  for (i = 0; i < VLAN_MAX_DEPTH; i++) {
    if (!proto_is_vlan(h_proto))
      break;

    if (vlh + 1 > (struct vlan_hdr *)data_end)
      break;

    h_proto = vlh->h_vlan_encapsulated_proto;
    if (vlans) /* collect VLAN ids */
      vlans->id[i] = (__constant_ntohs(vlh->h_vlan_TCI) & VLAN_VID_MASK);

    vlh++;
  }

  nh->pos = vlh;
  return h_proto; /* network-byte-order */
}

static __always_inline int parse_ethhdr(struct hdr_cursor *nh, void *data_end,
                                        struct ethhdr **ethhdr) {
  return parse_ethhdr_vlan(nh, data_end, ethhdr, NULL);
}

static __always_inline int parse_ip6hdr(struct hdr_cursor *nh, void *data_end,
                                        struct ipv6hdr **ip6hdr) {
  struct ipv6hdr *ip6h = nh->pos;

  if (ip6h + 1 > (struct ipv6hdr *)data_end)
    return -1;

  nh->pos = ip6h + 1;
  *ip6hdr = ip6h;

  return ip6h->nexthdr;
}

static __always_inline int parse_iphdr(struct hdr_cursor *nh, void *data_end,
                                       struct iphdr **iphdr) {
  struct iphdr *iph = (struct iphdr *)nh->pos;
  size_t hdrsize;

  if (iph + 1 > (struct iphdr *)data_end)
    return -1;

  hdrsize = iph->ihl * 4;
  if (hdrsize < sizeof(*iph))
    return -1;

  /* Variable-length IPv4 header, need to use byte-based arithmetic */
  if (nh->pos + hdrsize > data_end)
    return -1;

  nh->pos += hdrsize;
  *iphdr = iph;

  return iph->protocol;
}

static __always_inline int parse_udphdr(struct hdr_cursor *nh, void *data_end,
                                        struct udphdr **udphdr) {
  int len;
  struct udphdr *h = nh->pos;

  if (h + 1 > (struct udphdr *)data_end)
    return -1;

  nh->pos = h + 1;
  *udphdr = h;

  len = __constant_ntohs(h->len) - sizeof(struct udphdr);
  if (len < 0)
    return -1;

  return len;
}

#endif