/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2026 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 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 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/ioctl.h>
#include <sys/socket.h>
#include <sys/sysctl.h>
#include <sys/mman.h>
#include <netinet/in.h>
#include <fcntl.h>
#include <stdlib.h>
#include <poll.h>
#include <pthread.h>

#include <openssl/ssl.h>
#include <openssl/err.h>
#include <openssl/bio.h>
#include <openssl/x509v3.h>

#include <atf-c.h>

#define FSIZE   (size_t)(2 * 1024 * 1024)

struct ctx {
        EVP_PKEY *pkey;         /* Self-signed key ... */
        X509 *cert;             /* ... and certificate */
        SSL_CTX *ctx;           /* client context */
        SSL *cln;               /* client connection */
        SSL *srv;               /* server connection */
        int cs;                 /* client socket */
        int ss;                 /* server socket */
        int fd;                 /* test file descriptor */
        void *mfd;              /* mapped contents of the test file */
        uint16_t port;          /* server listening port */
        pthread_t thr;          /* server thread */
        off_t offset;           /* SSL_sendfile offset */
        size_t size;            /* SSL_sendfile size */
        bool nb;                /* SSL_sendfile mode */
        ossl_ssize_t sbytes;    /* SSL_sendfile returned sbytes */
        enum {
                INIT,
                READY,
                RUNNING,
                EXITING,
        } state;
        pthread_mutex_t mtx;
        pthread_cond_t cv;
};

static void *server_thread(void *arg);

static void
common_init(struct ctx *c)
{
        char hostname[sizeof("localhost:65536")];
        char tempname[] = "/tmp/ssl_sendfile_test.XXXXXXXXXX";
        X509_NAME *name;
        X509_EXTENSION *ext;
        SSL *ssl;
        bool enable;
        size_t len = sizeof(enable);

        if (sysctlbyname("kern.ipc.tls.enable", &enable, &len, NULL, 0) == -1) {
                if (errno == ENOENT)
                        atf_tc_skip("kernel does not have options KERN_TLS");
                atf_libc_error(errno, "Failed to read kern.ipc.tls.enable");
        }
        if (!enable)
                atf_tc_skip("kern.ipc.tls.enable is off");

        c->state = INIT;

        /*
         * Generate self signed key & certificate.
         */
        SSL_library_init();
        OpenSSL_add_all_algorithms();
        SSL_load_error_strings();
        ATF_REQUIRE(c->pkey = EVP_RSA_gen(2048));
        ATF_REQUIRE(c->cert = X509_new());
        ASN1_INTEGER_set(X509_get_serialNumber(c->cert), 1);
        X509_set_version(c->cert, 2);
        X509_gmtime_adj(X509_get_notBefore(c->cert), 0);
        X509_gmtime_adj(X509_get_notAfter(c->cert), 60L*60*24*365);
        X509_set_pubkey(c->cert, c->pkey);
        name = X509_get_subject_name(c->cert);
        X509_NAME_add_entry_by_txt(name, "CN", MBSTRING_ASC,
            (unsigned char *)"localhost", -1, -1, 0);
        X509_set_issuer_name(c->cert, name);
        ext = X509V3_EXT_conf_nid(NULL, NULL, NID_basic_constraints,
            "critical,CA:FALSE");
        X509_add_ext(c->cert, ext, -1);
        X509_EXTENSION_free(ext);
        ATF_REQUIRE(X509_sign(c->cert, c->pkey, EVP_sha256()) > 0);

        /*
         * Create random filled file with memory mapping.
         */
        ATF_REQUIRE((c->fd = mkstemp(tempname)) > 0);
        ATF_REQUIRE(unlink(tempname) == 0);
        ATF_REQUIRE(ftruncate(c->fd, FSIZE) == 0);
        ATF_REQUIRE((c->mfd = mmap(NULL, FSIZE, PROT_READ | PROT_WRITE,
            MAP_SHARED, c->fd, 0)) != MAP_FAILED);
        arc4random_buf(c->mfd, FSIZE);

        ATF_REQUIRE(pthread_mutex_init(&c->mtx, NULL) == 0);
        ATF_REQUIRE(pthread_cond_init(&c->cv, NULL) == 0);

        /*
         * Start server and wait for it to finish bind(2) + listen(2).
         */
        ATF_REQUIRE(pthread_mutex_lock(&c->mtx) == 0);
        ATF_REQUIRE(pthread_create(&c->thr, NULL, server_thread, c) == 0);
        if (c->state != READY)
                ATF_REQUIRE(pthread_cond_wait(&c->cv, &c->mtx) == 0);
        ATF_REQUIRE(c->state == READY);
        ATF_REQUIRE(pthread_mutex_unlock(&c->mtx) == 0);

        /*
         * Connect client.
         */
        ATF_REQUIRE(c->ctx = SSL_CTX_new(TLS_client_method()));
        ATF_REQUIRE(X509_STORE_add_cert(SSL_CTX_get_cert_store(c->ctx),
            c->cert));
        ATF_REQUIRE(ssl = c->cln = SSL_new(c->ctx));
        ATF_REQUIRE((c->cs = socket(AF_INET, SOCK_STREAM, 0)) > 0);
        ATF_REQUIRE(connect(c->cs, (struct sockaddr *)&(struct sockaddr_in)
            { .sin_family = AF_INET, .sin_len = sizeof(struct sockaddr_in),
              .sin_addr.s_addr = htonl(INADDR_LOOPBACK), .sin_port = c->port },
            sizeof(struct sockaddr_in)) == 0);
        ATF_REQUIRE(SSL_set_fd(ssl, c->cs) == 1);
        ATF_REQUIRE(snprintf(hostname, sizeof(hostname), "localhost:%u",
            ntohs(c->port)) >= (int)sizeof("localhost:0"));
        ATF_REQUIRE(SSL_set_tlsext_host_name(ssl, hostname) == 1);
        SSL_set_verify(ssl, SSL_VERIFY_PEER, NULL);
        ATF_REQUIRE(SSL_connect(ssl) == 1);
        SSL_set_mode(ssl, SSL_MODE_AUTO_RETRY);
        ATF_REQUIRE(fcntl(c->cs, F_SETFL, O_NONBLOCK) != -1);
}

static void
common_cleanup(struct ctx *c)
{

        ATF_REQUIRE(pthread_mutex_lock(&c->mtx) == 0);
        c->state = EXITING;
        ATF_REQUIRE(pthread_cond_signal(&c->cv) == 0);
        ATF_REQUIRE(pthread_mutex_unlock(&c->mtx) == 0);
        ATF_REQUIRE(pthread_join(c->thr, NULL) == 0);

        ATF_REQUIRE(pthread_mutex_destroy(&c->mtx) == 0);
        ATF_REQUIRE(pthread_cond_destroy(&c->cv) == 0);

        SSL_free(c->cln);
        SSL_CTX_free(c->ctx);
        X509_free(c->cert);
        EVP_PKEY_free(c->pkey);
}

static void *
server_thread(void *arg) {
        struct ctx *c = arg;
        SSL_CTX *srv;
        SSL *ssl;
        struct sockaddr_in sin = {
                .sin_family = AF_INET,
                .sin_len = sizeof(sin),
                .sin_addr.s_addr = htonl(INADDR_LOOPBACK),
        };
        int s;

        ATF_REQUIRE(srv = SSL_CTX_new(TLS_server_method()));
        ATF_REQUIRE(SSL_CTX_set_options(srv, SSL_OP_ENABLE_KTLS) &
            SSL_OP_ENABLE_KTLS);
        SSL_CTX_use_PrivateKey(srv, c->pkey);
        SSL_CTX_use_certificate(srv, c->cert);
        ATF_REQUIRE((s = socket(AF_INET, SOCK_STREAM, 0)) > 0);
        ATF_REQUIRE(setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &(socklen_t){1},
             sizeof(int)) == 0);
        ATF_REQUIRE(bind(s, (struct sockaddr *)&sin, sizeof(sin)) == 0);
        ATF_REQUIRE(getsockname(s, (struct sockaddr *)&sin,
            &(socklen_t){ sizeof(sin) }) == 0);
        ATF_REQUIRE(listen(s, -1) == 0);

        ATF_REQUIRE(pthread_mutex_lock(&c->mtx) == 0);
        c->port = sin.sin_port;
        c->state = READY;
        ATF_REQUIRE(pthread_cond_signal(&c->cv) == 0);
        ATF_REQUIRE(pthread_mutex_unlock(&c->mtx) == 0);

        ATF_REQUIRE((c->ss = accept(s, NULL, NULL)) > 0);
        ssl = c->srv = SSL_new(srv);
        SSL_set_fd(ssl, c->ss);
        SSL_set_mode(ssl, SSL_MODE_AUTO_RETRY);
        ATF_REQUIRE(SSL_accept(ssl) > 0);

        ATF_REQUIRE(pthread_mutex_lock(&c->mtx) == 0);
        while (c->state != EXITING) {
                if (c->state == RUNNING) {
                        ATF_REQUIRE(fcntl(c->ss, F_SETFL,
                            c->nb ? O_NONBLOCK : 0) != -1);
                        c->sbytes = SSL_sendfile(ssl, c->fd, c->offset,
                            c->size, 0);
                        c->state = READY;
                }
                ATF_REQUIRE(c->state == READY);
                ATF_REQUIRE(pthread_cond_signal(&c->cv) == 0);
                ATF_REQUIRE(pthread_cond_wait(&c->cv, &c->mtx) == 0);
        }
        ATF_REQUIRE(pthread_mutex_unlock(&c->mtx) == 0);

        SSL_shutdown(ssl);
        SSL_free(ssl);
        close(c->ss);
        SSL_CTX_free(srv);
        close(s);

        return (NULL);
}

static void
sendme_locked(struct ctx *c, off_t offset, size_t size, bool nb)
{
        ATF_REQUIRE(c->state == READY);
        c->state = RUNNING;
        c->offset = offset;
        c->size = size;
        c->nb = nb;
        ATF_REQUIRE(pthread_cond_signal(&c->cv) == 0);
}

static void
sendme_locked_wait(struct ctx *c, off_t offset, size_t size, bool nb)
{
        sendme_locked(c, offset, size, nb);
        while (c->state != READY)
                ATF_REQUIRE(pthread_cond_wait(&c->cv, &c->mtx) == 0);
}

static void
sendme(struct ctx *c, off_t offset, size_t size, bool nb)
{
        ATF_REQUIRE(pthread_mutex_lock(&c->mtx) == 0);
        sendme_locked(c, offset, size, nb);
        ATF_REQUIRE(pthread_mutex_unlock(&c->mtx) == 0);
}

/*
 * Block until non-blocking socket has at least a byte.
 */
static int
SSL_read_b(SSL *ssl, void *buf, int size)
{
        int rv, fd;

        ATF_REQUIRE((fd = SSL_get_fd(ssl)) > 0);
        while ((rv = SSL_read(ssl, buf, size)) <= 0 &&
            SSL_get_error(ssl, rv) == SSL_ERROR_WANT_READ)
                ATF_REQUIRE(poll(&(struct pollfd){ .fd = fd, .events = POLLIN },
                    1, INFTIM) == 1);

        return (rv);
}

static void
require_sbytes(struct ctx *c, ssize_t expect)
{
        ATF_REQUIRE(pthread_mutex_lock(&c->mtx) == 0);
        ATF_REQUIRE(c->sbytes == expect);
        ATF_REQUIRE(pthread_mutex_unlock(&c->mtx) == 0);
}

ATF_TC_WITHOUT_HEAD(basic);
ATF_TC_BODY(basic, tc)
{
        struct ctx c;
        char buf[64];
        size_t nread;
        int n;

        common_init(&c);

        sendme(&c, 0, 0, false);
        nread = 0;
        while (nread < FSIZE && (n = SSL_read_b(c.cln, buf, sizeof(buf))) > 0) {
                ATF_REQUIRE(memcmp((char *)c.mfd + nread, buf, n) == 0);
                nread += n;
        }
        ATF_REQUIRE(nread == FSIZE);
        require_sbytes(&c, FSIZE);

        common_cleanup(&c);
}

ATF_TC_WITHOUT_HEAD(random);
ATF_TC_BODY(random, tc)
{
        struct ctx c;
#define RSIZE   (256*1024)

        common_init(&c);

        for (u_int i = 0; i < 10; i++) {
                char buf[RSIZE];
                off_t offset;
                size_t size, n, nread, expect;

                offset = arc4random() % FSIZE;
                size = arc4random() % RSIZE;
                sendme(&c, offset, size, false);
                expect = offset + size < FSIZE ? size : FSIZE - offset;
                nread = 0;
                while (nread < expect &&
                    (n = SSL_read_b(c.cln, buf, sizeof(buf))) > 0) {
                        ATF_REQUIRE(memcmp((char *)c.mfd + offset + nread, buf,
                            n) == 0);
                        nread += n;
                }
                ATF_REQUIRE(nread == expect);
                require_sbytes(&c, (ssize_t)expect);
        }

        common_cleanup(&c);
}

/* Truncate the file while sendfile(2) is working on it. */
ATF_TC_WITHOUT_HEAD(truncate);
ATF_TC_BODY(truncate, tc)
{
        struct ctx c;
        char buf[128 * 1024];
        size_t nread;
        int n;
#define TRUNC   (FSIZE - 1024)

        common_init(&c);

        ATF_REQUIRE(setsockopt(c.ss, SOL_SOCKET, SO_SNDBUF, &(int){FSIZE / 16},
            sizeof(int)) == 0);
        ATF_REQUIRE(setsockopt(c.cs, SOL_SOCKET, SO_RCVBUF, &(int){FSIZE / 16},
            sizeof(int)) == 0);

        sendme(&c, 0, 0, false);
        /* Make sure sender is waiting on the socket buffer. */
        while (poll(&(struct pollfd){ .fd = c.ss, .events = POLLOUT }, 1, 1)
            != 0)
                ;
        ATF_REQUIRE(ftruncate(c.fd, TRUNC) == 0);
        nread = 0;
        while (nread < TRUNC && (n = SSL_read_b(c.cln, buf, sizeof(buf))) > 0) {
                ATF_REQUIRE(memcmp((char *)c.mfd + nread, buf, n) == 0);
                nread += n;
        }
        ATF_REQUIRE(nread == TRUNC);
        require_sbytes(&c, TRUNC);

        common_cleanup(&c);
}

/* Grow the file while sendfile(2) is working on it. */
ATF_TC_WITHOUT_HEAD(grow);
ATF_TC_BODY(grow, tc)
{
        struct ctx c;
        char buf[128 * 1024];
        size_t nread;
        void *map;
        int n;
#define GROW    (FSIZE/2)

        common_init(&c);

        ATF_REQUIRE(setsockopt(c.ss, SOL_SOCKET, SO_SNDBUF, &(int){FSIZE / 16},
            sizeof(int)) == 0);
        ATF_REQUIRE(setsockopt(c.cs, SOL_SOCKET, SO_RCVBUF, &(int){FSIZE / 16},
            sizeof(int)) == 0);

        sendme(&c, 0, 0, false);
        /* Make sure sender is waiting on the socket buffer. */
        while (poll(&(struct pollfd){ .fd = c.ss, .events = POLLOUT }, 1, 1)
            != 0)
                ;
        /* Grow the file and create second map. */
        ATF_REQUIRE(ftruncate(c.fd, FSIZE + GROW) == 0);
        ATF_REQUIRE((map = mmap(NULL, GROW, PROT_READ | PROT_WRITE,
            MAP_SHARED, c.fd, FSIZE)) != MAP_FAILED);
        arc4random_buf(map, GROW);

        /* Read out original part. */
        nread = 0;
        while (nread < FSIZE && (n = SSL_read_b(c.cln, buf,
            FSIZE - nread > sizeof(buf) ? sizeof(buf) : FSIZE - nread)) > 0) {
                ATF_REQUIRE(memcmp((char *)c.mfd + nread, buf, n) == 0);
                nread += n;
        }
        ATF_REQUIRE(nread == FSIZE);
        /* Read out grown part. */
        nread = 0;
        while (nread < GROW && (n = SSL_read_b(c.cln, buf, sizeof(buf))) > 0) {
                ATF_REQUIRE(memcmp((char *)map + nread, buf, n) == 0);
                nread += n;
        }
        ATF_REQUIRE(nread == GROW);
        require_sbytes(&c, FSIZE + GROW);

        common_cleanup(&c);
}

ATF_TC_WITHOUT_HEAD(offset_beyond_eof);
ATF_TC_BODY(offset_beyond_eof, tc)
{
        struct ctx c;

        common_init(&c);

        c.sbytes = -1;
        sendme(&c, FSIZE + 1, 0, false);
        ATF_REQUIRE(pthread_mutex_lock(&c.mtx) == 0);
        while (c.state != READY)
                ATF_REQUIRE(pthread_cond_wait(&c.cv, &c.mtx) == 0);
        ATF_REQUIRE(c.sbytes == 0);
        ATF_REQUIRE(pthread_mutex_unlock(&c.mtx) == 0);

        common_cleanup(&c);
}

/*
 * Prove that we can differentiate a short write due to EAGAIN from one due to
 * end of file.
 */
ATF_TC_WITHOUT_HEAD(eagain_vs_eof);
ATF_TC_BODY(eagain_vs_eof, tc)
{
        struct ctx c;
        char buf[16 * 1024];
        ssize_t nread;
        int n;

        common_init(&c);

        /*
         * Exercise short write due to no buffer space on non-blocking
         * socket.  Internall sendfile(2) returns -1 and errno == EAGAIN.
         */
        ATF_REQUIRE(pthread_mutex_lock(&c.mtx) == 0);
        sendme_locked_wait(&c, 0, FSIZE, true);
        ATF_REQUIRE(c.sbytes > 0);
        ATF_REQUIRE(SSL_get_error(c.srv, c.sbytes) == 0);
#if 0   /* see https://github.com/openssl/openssl/issues/29742 */
        ATF_REQUIRE(BIO_should_retry(SSL_get_wbio(c.srv)));
#endif

        /*
         * Exercise second attempt on already full buffer.
         */
        sendme_locked_wait(&c, 0, FSIZE, true);
        ATF_REQUIRE(c.sbytes == -1);
        ATF_REQUIRE(SSL_get_error(c.srv, c.sbytes) == SSL_ERROR_WANT_WRITE);
        ATF_REQUIRE(BIO_should_retry(SSL_get_wbio(c.srv)));

        /* Clear the buffer. */
        nread = 0;
        while (nread < c.sbytes &&
            (n = SSL_read_b(c.cln, buf, sizeof(buf))) > 0) {
                ATF_REQUIRE(memcmp((char *)c.mfd + nread, buf, n) == 0);
                nread += n;
        }

        /*
         * Exercise zero length write: offset == file size.
         *
         * SSL_ERROR_SYSCALL seems a strange error code, as the syscall did not
         * fail, and errno is clear, because a request to send 0 bytes is
         * legitimate one.  This test just documents the existing behavior
         * rather than asserts that this is a correct behavior.
         */
        sendme_locked_wait(&c, FSIZE, 0, true);
        ATF_REQUIRE(c.sbytes == 0);
        ATF_REQUIRE(SSL_get_error(c.srv, c.sbytes) == SSL_ERROR_SYSCALL);
#if 0   /* see https://github.com/openssl/openssl/issues/29742 */
        ATF_REQUIRE(!BIO_should_retry(SSL_get_wbio(c.srv)));
#endif

        /*
         * Exercise short write due to end of file.
         */
        sendme_locked_wait(&c, FSIZE - 100, 0, true);
        ATF_REQUIRE(c.sbytes == 100);
        ATF_REQUIRE(SSL_get_error(c.srv, c.sbytes) == 0);
#if 0   /* see https://github.com/openssl/openssl/issues/29742 */
        ATF_REQUIRE(!BIO_should_retry(SSL_get_wbio(c.srv)));
#endif

        ATF_REQUIRE(pthread_mutex_unlock(&c.mtx) == 0);

        common_cleanup(&c);
}

ATF_TP_ADD_TCS(tp)
{
        ATF_TP_ADD_TC(tp, basic);
        ATF_TP_ADD_TC(tp, random);
        ATF_TP_ADD_TC(tp, truncate);
        ATF_TP_ADD_TC(tp, grow);
        ATF_TP_ADD_TC(tp, offset_beyond_eof);
        ATF_TP_ADD_TC(tp, eagain_vs_eof);

        return atf_no_error();
}