#define _GNU_SOURCE
#include <errno.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <ifaddrs.h>
#include <net/if.h>
#include "__netlink.h"

/* getifaddrs() uses PF_PACKET to relay hardware addresses.
 * But Infiniband socket address length is longer, so use this hack
 * (like glibc) to return it anyway. */
struct sockaddr_ll_hack {
	unsigned short sll_family, sll_protocol;
	int sll_ifindex;
	unsigned short sll_hatype;
	unsigned char sll_pkttype, sll_halen;
	unsigned char sll_addr[24];
};

union sockany {
	struct sockaddr sa;
	struct sockaddr_ll_hack ll;
	struct sockaddr_in v4;
	struct sockaddr_in6 v6;
};

struct ifaddrs_storage {
	struct ifaddrs ifa;
	struct ifaddrs_storage *hash_next;
	union sockany addr, netmask, ifu;
	unsigned int index;
	char name[IFNAMSIZ+1];
};

#define IFADDRS_HASH_SIZE 64
struct ifaddrs_ctx {
	struct ifaddrs_storage *first;
	struct ifaddrs_storage *last;
	struct ifaddrs_storage *hash[IFADDRS_HASH_SIZE];
};

void freeifaddrs(struct ifaddrs *ifp)
{
	struct ifaddrs *n;
	while (ifp) {
		n = ifp->ifa_next;
		free(ifp);
		ifp = n;
	}
}

static void addifaddrs(struct ifaddrs_ctx *ctx, struct ifaddrs_storage *add)
{
	if (!add->ifa.ifa_name) {
		free(add);
		return;
	}
	if (!ctx->first) ctx->first = add;
	if (ctx->last) ctx->last->ifa.ifa_next = &add->ifa;
	ctx->last = add;
}

static struct sockaddr* copy_lladdr(union sockany *sa, struct rtattr *rta, struct ifinfomsg *ifi)
{
	if (RTA_DATALEN(rta) > sizeof(sa->ll.sll_addr)) return 0;
	sa->ll.sll_family = AF_PACKET;
	sa->ll.sll_ifindex = ifi->ifi_index;
	sa->ll.sll_hatype = ifi->ifi_type;
	sa->ll.sll_halen = RTA_DATALEN(rta);
	memcpy(sa->ll.sll_addr, RTA_DATA(rta), RTA_DATALEN(rta));
	return &sa->sa;
}

static uint8_t *sockany_addr(int af, union sockany *sa, int *len)
{
	switch (af) {
	case AF_INET: *len = 4; return (uint8_t*) &sa->v4.sin_addr;
	case AF_INET6: *len = 16; return (uint8_t*) &sa->v6.sin6_addr;
	}
	return 0;
}

static struct sockaddr* copy_addr(int af, union sockany *sa, struct rtattr *rta)
{
	int len;
	uint8_t *dst = sockany_addr(af, sa, &len);
	if (!dst || RTA_DATALEN(rta) != len) return 0;
	sa->sa.sa_family = af;
	memcpy(dst, RTA_DATA(rta), len);
	return &sa->sa;
}

static struct sockaddr *gen_netmask(int af, union sockany *sa, int prefixlen)
{
	int maxlen, i;
	uint8_t *dst = sockany_addr(af, sa, &maxlen);
	if (!dst) return 0;
	sa->sa.sa_family = af;
	if (prefixlen > 8*maxlen) prefixlen = 8*maxlen;
	i = prefixlen / 8;
	memset(dst, 0xff, i);
	if (i<maxlen) {
		dst[i++] = 0xff << (8 - (prefixlen % 8));
		if (i<maxlen) memset(&dst[i+1], 0x00, maxlen-i);
	}
	return &sa->sa;
}

static int __handle_link(void *pctx, struct nlmsghdr *h)
{
	struct ifaddrs_ctx *ctx = pctx;
	struct ifaddrs_storage *ifs;
	struct ifinfomsg *ifi = NLMSG_DATA(h);
	struct rtattr *rta;
	int stats_len = 0;

	for (rta = NLMSG_RTA(h, sizeof(*ifi)); NLMSG_RTAOK(rta, h); rta = RTA_NEXT(rta)) {
		if (rta->rta_type != IFLA_STATS) continue;
		stats_len = RTA_DATALEN(rta);
		break;
	}

	ifs = calloc(1, sizeof(struct ifaddrs_storage) + stats_len);
	if (ifs == 0) return -1;

	ifs->index = ifi->ifi_index;
	ifs->ifa.ifa_flags = ifi->ifi_flags;

	for (rta = NLMSG_RTA(h, sizeof(*ifi)); NLMSG_RTAOK(rta, h); rta = RTA_NEXT(rta)) {
		switch (rta->rta_type) {
		case IFLA_IFNAME:
			if (RTA_DATALEN(rta) <= IFNAMSIZ) {
				strncpy(ifs->name, RTA_DATA(rta), RTA_DATALEN(rta));
				ifs->ifa.ifa_name = ifs->name;
			}
			break;
		case IFLA_ADDRESS:
			ifs->ifa.ifa_addr = copy_lladdr(&ifs->addr, rta, ifi);
			break;
		case IFLA_BROADCAST:
			ifs->ifa.ifa_broadaddr = copy_lladdr(&ifs->ifu, rta, ifi);
			break;
		case IFLA_STATS:
			ifs->ifa.ifa_data = (void*)(ifs+1);
			memcpy(ifs->ifa.ifa_data, RTA_DATA(rta), RTA_DATALEN(rta));
			break;
		}
	}
	if (ifs->ifa.ifa_name) {
		ifs->hash_next = ctx->hash[ifs->index%IFADDRS_HASH_SIZE];
		ctx->hash[ifs->index%IFADDRS_HASH_SIZE] = ifs;
	}
	addifaddrs(ctx, ifs);
	return 0;
}

static int __handle_addr(void *pctx, struct nlmsghdr *h)
{
	struct ifaddrs_ctx *ctx = pctx;
	struct ifaddrs_storage *ifs, *ifs0;
	struct ifaddrmsg *ifa = NLMSG_DATA(h);
	struct rtattr *rta;

	ifs = calloc(1, sizeof(struct ifaddrs_storage));
	if (ifs == 0) return -1;

	for (ifs0 = ctx->hash[ifa->ifa_index%IFADDRS_HASH_SIZE]; ifs0; ifs0 = ifs0->hash_next)
		if (ifs0->index == ifa->ifa_index)
			break;
	if (!ifs0) return 0;

	ifs->ifa.ifa_name = ifs0->ifa.ifa_name;
	ifs->ifa.ifa_flags = ifs0->ifa.ifa_flags;
	for (rta = NLMSG_RTA(h, sizeof(*ifa)); NLMSG_RTAOK(rta, h); rta = RTA_NEXT(rta)) {
		switch (rta->rta_type) {
		case IFA_ADDRESS:
			ifs->ifa.ifa_addr = copy_addr(ifa->ifa_family, &ifs->addr, rta);
			if (ifs->ifa.ifa_addr)
				ifs->ifa.ifa_netmask = gen_netmask(ifa->ifa_family, &ifs->netmask, ifa->ifa_prefixlen);
			break;
		case IFA_BROADCAST:
			/* For point-to-point links this is peer, but ifa_broadaddr
			 * and ifa_dstaddr are union, so this works for both.  */
			ifs->ifa.ifa_broadaddr = copy_addr(ifa->ifa_family, &ifs->ifu, rta);
			break;
		}
	}

	addifaddrs(ctx, ifs);
	return 0;
}

int getifaddrs(struct ifaddrs **ifap)
{
	struct ifaddrs_ctx _ctx, *ctx = &_ctx;
	struct __netlink_handle *nh;
	int r = 0;

	nh = __netlink_open(NETLINK_ROUTE);
	if (!nh) return -1;
	memset(ctx, 0, sizeof(*ctx));
	if (__netlink_enumerate(nh, RTM_GETLINK, __handle_link, ctx)) r = -1;
	if (__netlink_enumerate(nh, RTM_GETADDR, __handle_addr, ctx)) r = -1;
	__netlink_close(nh);
	if (r == 0) *ifap = &ctx->first->ifa;
	else freeifaddrs(&ctx->first->ifa);
	return r;
}