/*      $NetBSD: t_inetd.c,v 1.3 2025/01/05 08:23:33 riastradh Exp $    */

/*-
* Copyright (c) 2021 The NetBSD Foundation, Inc.
* All rights reserved.
*
* This code is derived from software contributed to The NetBSD Foundation
* by James Browning, Gabe Coffland, Alex Gavin, and Solomon Ritzow.
*
* 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 NETBSD FOUNDATION, INC. 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 FOUNDATION 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/cdefs.h>
__RCSID("$NetBSD: t_inetd.c,v 1.3 2025/01/05 08:23:33 riastradh Exp $");

#include <atf-c.h>
#include <spawn.h>
#include <errno.h>
#include <stdlib.h>
#include <signal.h>
#include <stdio.h>
#include <string.h>
#include <err.h>
#include <sys/wait.h>
#include <sys/time.h>
#include <sys/socket.h>
#include <netdb.h>

#define CHECK_ERROR(expr) ATF_REQUIRE_MSG((expr) != -1,\
   "%s", strerror(errno))

#define TCP 6
#define UDP 17

static pid_t    run(const char *, char *const *);
static char     *concat(const char *restrict, const char *restrict);
static void     waitfor(pid_t, const char *);
static bool     run_udp_client(const char *);
static int      create_socket(const char *, const char *, int, int, time_t, struct sockaddr_storage *);
static bool     run_tcp_client(const char *);

/* This test should take around 5 to 7 seconds to complete. */
ATF_TC(test_ratelimit);

ATF_TC_HEAD(test_ratelimit, tc)
{
       atf_tc_set_md_var(tc, "descr", "Test inetd rate limiting values, "
       "uses UDP/TCP ports 5432-5439 with localhost.");
       /* Need to run as root so inetd can set uid */
       atf_tc_set_md_var(tc, "require.user", "root");
       atf_tc_set_md_var(tc, "require.progs", "inetd");
       /* Time out after 10 seconds, just in case */
       atf_tc_set_md_var(tc, "timeout", "10");
}

ATF_TC_BODY(test_ratelimit, tc)
{
       pid_t proc;

       /* Copy test server to relative path used in inetd_ratelimit.conf */
       atf_utils_copy_file(
           concat(atf_tc_get_config_var(tc, "srcdir"), "/test_server"),
           "test_server"
       );

       /* Run inetd in debug mode using specified config file */
       proc = run("inetd", (char* const []) {
           __UNCONST("inetd"), __UNCONST("-d"),
           concat(atf_tc_get_config_var(tc, "srcdir"),
           "/inetd_ratelimit.conf"),
           NULL
       });

       /* Wait for inetd to load services */
       sleep(1);

       /*
        * TODO test dgram/nowait? Specified in manpage but doesn't seem to
        * work
        */

       /* dgram/wait ip_max of 3, should receive these 3 responses */
       for (int i = 0; i < 3; i++) {
               ATF_REQUIRE(run_udp_client("5432"));
       }

       /* Rate limiting should prevent a response to this request */
       ATF_REQUIRE(!run_udp_client("5432"));

       /* dgram/wait ip_max of 0 */
       ATF_REQUIRE(!run_udp_client("5433"));

       /* dgram/wait service_max of 2 */
       ATF_REQUIRE(run_udp_client("5434"));
       ATF_REQUIRE(run_udp_client("5434"));
       ATF_REQUIRE(!run_udp_client("5434"));

       /* dgram/wait service_max of 0 */
       ATF_REQUIRE(!run_udp_client("5435"));

       /* stream/wait service_max of 2 */
       ATF_REQUIRE(run_tcp_client("5434"));
       ATF_REQUIRE(run_tcp_client("5434"));
       ATF_REQUIRE(!run_tcp_client("5434"));

       /* stream/wait service_max of 0 */
       ATF_REQUIRE(!run_tcp_client("5435"));

       /* stream/nowait ip_max of 3 */
       for (int i = 0; i < 3; i++) {
               ATF_REQUIRE(run_tcp_client("5436"));
       }
       ATF_REQUIRE(!run_tcp_client("5436"));

       /* stream/nowait ip_max of 0 */
       ATF_REQUIRE(!run_tcp_client("5437"));

       /* dgram/wait service_max of 2 */
       ATF_REQUIRE(run_tcp_client("5438"));
       ATF_REQUIRE(run_tcp_client("5438"));
       ATF_REQUIRE(!run_tcp_client("5438"));

       /* dgram/wait service_max of 0 */
       ATF_REQUIRE(!run_tcp_client("5439"));

       /* Exit inetd */
       CHECK_ERROR(kill(proc, SIGTERM));

       waitfor(proc, "inetd");
}

ATF_TP_ADD_TCS(tp)
{
       ATF_TP_ADD_TC(tp, test_ratelimit);

       return atf_no_error();
}

/* Return true if successfully received message, false if timeout */
static bool
run_udp_client(const char *port)
{
       char buffer[] = "test";
       struct sockaddr_storage addr;

       int udp = create_socket("127.0.0.1", port, SOCK_DGRAM, UDP, 1, &addr);

       CHECK_ERROR(sendto(udp, buffer, sizeof(buffer), 0,
           (struct sockaddr *)&addr, addr.ss_len));

       struct iovec iov = {
               .iov_base = buffer,
               .iov_len = sizeof(buffer)
       };

       struct msghdr msg = {
               .msg_name = &addr,
               .msg_namelen = addr.ss_len, /* is this correct? */
               .msg_iov = &iov,
               .msg_iovlen = 1
       };

       ssize_t count = recvmsg(udp, &msg, 0);
       if (count == -1) {
               if (errno == EAGAIN) {
                       /* Timed out, return false */
                       CHECK_ERROR(close(udp));
                       return false;
               } else {
                       /* All other errors fatal */
                       CHECK_ERROR(-1);
               }
       }
       CHECK_ERROR(close(udp));
       return true;
}

/* Run localhost tcp echo, return true if successful, false if timeout/disconnect */
static bool
run_tcp_client(const char *port)
{
       struct sockaddr_storage remote;
       ssize_t count;
       int tcp;
       char buffer[] = "test";

       tcp = create_socket("127.0.0.1", port, SOCK_STREAM, TCP, 1, &remote);
       CHECK_ERROR(connect(tcp, (const struct sockaddr *)&remote,
           remote.ss_len));
       CHECK_ERROR(send(tcp, buffer, sizeof(buffer), 0));
       count = recv(tcp, buffer, sizeof(buffer), 0);
       if (count == -1) {
               /*
                * Connection reset by peer indicates the connection was
                * dropped. EAGAIN indicates the timeout expired. Any other
                * error is unexpected for this client program test.
                */
               if(errno == ECONNRESET || errno == EAGAIN) {
                       return false;
               } else {
                       CHECK_ERROR(-1);
                       return false;
               }
       }

       if (count == 0) {
               /* socket was shutdown by inetd, no more data available */
               return false;
       }
       return true;
}

/*
* Create a socket with the characteristics inferred by the args, return parsed
* socket address in dst.
*/
static int
create_socket(const char *address, const char *port,
   int socktype, int proto, time_t timeout_sec, struct sockaddr_storage *dst)
{
       struct addrinfo hints = {
               .ai_flags = AI_NUMERICHOST,
               .ai_socktype = socktype,
               .ai_protocol = proto
       };
       struct addrinfo * res;
       int error, fd;

       ATF_REQUIRE_EQ_MSG(error = getaddrinfo(address, port, &hints, &res), 0,
           "%s", gai_strerror(error));

       /* Make sure there is only one possible bind address */
       ATF_REQUIRE_MSG(res->ai_next == NULL, "Ambiguous create_socket args");
       CHECK_ERROR(fd = socket(res->ai_family,
           res->ai_socktype, res->ai_protocol));
       struct timeval timeout = { timeout_sec, 0 };
       CHECK_ERROR(setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout,
           sizeof(timeout)));
       memcpy(dst, res->ai_addr, res->ai_addrlen);
       freeaddrinfo(res);
       return fd;
}

/* Run program with args */
static pid_t
run(const char *prog, char *const *args)
{
       pid_t proc;
       extern char **environ;
       ATF_REQUIRE_EQ(posix_spawnp(&proc, prog,
           NULL, NULL, args, environ), 0);
       return proc;
}

/* Wait for a process to exit, check return value */
static void
waitfor(pid_t pid, const char *taskname)
{
       int status;
       int rpid = waitpid(pid, &status, WALLSIG);
       ATF_REQUIRE_MSG(rpid == pid, "wait %d != %d %s",
           rpid, pid, strerror(errno));

       ATF_REQUIRE_EQ_MSG(WEXITSTATUS(status), EXIT_SUCCESS,
           "%s failed with "
           "exit status %d", taskname, WEXITSTATUS(status));
}

/* Concatenate two const strings, do not free arguments */
static char *
concat(const char *restrict left, const char *restrict right)
{
       char *res;
       if (asprintf(&res, "%s%s", left, right) == -1)
               return NULL;
       return res;
}