// SPDX-License-Identifier: GPL-2.0
// Copyright (c) 2018 Facebook
// Copyright (c) 2019 Cloudflare
// Copyright (c) 2020 Isovalent, Inc.
/*
 * Test that the socket assign program is able to redirect traffic towards a
 * socket, regardless of whether the port or address destination of the traffic
 * matches the port.
 */

#define _GNU_SOURCE
#include <fcntl.h>
#include <signal.h>
#include <stdlib.h>
#include <unistd.h>

#include "test_progs.h"
#include "network_helpers.h"

#define BIND_PORT 1234
#define CONNECT_PORT 4321
#define TEST_DADDR (0xC0A80203)
#define NS_SELF "/proc/self/ns/net"
#define SERVER_MAP_PATH "/sys/fs/bpf/tc/globals/server_map"

static int stop, duration;

static bool
configure_stack(void)
{
        char tc_version[128];
        char tc_cmd[BUFSIZ];
        char *prog;
        FILE *tc;

        /* Check whether tc is built with libbpf. */
        tc = popen("tc -V", "r");
        if (CHECK_FAIL(!tc))
                return false;
        if (CHECK_FAIL(!fgets(tc_version, sizeof(tc_version), tc))) {
                pclose(tc);
                return false;
        }
        if (strstr(tc_version, ", libbpf "))
                prog = "test_sk_assign_libbpf.bpf.o";
        else
                prog = "test_sk_assign.bpf.o";
        if (CHECK_FAIL(pclose(tc)))
                return false;

        /* Move to a new networking namespace */
        if (CHECK_FAIL(unshare(CLONE_NEWNET)))
                return false;

        /* Configure necessary links, routes */
        if (CHECK_FAIL(system("ip link set dev lo up")))
                return false;
        if (CHECK_FAIL(system("ip route add local default dev lo")))
                return false;
        if (CHECK_FAIL(system("ip -6 route add local default dev lo")))
                return false;

        /* Load qdisc, BPF program */
        if (CHECK_FAIL(system("tc qdisc add dev lo clsact")))
                return false;
        sprintf(tc_cmd, "%s %s %s %s %s", "tc filter add dev lo ingress bpf",
                       "direct-action object-file", prog,
                       "section tc",
                       (env.verbosity < VERBOSE_VERY) ? " 2>/dev/null" : "verbose");
        if (CHECK(system(tc_cmd), "BPF load failed;",
                  "run with -vv for more info\n"))
                return false;

        return true;
}

static in_port_t
get_port(int fd)
{
        struct sockaddr_storage ss;
        socklen_t slen = sizeof(ss);
        in_port_t port = 0;

        if (CHECK_FAIL(getsockname(fd, (struct sockaddr *)&ss, &slen)))
                return port;

        switch (ss.ss_family) {
        case AF_INET:
                port = ((struct sockaddr_in *)&ss)->sin_port;
                break;
        case AF_INET6:
                port = ((struct sockaddr_in6 *)&ss)->sin6_port;
                break;
        default:
                CHECK(1, "Invalid address family", "%d\n", ss.ss_family);
        }
        return port;
}

static ssize_t
rcv_msg(int srv_client, int type)
{
        char buf[BUFSIZ];

        if (type == SOCK_STREAM)
                return read(srv_client, &buf, sizeof(buf));
        else
                return recvfrom(srv_client, &buf, sizeof(buf), 0, NULL, NULL);
}

static int
run_test(int server_fd, const struct sockaddr *addr, socklen_t len, int type)
{
        int client = -1, srv_client = -1;
        char buf[] = "testing";
        in_port_t port;
        int ret = 1;

        client = connect_to_addr(type, (struct sockaddr_storage *)addr, len, NULL);
        if (client == -1) {
                perror("Cannot connect to server");
                goto out;
        }

        if (type == SOCK_STREAM) {
                srv_client = accept(server_fd, NULL, NULL);
                if (CHECK_FAIL(srv_client == -1)) {
                        perror("Can't accept connection");
                        goto out;
                }
        } else {
                srv_client = server_fd;
        }
        if (CHECK_FAIL(write(client, buf, sizeof(buf)) != sizeof(buf))) {
                perror("Can't write on client");
                goto out;
        }
        if (CHECK_FAIL(rcv_msg(srv_client, type) != sizeof(buf))) {
                perror("Can't read on server");
                goto out;
        }

        port = get_port(srv_client);
        if (CHECK_FAIL(!port))
                goto out;
        /* SOCK_STREAM is connected via accept(), so the server's local address
         * will be the CONNECT_PORT rather than the BIND port that corresponds
         * to the listen socket. SOCK_DGRAM on the other hand is connectionless
         * so we can't really do the same check there; the server doesn't ever
         * create a socket with CONNECT_PORT.
         */
        if (type == SOCK_STREAM &&
            CHECK(port != htons(CONNECT_PORT), "Expected", "port %u but got %u",
                  CONNECT_PORT, ntohs(port)))
                goto out;
        else if (type == SOCK_DGRAM &&
                 CHECK(port != htons(BIND_PORT), "Expected",
                       "port %u but got %u", BIND_PORT, ntohs(port)))
                goto out;

        ret = 0;
out:
        close(client);
        if (srv_client != server_fd)
                close(srv_client);
        if (ret)
                WRITE_ONCE(stop, 1);
        return ret;
}

static void
prepare_addr(struct sockaddr *addr, int family, __u16 port, bool rewrite_addr)
{
        struct sockaddr_in *addr4;
        struct sockaddr_in6 *addr6;

        switch (family) {
        case AF_INET:
                addr4 = (struct sockaddr_in *)addr;
                memset(addr4, 0, sizeof(*addr4));
                addr4->sin_family = family;
                addr4->sin_port = htons(port);
                if (rewrite_addr)
                        addr4->sin_addr.s_addr = htonl(TEST_DADDR);
                else
                        addr4->sin_addr.s_addr = htonl(INADDR_LOOPBACK);
                break;
        case AF_INET6:
                addr6 = (struct sockaddr_in6 *)addr;
                memset(addr6, 0, sizeof(*addr6));
                addr6->sin6_family = family;
                addr6->sin6_port = htons(port);
                addr6->sin6_addr = in6addr_loopback;
                if (rewrite_addr)
                        addr6->sin6_addr.s6_addr32[3] = htonl(TEST_DADDR);
                break;
        default:
                fprintf(stderr, "Invalid family %d", family);
        }
}

struct test_sk_cfg {
        const char *name;
        int family;
        struct sockaddr *addr;
        socklen_t len;
        int type;
        bool rewrite_addr;
};

#define TEST(NAME, FAMILY, TYPE, REWRITE)                               \
{                                                                       \
        .name = NAME,                                                   \
        .family = FAMILY,                                               \
        .addr = (FAMILY == AF_INET) ? (struct sockaddr *)&addr4         \
                                    : (struct sockaddr *)&addr6,        \
        .len = (FAMILY == AF_INET) ? sizeof(addr4) : sizeof(addr6),     \
        .type = TYPE,                                                   \
        .rewrite_addr = REWRITE,                                        \
}

void test_sk_assign(void)
{
        struct sockaddr_in addr4;
        struct sockaddr_in6 addr6;
        struct test_sk_cfg tests[] = {
                TEST("ipv4 tcp port redir", AF_INET, SOCK_STREAM, false),
                TEST("ipv4 tcp addr redir", AF_INET, SOCK_STREAM, true),
                TEST("ipv6 tcp port redir", AF_INET6, SOCK_STREAM, false),
                TEST("ipv6 tcp addr redir", AF_INET6, SOCK_STREAM, true),
                TEST("ipv4 udp port redir", AF_INET, SOCK_DGRAM, false),
                TEST("ipv4 udp addr redir", AF_INET, SOCK_DGRAM, true),
                TEST("ipv6 udp port redir", AF_INET6, SOCK_DGRAM, false),
                TEST("ipv6 udp addr redir", AF_INET6, SOCK_DGRAM, true),
        };
        __s64 server = -1;
        int server_map;
        int self_net;
        int i;

        self_net = open(NS_SELF, O_RDONLY);
        if (CHECK_FAIL(self_net < 0)) {
                perror("Unable to open "NS_SELF);
                return;
        }

        if (!configure_stack()) {
                perror("configure_stack");
                goto cleanup;
        }

        server_map = bpf_obj_get(SERVER_MAP_PATH);
        if (CHECK_FAIL(server_map < 0)) {
                perror("Unable to open " SERVER_MAP_PATH);
                goto cleanup;
        }

        for (i = 0; i < ARRAY_SIZE(tests) && !READ_ONCE(stop); i++) {
                struct test_sk_cfg *test = &tests[i];
                const struct sockaddr *addr;
                const int zero = 0;
                int err;

                if (!test__start_subtest(test->name))
                        continue;
                prepare_addr(test->addr, test->family, BIND_PORT, false);
                addr = (const struct sockaddr *)test->addr;
                server = start_server_addr(test->type,
                                           (const struct sockaddr_storage *)addr,
                                           test->len, NULL);
                if (server == -1)
                        goto close;

                err = bpf_map_update_elem(server_map, &zero, &server, BPF_ANY);
                if (CHECK_FAIL(err)) {
                        perror("Unable to update server_map");
                        goto close;
                }

                /* connect to unbound ports */
                prepare_addr(test->addr, test->family, CONNECT_PORT,
                             test->rewrite_addr);
                if (run_test(server, addr, test->len, test->type))
                        goto close;

                close(server);
                server = -1;
        }

close:
        close(server);
        close(server_map);
cleanup:
        if (CHECK_FAIL(unlink(SERVER_MAP_PATH)))
                perror("Unable to unlink " SERVER_MAP_PATH);
        if (CHECK_FAIL(setns(self_net, CLONE_NEWNET)))
                perror("Failed to setns("NS_SELF")");
        close(self_net);
}