/*
 * Copyright (c) 2025 Mark Johnston <markj@FreeBSD.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <sys/mman.h>
#include <sys/stat.h>

#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <atf-c.h>
#include <sha256.h>

/*
 * Create a file with random data and size between 1B and 32MB.  Return a file
 * descriptor for the file.
 */
static int
genfile(void)
{
        char buf[256], file[NAME_MAX];
        size_t sz;
        int fd;

        sz = (random() % (32 * 1024 * 1024ul)) + 1;

        snprintf(file, sizeof(file), "testfile.XXXXXX");
        fd = mkstemp(file);
        ATF_REQUIRE(fd != -1);

        while (sz > 0) {
                ssize_t n;
                int error;

                error = getentropy(buf, sizeof(buf));
                ATF_REQUIRE(error == 0);
                n = write(fd, buf, sizeof(buf) < sz ? sizeof(buf) : sz);
                ATF_REQUIRE(n > 0);

                sz -= n;
        }

        ATF_REQUIRE(lseek(fd, 0, SEEK_SET) == 0);
        return (fd);
}

/*
 * Return true if the file data in the two file descriptors is the same,
 * false otherwise.
 */
static bool
cmpfile(int fd1, int fd2)
{
        struct stat st1, st2;
        void *addr1, *addr2;
        size_t sz;
        int res;

        ATF_REQUIRE(fstat(fd1, &st1) == 0);
        ATF_REQUIRE(fstat(fd2, &st2) == 0);
        if (st1.st_size != st2.st_size)
                return (false);

        sz = st1.st_size;
        addr1 = mmap(NULL, sz, PROT_READ, MAP_PRIVATE, fd1, 0);
        ATF_REQUIRE(addr1 != MAP_FAILED);
        addr2 = mmap(NULL, sz, PROT_READ, MAP_PRIVATE, fd2, 0);
        ATF_REQUIRE(addr2 != MAP_FAILED);

        res = memcmp(addr1, addr2, sz);

        ATF_REQUIRE(munmap(addr1, sz) == 0);
        ATF_REQUIRE(munmap(addr2, sz) == 0);

        return (res == 0);
}

/*
 * Exercise a few error paths in the copy_file_range() syscall.
 */
ATF_TC_WITHOUT_HEAD(copy_file_range_invalid);
ATF_TC_BODY(copy_file_range_invalid, tc)
{
        off_t off1, off2;
        int fd1, fd2;

        fd1 = genfile();
        fd2 = genfile();

        /* Can't copy a file to itself without explicit offsets. */
        ATF_REQUIRE_ERRNO(EINVAL,
            copy_file_range(fd1, NULL, fd1, NULL, SSIZE_MAX, 0) == -1);

        /* When copying a file to itself, ranges cannot overlap. */
        off1 = off2 = 0;
        ATF_REQUIRE_ERRNO(EINVAL,
            copy_file_range(fd1, &off1, fd1, &off2, 1, 0) == -1);

        /* Negative offsets are not allowed. */
        off1 = -1;
        off2 = 0;
        ATF_REQUIRE_ERRNO(EINVAL,
            copy_file_range(fd1, &off1, fd2, &off2, 42, 0) == -1);
        ATF_REQUIRE_ERRNO(EINVAL,
            copy_file_range(fd2, &off2, fd1, &off1, 42, 0) == -1);
}

/*
 * Make sure that copy_file_range() updates the file offsets passed to it.
 */
ATF_TC_WITHOUT_HEAD(copy_file_range_offset);
ATF_TC_BODY(copy_file_range_offset, tc)
{
        struct stat sb;
        off_t off1, off2;
        ssize_t n;
        int fd1, fd2;

        off1 = off2 = 0;

        fd1 = genfile();
        fd2 = open("copy", O_RDWR | O_CREAT, 0644);
        ATF_REQUIRE(fd2 != -1);

        ATF_REQUIRE(fstat(fd1, &sb) == 0);

        ATF_REQUIRE(lseek(fd1, 0, SEEK_CUR) == 0);
        ATF_REQUIRE(lseek(fd2, 0, SEEK_CUR) == 0);

        do {
                off_t ooff1, ooff2;

                ooff1 = off1;
                ooff2 = off2;
                n = copy_file_range(fd1, &off1, fd2, &off2, sb.st_size, 0);
                ATF_REQUIRE(n >= 0);
                ATF_REQUIRE_EQ(off1, ooff1 + n);
                ATF_REQUIRE_EQ(off2, ooff2 + n);
        } while (n != 0);

        /* Offsets should have been adjusted by copy_file_range(). */
        ATF_REQUIRE_EQ(off1, sb.st_size);
        ATF_REQUIRE_EQ(off2, sb.st_size);
        /* Seek offsets should have been left alone. */
        ATF_REQUIRE(lseek(fd1, 0, SEEK_CUR) == 0);
        ATF_REQUIRE(lseek(fd2, 0, SEEK_CUR) == 0);
        /* Make sure the file contents are the same. */
        ATF_REQUIRE_MSG(cmpfile(fd1, fd2), "file contents differ");

        ATF_REQUIRE(close(fd1) == 0);
        ATF_REQUIRE(close(fd2) == 0);
}

/*
 * Make sure that copying to a larger file doesn't cause it to be truncated.
 */
ATF_TC_WITHOUT_HEAD(copy_file_range_truncate);
ATF_TC_BODY(copy_file_range_truncate, tc)
{
        struct stat sb, sb1, sb2;
        char digest1[65], digest2[65];
        off_t off;
        ssize_t n;
        int fd1, fd2;

        fd1 = genfile();
        fd2 = genfile();

        ATF_REQUIRE(fstat(fd1, &sb1) == 0);
        ATF_REQUIRE(fstat(fd2, &sb2) == 0);

        /* fd1 refers to the smaller file. */
        if (sb1.st_size > sb2.st_size) {
                int tmp;

                tmp = fd1;
                fd1 = fd2;
                fd2 = tmp;
                ATF_REQUIRE(fstat(fd1, &sb1) == 0);
                ATF_REQUIRE(fstat(fd2, &sb2) == 0);
        }

        /*
         * Compute a hash of the bytes in the larger file which lie beyond the
         * length of the smaller file.
         */
        SHA256_FdChunk(fd2, digest1, sb1.st_size, sb2.st_size - sb1.st_size);
        ATF_REQUIRE(lseek(fd2, 0, SEEK_SET) == 0);

        do {
                n = copy_file_range(fd1, NULL, fd2, NULL, SSIZE_MAX, 0);
                ATF_REQUIRE(n >= 0);
        } while (n != 0);

        /* Validate file offsets after the copy. */
        off = lseek(fd1, 0, SEEK_CUR);
        ATF_REQUIRE(off == sb1.st_size);
        off = lseek(fd2, 0, SEEK_CUR);
        ATF_REQUIRE(off == sb1.st_size);

        /* The larger file's size should remain the same. */
        ATF_REQUIRE(fstat(fd2, &sb) == 0);
        ATF_REQUIRE(sb.st_size == sb2.st_size);

        /* The bytes beyond the end of the copy should be unchanged. */
        SHA256_FdChunk(fd2, digest2, sb1.st_size, sb2.st_size - sb1.st_size);
        ATF_REQUIRE_MSG(strcmp(digest1, digest2) == 0,
            "trailing file contents differ after copy_file_range()");

        /*
         * Verify that the copy actually replicated bytes from the smaller file.
         */
        ATF_REQUIRE(ftruncate(fd2, sb1.st_size) == 0);
        ATF_REQUIRE(cmpfile(fd1, fd2));
}

ATF_TP_ADD_TCS(tp)
{
        ATF_TP_ADD_TC(tp, copy_file_range_invalid);
        ATF_TP_ADD_TC(tp, copy_file_range_offset);
        ATF_TP_ADD_TC(tp, copy_file_range_truncate);

        return (atf_no_error());
}