// SPDX-License-Identifier: GPL-2.0

#include <arpa/inet.h>
#include <error.h>
#include <linux/errqueue.h>
#include <linux/icmp.h>
#include <linux/icmpv6.h>
#include <linux/in6.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <netinet/in.h>
#include <netinet/udp.h>
#include <poll.h>
#include <sched.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/ioctl.h>
#include <sys/socket.h>

#include "../kselftest_harness.h"

static const unsigned short src_port = 44444;
static const unsigned short dst_port = 55555;
static const int min_orig_dgram_len = 128;
static const int min_payload_len_v4 =
        min_orig_dgram_len - sizeof(struct iphdr) - sizeof(struct udphdr);
static const int min_payload_len_v6 =
        min_orig_dgram_len - sizeof(struct ipv6hdr) - sizeof(struct udphdr);
static const uint8_t orig_payload_byte =  0xAA;

struct sockaddr_inet {
        union {
                struct sockaddr_in6 v6;
                struct sockaddr_in v4;
                struct sockaddr sa;
        };
        socklen_t len;
};

struct ip_case_info {
        int     domain;
        int     level;
        int     opt1;
        int     opt2;
        int     proto;
        int     (*build_func)(uint8_t *buf, ssize_t buflen, bool with_ext,
                              int payload_len, bool bad_csum, bool bad_len,
                              bool smaller_len);
        int     min_payload;
};

static int bringup_loopback(void)
{
        struct ifreq ifr = {
                .ifr_name = "lo"
        };
        int fd;

        fd = socket(AF_INET, SOCK_DGRAM, 0);
        if (fd < 0)
                return -1;

        if (ioctl(fd, SIOCGIFFLAGS, &ifr) < 0)
                goto err;

        ifr.ifr_flags = ifr.ifr_flags | IFF_UP;

        if (ioctl(fd, SIOCSIFFLAGS, &ifr) < 0)
                goto err;

        close(fd);
        return 0;

err:
        close(fd);
        return -1;
}

static uint16_t csum(const void *buf, size_t len)
{
        const uint8_t *data = buf;
        uint32_t sum = 0;

        while (len > 1) {
                sum += (data[0] << 8) | data[1];
                data += 2;
                len -= 2;
        }

        if (len == 1)
                sum += data[0] << 8;

        while (sum >> 16)
                sum = (sum & 0xFFFF) + (sum >> 16);

        return ~sum & 0xFFFF;
}

static int poll_err(int fd)
{
        struct pollfd pfd;

        memset(&pfd, 0, sizeof(pfd));
        pfd.fd = fd;

        if (poll(&pfd, 1, 5000) != 1 || pfd.revents != POLLERR)
                return -1;

        return 0;
}

static void set_addr(struct sockaddr_inet *addr, int domain,
                     unsigned short port)
{
        memset(addr, 0, sizeof(*addr));

        switch (domain) {
        case AF_INET:
                addr->v4.sin_family = AF_INET;
                addr->v4.sin_port = htons(port);
                addr->v4.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
                addr->len = sizeof(addr->v4);
                break;
        case AF_INET6:
                addr->v6.sin6_family = AF_INET6;
                addr->v6.sin6_port = htons(port);
                addr->v6.sin6_addr = in6addr_loopback;
                addr->len = sizeof(addr->v6);
                break;
        }
}

static int bind_and_setsockopt(int fd, const struct ip_case_info *info)
{
        struct sockaddr_inet addr;
        int opt = 1;

        set_addr(&addr, info->domain, src_port);

        if (setsockopt(fd, info->level, info->opt1, &opt, sizeof(opt)) < 0)
                return -1;

        if (setsockopt(fd, info->level, info->opt2, &opt, sizeof(opt)) < 0)
                return -1;

        return bind(fd, &addr.sa, addr.len);
}

static int build_rfc4884_ext(uint8_t *buf, size_t buflen, bool bad_csum,
                             bool bad_len, bool smaller_len)
{
        struct icmp_extobj_hdr *objh;
        struct icmp_ext_hdr *exthdr;
        size_t obj_len, ext_len;
        uint16_t sum;

        /* Use an object payload of 4 bytes */
        obj_len = sizeof(*objh) + sizeof(uint32_t);
        ext_len = sizeof(*exthdr) + obj_len;

        if (ext_len > buflen)
                return -EINVAL;

        exthdr = (struct icmp_ext_hdr *)buf;
        objh = (struct icmp_extobj_hdr *)(buf + sizeof(*exthdr));

        exthdr->version = 2;
        /* When encoding a bad object length, either encode a length too small
         * to fit the object header or too big to fit in the packet.
         */
        if (bad_len)
                obj_len = smaller_len ? sizeof(*objh) - 1 : obj_len * 2;
        objh->length = htons(obj_len);

        sum = csum(buf, ext_len);
        exthdr->checksum = htons(bad_csum ? sum - 1 : sum);

        return ext_len;
}

static int build_orig_dgram_v4(uint8_t *buf, ssize_t buflen, int payload_len)
{
        struct udphdr *udph;
        struct iphdr *iph;
        size_t len = 0;

        len = sizeof(*iph) + sizeof(*udph) + payload_len;
        if (len > buflen)
                return -EINVAL;

        iph = (struct iphdr *)buf;
        udph = (struct udphdr *)(buf + sizeof(*iph));

        iph->version = 4;
        iph->ihl = 5;
        iph->protocol = IPPROTO_UDP;
        iph->saddr = htonl(INADDR_LOOPBACK);
        iph->daddr = htonl(INADDR_LOOPBACK);
        iph->tot_len = htons(len);
        iph->check = htons(csum(iph, sizeof(*iph)));

        udph->source = htons(src_port);
        udph->dest = htons(dst_port);
        udph->len = htons(sizeof(*udph) + payload_len);

        memset(buf + sizeof(*iph) + sizeof(*udph), orig_payload_byte,
               payload_len);

        return len;
}

static int build_orig_dgram_v6(uint8_t *buf, ssize_t buflen, int payload_len)
{
        struct udphdr *udph;
        struct ipv6hdr *iph;
        size_t len = 0;

        len = sizeof(*iph) + sizeof(*udph) + payload_len;
        if (len > buflen)
                return -EINVAL;

        iph = (struct ipv6hdr *)buf;
        udph = (struct udphdr *)(buf + sizeof(*iph));

        iph->version = 6;
        iph->payload_len = htons(sizeof(*udph) + payload_len);
        iph->nexthdr = IPPROTO_UDP;
        iph->saddr = in6addr_loopback;
        iph->daddr = in6addr_loopback;

        udph->source = htons(src_port);
        udph->dest = htons(dst_port);
        udph->len = htons(sizeof(*udph) + payload_len);

        memset(buf + sizeof(*iph) + sizeof(*udph), orig_payload_byte,
               payload_len);

        return len;
}

static int build_icmpv4_pkt(uint8_t *buf, ssize_t buflen, bool with_ext,
                            int payload_len, bool bad_csum, bool bad_len,
                            bool smaller_len)
{
        struct icmphdr *icmph;
        int len, ret;

        len = sizeof(*icmph);
        memset(buf, 0, buflen);

        icmph = (struct icmphdr *)buf;
        icmph->type = ICMP_DEST_UNREACH;
        icmph->code = ICMP_PORT_UNREACH;
        icmph->checksum = 0;

        ret = build_orig_dgram_v4(buf + len, buflen - len, payload_len);
        if (ret < 0)
                return ret;

        len += ret;

        icmph->un.reserved[1] = (len - sizeof(*icmph)) / sizeof(uint32_t);

        if (with_ext) {
                ret = build_rfc4884_ext(buf + len, buflen - len,
                                        bad_csum, bad_len, smaller_len);
                if (ret < 0)
                        return ret;

                len += ret;
        }

        icmph->checksum = htons(csum(icmph, len));
        return len;
}

static int build_icmpv6_pkt(uint8_t *buf, ssize_t buflen, bool with_ext,
                            int payload_len, bool bad_csum, bool bad_len,
                            bool smaller_len)
{
        struct icmp6hdr *icmph;
        int len, ret;

        len = sizeof(*icmph);
        memset(buf, 0, buflen);

        icmph = (struct icmp6hdr *)buf;
        icmph->icmp6_type = ICMPV6_DEST_UNREACH;
        icmph->icmp6_code = ICMPV6_PORT_UNREACH;
        icmph->icmp6_cksum = 0;

        ret = build_orig_dgram_v6(buf + len, buflen - len, payload_len);
        if (ret < 0)
                return ret;

        len += ret;

        icmph->icmp6_datagram_len = (len - sizeof(*icmph)) / sizeof(uint64_t);

        if (with_ext) {
                ret = build_rfc4884_ext(buf + len, buflen - len,
                                        bad_csum, bad_len, smaller_len);
                if (ret < 0)
                        return ret;

                len += ret;
        }

        icmph->icmp6_cksum = htons(csum(icmph, len));
        return len;
}

FIXTURE(rfc4884) {};

FIXTURE_SETUP(rfc4884)
{
        int ret;

        ret = unshare(CLONE_NEWNET);
        ASSERT_EQ(ret, 0) {
                TH_LOG("unshare(CLONE_NEWNET) failed: %s", strerror(errno));
        }

        ret = bringup_loopback();
        ASSERT_EQ(ret, 0) TH_LOG("Failed to bring up loopback interface");
}

FIXTURE_TEARDOWN(rfc4884)
{
}

const struct ip_case_info ipv4_info = {
        .domain         = AF_INET,
        .level          = SOL_IP,
        .opt1           = IP_RECVERR,
        .opt2           = IP_RECVERR_RFC4884,
        .proto          = IPPROTO_ICMP,
        .build_func     = build_icmpv4_pkt,
        .min_payload    = min_payload_len_v4,
};

const struct ip_case_info ipv6_info = {
        .domain         = AF_INET6,
        .level          = SOL_IPV6,
        .opt1           = IPV6_RECVERR,
        .opt2           = IPV6_RECVERR_RFC4884,
        .proto          = IPPROTO_ICMPV6,
        .build_func     = build_icmpv6_pkt,
        .min_payload    = min_payload_len_v6,
};

FIXTURE_VARIANT(rfc4884) {
        /* IPv4/v6 related information */
        struct ip_case_info     info;
        /* Whether to append an ICMP extension or not */
        bool                    with_ext;
        /* UDP payload length */
        int                     payload_len;
        /* Whether to generate a bad checksum in the ICMP extension structure */
        bool                    bad_csum;
        /* Whether to generate a bad length in the ICMP object header */
        bool                    bad_len;
        /* Whether it is too small to fit the object header or too big to fit
         * in the packet
         */
        bool                    smaller_len;
};

/* Tests that a valid ICMPv4 error message with extension and the original
 * datagram is smaller than 128 bytes, generates an error with zero offset,
 * and does not raise the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv4_ext_small_payload) {
        .info           = ipv4_info,
        .with_ext       = true,
        .payload_len    = 64,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that a valid ICMPv4 error message with extension and 128 bytes original
 * datagram, generates an error with the expected offset, and does not raise the
 * SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv4_ext) {
        .info           = ipv4_info,
        .with_ext       = true,
        .payload_len    = min_payload_len_v4,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that a valid ICMPv4 error message with extension and the original
 * datagram is larger than 128 bytes, generates an error with the expected
 * offset, and does not raise the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv4_ext_large_payload) {
        .info           = ipv4_info,
        .with_ext       = true,
        .payload_len    = 256,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that a valid ICMPv4 error message without extension and the original
 * datagram is smaller than 128 bytes, generates an error with zero offset,
 * and does not raise the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv4_no_ext_small_payload) {
        .info           = ipv4_info,
        .with_ext       = false,
        .payload_len    = 64,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that a valid ICMPv4 error message without extension and 128 bytes
 * original datagram, generates an error with zero offset, and does not raise
 * the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv4_no_ext_min_payload) {
        .info           = ipv4_info,
        .with_ext       = false,
        .payload_len    = min_payload_len_v4,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that a valid ICMPv4 error message without extension and the original
 * datagram is larger than 128 bytes, generates an error with zero offset,
 * and does not raise the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv4_no_ext_large_payload) {
        .info           = ipv4_info,
        .with_ext       = false,
        .payload_len    = 256,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that an ICMPv4 error message with extension and an invalid checksum,
 * generates an error with the expected offset, and raises the
 * SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv4_invalid_ext_checksum) {
        .info           = ipv4_info,
        .with_ext       = true,
        .payload_len    = min_payload_len_v4,
        .bad_csum       = true,
        .bad_len        = false,
};

/* Tests that an ICMPv4 error message with extension and an object length
 * smaller than the object header, generates an error with the expected offset,
 * and raises the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv4_invalid_ext_length_small) {
        .info           = ipv4_info,
        .with_ext       = true,
        .payload_len    = min_payload_len_v4,
        .bad_csum       = false,
        .bad_len        = true,
        .smaller_len    = true,
};

/* Tests that an ICMPv4 error message with extension and an object length that
 * is too big to fit in the packet, generates an error with the expected offset,
 * and raises the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv4_invalid_ext_length_large) {
        .info           = ipv4_info,
        .with_ext       = true,
        .payload_len    = min_payload_len_v4,
        .bad_csum       = false,
        .bad_len        = true,
        .smaller_len    = false,
};

/* Tests that a valid ICMPv6 error message with extension and the original
 * datagram is smaller than 128 bytes, generates an error with zero offset,
 * and does not raise the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv6_ext_small_payload) {
        .info           = ipv6_info,
        .with_ext       = true,
        .payload_len    = 64,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that a valid ICMPv6 error message with extension and 128 bytes original
 * datagram, generates an error with the expected offset, and does not raise the
 * SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv6_ext) {
        .info           = ipv6_info,
        .with_ext       = true,
        .payload_len    = min_payload_len_v6,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that a valid ICMPv6 error message with extension and the original
 * datagram is larger than 128 bytes, generates an error with the expected
 * offset, and does not raise the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv6_ext_large_payload) {
        .info           = ipv6_info,
        .with_ext       = true,
        .payload_len    = 256,
        .bad_csum       = false,
        .bad_len        = false,
};
/* Tests that a valid ICMPv6 error message without extension and the original
 * datagram is smaller than 128 bytes, generates an error with zero offset,
 * and does not raise the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv6_no_ext_small_payload) {
        .info           = ipv6_info,
        .with_ext       = false,
        .payload_len    = 64,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that a valid ICMPv6 error message without extension and 128 bytes
 * original datagram, generates an error with zero offset, and does not
 * raise the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv6_no_ext_min_payload) {
        .info           = ipv6_info,
        .with_ext       = false,
        .payload_len    = min_payload_len_v6,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that a valid ICMPv6 error message without extension and the original
 * datagram is larger than 128 bytes, generates an error with zero offset,
 * and does not raise the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv6_no_ext_large_payload) {
        .info           = ipv6_info,
        .with_ext       = false,
        .payload_len    = 256,
        .bad_csum       = false,
        .bad_len        = false,
};

/* Tests that an ICMPv6 error message with extension and an invalid checksum,
 * generates an error with the expected offset, and raises the
 * SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv6_invalid_ext_checksum) {
        .info           = ipv6_info,
        .with_ext       = true,
        .payload_len    = min_payload_len_v6,
        .bad_csum       = true,
        .bad_len        = false,
};

/* Tests that an ICMPv6 error message with extension and an object length
 * smaller than the object header, generates an error with the expected offset,
 * and raises the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv6_invalid_ext_length_small) {
        .info           = ipv6_info,
        .with_ext       = true,
        .payload_len    = min_payload_len_v6,
        .bad_csum       = false,
        .bad_len        = true,
        .smaller_len    = true,
};

/* Tests that an ICMPv6 error message with extension and an object length that
 * is too big to fit in the packet, generates an error with the expected offset,
 * and raises the SO_EE_RFC4884_FLAG_INVALID flag.
 */
FIXTURE_VARIANT_ADD(rfc4884, ipv6_invalid_ext_length_large) {
        .info           = ipv6_info,
        .with_ext       = true,
        .payload_len    = min_payload_len_v6,
        .bad_csum       = false,
        .bad_len        = true,
        .smaller_len    = false,
};

static void
check_rfc4884_offset(struct __test_metadata *_metadata, int sock,
                     const FIXTURE_VARIANT(rfc4884) *v)
{
        char rxbuf[1024];
        char ctrl[1024];
        struct iovec iov = {
                .iov_base = rxbuf,
                .iov_len = sizeof(rxbuf)
        };
        struct msghdr msg = {
                .msg_iov = &iov,
                .msg_iovlen = 1,
                .msg_control = ctrl,
                .msg_controllen = sizeof(ctrl),
        };
        struct cmsghdr *cmsg;
        int recv;

        ASSERT_EQ(poll_err(sock), 0);

        recv = recvmsg(sock, &msg, MSG_ERRQUEUE);
        ASSERT_GE(recv, 0) TH_LOG("recvmsg(MSG_ERRQUEUE) failed");

        for (cmsg = CMSG_FIRSTHDR(&msg); cmsg; cmsg = CMSG_NXTHDR(&msg, cmsg)) {
                bool is_invalid, expected_invalid;
                struct sock_extended_err *ee;
                int expected_off;
                uint16_t off;

                if (cmsg->cmsg_level != v->info.level ||
                    cmsg->cmsg_type != v->info.opt1) {
                        TH_LOG("Unrelated cmsgs were encountered in recvmsg()");
                        continue;
                }

                ee = (struct sock_extended_err *)CMSG_DATA(cmsg);
                off = ee->ee_rfc4884.len;
                is_invalid = ee->ee_rfc4884.flags & SO_EE_RFC4884_FLAG_INVALID;

                expected_invalid = v->bad_csum || v->bad_len;
                ASSERT_EQ(is_invalid, expected_invalid) {
                        TH_LOG("Expected invalidity flag to be %d, but got %d",
                               expected_invalid, is_invalid);
                }

                expected_off =
                        (v->with_ext && v->payload_len >= v->info.min_payload) ?
                        v->payload_len : 0;
                ASSERT_EQ(off, expected_off) {
                        TH_LOG("Expected RFC4884 offset %u, got %u",
                               expected_off, off);
                }
                break;
        }
}

TEST_F(rfc4884, rfc4884)
{
        const typeof(variant) v = variant;
        struct sockaddr_inet addr;
        uint8_t pkt[1024];
        int dgram, raw;
        int len, sent;
        int err;

        dgram = socket(v->info.domain, SOCK_DGRAM, 0);
        ASSERT_GE(dgram, 0) TH_LOG("Opening datagram socket failed");

        err = bind_and_setsockopt(dgram, &v->info);
        ASSERT_EQ(err, 0) TH_LOG("Bind failed");

        raw = socket(v->info.domain, SOCK_RAW, v->info.proto);
        ASSERT_GE(raw, 0) TH_LOG("Opening raw socket failed");

        len = v->info.build_func(pkt, sizeof(pkt), v->with_ext, v->payload_len,
                                 v->bad_csum, v->bad_len, v->smaller_len);
        ASSERT_GT(len, 0) TH_LOG("Building packet failed");

        set_addr(&addr, v->info.domain, 0);
        sent = sendto(raw, pkt, len, 0, &addr.sa, addr.len);
        ASSERT_EQ(len, sent) TH_LOG("Sending packet failed");

        check_rfc4884_offset(_metadata, dgram, v);

        close(dgram);
        close(raw);
}

TEST_HARNESS_MAIN