/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2025 Gleb Smirnoff <glebius@FreeBSD.org>
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include <sys/param.h>
#include <sys/jail.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <errno.h>
#include <ifaddrs.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <atf-c.h>

static const char buf[] = "Hello";

static void
sendtolocalhost(int s)
{
        struct sockaddr_in dst = {
                .sin_family = AF_INET,
                .sin_len = sizeof(struct sockaddr_in),
                .sin_addr = { htonl(INADDR_LOOPBACK) },
                .sin_port = htons(1638),
        };

        ATF_REQUIRE(sendto(s, buf, sizeof(buf), 0, (struct sockaddr *)&dst,
            sizeof(dst)) == sizeof(buf));
}

/*
 * Echo back to the sender its own address in payload.
 */
static void *
echo(void *arg)
{
        int s = *(int *)arg;
        struct sockaddr_in sin;
        socklen_t slen = sizeof(sin);
        char rbuf[sizeof(buf)];

        ATF_REQUIRE(recvfrom(s, &rbuf, sizeof(rbuf), 0, (struct sockaddr *)&sin,
            &slen) == sizeof(rbuf));
        printf("Echo to %s:%u\n", inet_ntoa(sin.sin_addr), ntohs(sin.sin_port));
        ATF_REQUIRE(sendto(s, &sin, sizeof(sin), 0, (struct sockaddr *)&sin,
            sizeof(sin)) == sizeof(sin));
        return (NULL);
}

/*
 * Cycle through local addresses (normally there should be at least two
 * different IPv4 ones), and communicate to the echo server checking both
 * IP_SENDSRCADDR and IP_RECVDSTADDR.  Use same cmsg buffer for both send
 * and receive operation, this is a suggested in manual, given that
 * IP_RECVDSTADDR == IP_SENDSRCADDR.
 * At the setup phase check that IP_SENDSRCADDR doesn't work on unbound socket.
 */
ATF_TC_WITHOUT_HEAD(IP_SENDSRCADDR);
ATF_TC_BODY(IP_SENDSRCADDR, tc)
{
        struct sockaddr_in srv = {
                .sin_family = AF_INET,
                .sin_len = sizeof(struct sockaddr_in),
        }, dst;
        char cbuf[CMSG_SPACE(sizeof(struct in_addr))];
        struct iovec iov = {
                .iov_base = __DECONST(char *, buf),
                .iov_len = sizeof(buf),
        };
        struct iovec riov = {
                .iov_base = &dst,
                .iov_len = sizeof(dst),
        };
        struct msghdr msg = {
                .msg_iov = &iov,
                .msg_iovlen = 1,
                .msg_name = &srv,
                .msg_namelen = sizeof(srv),
                .msg_control = cbuf,
                .msg_controllen = sizeof(cbuf),
        };
        struct msghdr rmsg = {
                .msg_iov = &riov,
                .msg_iovlen = 1,
                .msg_control = cbuf,
                .msg_controllen = sizeof(cbuf),
        };
        struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
        *cmsg = (struct cmsghdr) {
                .cmsg_level = IPPROTO_IP,
                .cmsg_type = IP_SENDSRCADDR,
                .cmsg_len = CMSG_LEN(sizeof(struct in_addr)),
        };
        socklen_t slen = sizeof(struct sockaddr_in);
        struct ifaddrs *ifa0, *ifa;
        pthread_t tid;
        int s, e;

        /* First check that IP_SENDSRCADDR doesn't work on an unbound socket. */
        ATF_REQUIRE((s = socket(PF_INET, SOCK_DGRAM, 0)) > 0);
        ATF_REQUIRE_MSG(sendmsg(s, &msg, 0) == -1 && errno == EINVAL,
            "sendmsg(.cmsg_type = IP_SENDSRCADDR), errno %d", errno);

        /* Bind to random ports both sender and echo server. */
        ATF_REQUIRE(bind(s, (struct sockaddr *)&srv, sizeof(srv)) == 0);
        ATF_REQUIRE((e = socket(PF_INET, SOCK_DGRAM, 0)) > 0);
        ATF_REQUIRE(bind(e, (struct sockaddr *)&srv, sizeof(srv)) == 0);
        ATF_REQUIRE(getsockname(e, (struct sockaddr *)&srv, &slen) == 0);
        srv.sin_addr.s_addr = htonl(INADDR_LOOPBACK);

        ATF_REQUIRE(getifaddrs(&ifa0) == 0);
        for (ifa = ifa0; ifa != NULL; ifa = ifa->ifa_next) {
                struct sockaddr_in src;
                struct in_addr vrf;

                if (ifa->ifa_addr->sa_family != AF_INET)
                        continue;
                memcpy(&src, ifa->ifa_addr, sizeof(src));
                printf("Sending from %s\n", inet_ntoa(src.sin_addr));
                ATF_REQUIRE(pthread_create(&tid, NULL, echo, &e) == 0);
                memcpy(CMSG_DATA(cmsg), &src.sin_addr, sizeof(src.sin_addr));
                ATF_REQUIRE(sendmsg(s, &msg, 0) == sizeof(buf));
                ATF_REQUIRE(recvmsg(s, &rmsg, 0) == sizeof(struct sockaddr_in));
                memcpy(&vrf, CMSG_DATA(cmsg), sizeof(vrf));
                ATF_REQUIRE_MSG(dst.sin_addr.s_addr == src.sin_addr.s_addr,
                    "Sent from %s, but echo server reports %s",
                    inet_ntoa(src.sin_addr), inet_ntoa(dst.sin_addr));
                ATF_REQUIRE_MSG(vrf.s_addr == src.sin_addr.s_addr,
                    "Sent from %s, but IP_RECVDSTADDR reports %s",
                    inet_ntoa(src.sin_addr), inet_ntoa(vrf));
                ATF_REQUIRE(pthread_join(tid, NULL) == 0);
        }

        freeifaddrs(ifa0);
        close(s);
        close(e);
}

/*
 * Check gethostname(2) on a newborn socket, and then on an unconnected, but
 * used socket.  The first shall return all-zeroes, and second one should
 * return us our assigned port.
 */
ATF_TC_WITHOUT_HEAD(gethostname);
ATF_TC_BODY(gethostname, tc)
{
        struct sockaddr_in sin;
        socklen_t slen = sizeof(sin);
        int s;

        ATF_REQUIRE((s = socket(PF_INET, SOCK_DGRAM, 0)) > 0);
        ATF_REQUIRE(getsockname(s, (struct sockaddr *)&sin, &slen) == 0);
        ATF_REQUIRE_MSG(sin.sin_addr.s_addr == INADDR_ANY && sin.sin_port == 0,
            "newborn socket name %s:%u", inet_ntoa(sin.sin_addr),
            ntohs(sin.sin_port));
        sendtolocalhost(s);
        ATF_REQUIRE(getsockname(s, (struct sockaddr *)&sin, &slen) == 0);
        ATF_REQUIRE_MSG(sin.sin_addr.s_addr == INADDR_ANY && sin.sin_port != 0,
            "used unconnected socket name %s:%u", inet_ntoa(sin.sin_addr),
            ntohs(sin.sin_port));
        close(s);
}

ATF_TC_WITHOUT_HEAD(gethostname_jailed);
ATF_TC_BODY(gethostname_jailed, tc)
{
        struct in_addr laddr = { htonl(INADDR_LOOPBACK) };
        struct jail jconf = {
                .version = JAIL_API_VERSION,
                .path = __DECONST(char *, "/"),
                .hostname = __DECONST(char *,"test"),
                .ip4s = 1,
                .ip4 = &laddr,
        };
        struct sockaddr_in sin;
        socklen_t slen = sizeof(sin);
        int s;

        ATF_REQUIRE(jail(&jconf) > 0);
        ATF_REQUIRE((s = socket(PF_INET, SOCK_DGRAM, 0)) > 0);
        sendtolocalhost(s);
        ATF_REQUIRE(getsockname(s, (struct sockaddr *)&sin, &slen) == 0);
        ATF_REQUIRE_MSG(sin.sin_addr.s_addr == laddr.s_addr &&
            sin.sin_port != 0,
            "jailed unconnected socket name %s:%u", inet_ntoa(sin.sin_addr),
            ntohs(sin.sin_port));
        close(s);
}

/*
 * See bug 274009.
 */
ATF_TC_WITHOUT_HEAD(v4mapped);
ATF_TC_BODY(v4mapped, tc)
{
        struct sockaddr_in6 sa6 = {
                .sin6_family = AF_INET6,
                .sin6_len = sizeof(struct sockaddr_in6),
                .sin6_port = htons(1),
        };
        int s;

        ATF_REQUIRE((s = socket(PF_INET6, SOCK_DGRAM, 0)) > 0);
        ATF_REQUIRE(setsockopt(s, IPPROTO_IPV6, IPV6_V6ONLY, &(int){0},
            sizeof(int)) == 0);
        ATF_REQUIRE(inet_pton(AF_INET6, "::ffff:127.0.0.1", &(sa6.sin6_addr))
            == 1);
        ATF_REQUIRE(sendto(s, buf, sizeof(buf), 0, (struct sockaddr *)&sa6,
            sizeof(sa6)) == sizeof(buf));
        close(s);
}

ATF_TP_ADD_TCS(tp)
{
        ATF_TP_ADD_TC(tp, v4mapped);
        ATF_TP_ADD_TC(tp, gethostname);
        ATF_TP_ADD_TC(tp, gethostname_jailed);
        ATF_TP_ADD_TC(tp, IP_SENDSRCADDR);

        return (atf_no_error());
}