// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only OR HDF5

// This file may be used, modified, and distributed under either GPL-3.0-only
// or the HDF5 license, at the recipient's option.

// Kindly acknowledge modifications from the HDF5 Group

/*
 * enospc_shim — LD_PRELOAD shim that forces ENOSPC after a fixed write
 * budget. Pair with the H5FDpoison_sec2 driver to drive HDF5's
 * out-of-space recovery path against a real disk.
 *
 * Configuration:
 *   ENOSPC_AFTER (env var) — total write bytes allowed before subsequent
 *                            writes/syncs return ENOSPC. Default 10 MiB.
 *                            Read once at constructor time; use the
 *                            runtime API below to change it per test.
 *
 * Runtime API (callable via dlsym after LD_PRELOAD):
 *   void   enospc_shim_reset(size_t new_fail_after);
 *                            Zero the byte counter. If new_fail_after != 0,
 *                            also replace the cap.
 *   size_t enospc_shim_get_total(void);
 *                            Bytes charged against the budget so far.
 *   size_t enospc_shim_get_fail_after(void);
 *                            Current cap.
 *
 * Intercepts: write, pwrite, pwrite64, writev, pwritev, pwritev64,
 *             fallocate, fallocate64, posix_fallocate,
 *             posix_fallocate64, fsync, fdatasync, ftruncate,
 *             ftruncate64, dprintf, vdprintf, and relevant FORTIFY
 *             *_chk write/dprintf variants.
 *
 * Caveats:
 *   - mmap-backed I/O is not intercepted; do not test memory-mapped VFDs
 *     with this shim.
 *   - aio_write, copy_file_range, sendfile, and stdio FILE* output
 *     are not intercepted.
 *   - The byte budget is process-global, not per-fd. Mixing real I/O
 *     (logging, /dev/null, etc.) with the file under test will exhaust
 *     the budget early.
 */

#define _GNU_SOURCE
/*
 * This interposer exports both default and *64 ELF symbols explicitly.
 * Avoid header-level redirects that would rename the default definitions
 * under -D_FILE_OFFSET_BITS=64 and collide with the explicit *64 wrappers.
 */
#ifdef _FILE_OFFSET_BITS
#undef _FILE_OFFSET_BITS
#endif
#ifdef _TIME_BITS
#undef _TIME_BITS
#endif
#include <dlfcn.h>
#include <errno.h>
#include <fcntl.h>
#include <linux/falloc.h>
#include <pthread.h>
#include <stdint.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <unistd.h>

static ssize_t (*real_write)(int, const void *, size_t) = NULL;
static ssize_t (*real_pwrite)(int, const void *, size_t, off_t) = NULL;
static ssize_t (*real_pwrite64)(int, const void *, size_t, off64_t) = NULL;
static ssize_t (*real_writev)(int, const struct iovec *, int) = NULL;
static ssize_t (*real_pwritev)(int, const struct iovec *, int, off_t) = NULL;
static ssize_t (*real_pwritev64)(int, const struct iovec *, int, off64_t) = NULL;
static int (*real_fsync)(int) = NULL;
static int (*real_fdatasync)(int) = NULL;
static int (*real_ftruncate)(int, off_t) = NULL;
static int (*real_ftruncate64)(int, off64_t) = NULL;
static int (*real_fallocate)(int, int, off_t, off_t) = NULL;
static int (*real_fallocate64)(int, int, off64_t, off64_t) = NULL;
static int (*real_posix_fallocate)(int, off_t, off_t) = NULL;
static int (*real_posix_fallocate64)(int, off64_t, off64_t) = NULL;
static int (*real___vasprintf_chk)(char **, int, const char *, va_list) = NULL;

static size_t total_written = 0;
static size_t fail_after = 10ULL * 1024 * 1024;

static pthread_once_t init_once = PTHREAD_ONCE_INIT;
static pthread_mutex_t lock = PTHREAD_MUTEX_INITIALIZER;

#define LOAD_SYMBOL(function_pointer, symbol_name) \
    (*(void **)(&(function_pointer)) = dlsym(RTLD_NEXT, (symbol_name)))

static void init_symbols(void) {
    const char *env = getenv("ENOSPC_AFTER");

    LOAD_SYMBOL(real_write, "write");
    LOAD_SYMBOL(real_pwrite, "pwrite");
    LOAD_SYMBOL(real_pwrite64, "pwrite64");
    LOAD_SYMBOL(real_writev, "writev");
    LOAD_SYMBOL(real_pwritev, "pwritev");
    LOAD_SYMBOL(real_pwritev64, "pwritev64");
    LOAD_SYMBOL(real_fsync, "fsync");
    LOAD_SYMBOL(real_fdatasync, "fdatasync");
    LOAD_SYMBOL(real_ftruncate, "ftruncate");
    LOAD_SYMBOL(real_ftruncate64, "ftruncate64");
    LOAD_SYMBOL(real_fallocate, "fallocate");
    LOAD_SYMBOL(real_fallocate64, "fallocate64");
    LOAD_SYMBOL(real_posix_fallocate, "posix_fallocate");
    LOAD_SYMBOL(real_posix_fallocate64, "posix_fallocate64");
    LOAD_SYMBOL(real___vasprintf_chk, "__vasprintf_chk");

    if (env != NULL)
        fail_after = strtoull(env, NULL, 10);
}

__attribute__((constructor))
static void init(void) {
    pthread_once(&init_once, init_symbols);
}

/*
 * Reserve up to `requested` bytes against the global budget.
 *
 *   - returns -1 with errno=ENOSPC if the budget is already exhausted.
 *   - on success, *allowed is the number of bytes the caller may now
 *     pass to the real write syscall. It may be smaller than requested,
 *     producing a deliberate short write that mirrors how a kernel
 *     reports a partial fill before the next call hits ENOSPC.
 */
static int reserve_bytes(size_t requested, size_t *allowed) {
    pthread_mutex_lock(&lock);

    if (requested == 0) {
        *allowed = 0;
        pthread_mutex_unlock(&lock);
        return 0;
    }

    if (total_written >= fail_after) {
        pthread_mutex_unlock(&lock);
        errno = ENOSPC;
        return -1;
    }

    *allowed = requested;
    if (requested > fail_after - total_written)
        *allowed = fail_after - total_written;

    total_written += *allowed;
    pthread_mutex_unlock(&lock);
    return 0;
}

/*
 * If the real syscall wrote fewer bytes than reserved (kernel short
 * write or outright failure), refund the difference so the next caller
 * gets a faithful budget. Without this, a transient short write would
 * permanently shrink the effective limit below ENOSPC_AFTER.
 */
static void release_unused_reservation(size_t reserved, ssize_t written) {
    size_t actual = written > 0 ? (size_t)written : 0;
    size_t unused;

    if (actual >= reserved)
        return;

    unused = reserved - actual;

    pthread_mutex_lock(&lock);
    if (unused > total_written)
        total_written = 0;
    else
        total_written -= unused;
    pthread_mutex_unlock(&lock);
}

static size_t iov_total(const struct iovec *iov, int iovcnt) {
    size_t total = 0;

    for (int i = 0; i < iovcnt; i++) {
        if (SIZE_MAX - total < iov[i].iov_len)
            return SIZE_MAX;
        total += iov[i].iov_len;
    }

    return total;
}

/*
 * Build a copy of an iovec truncated to at most `allowed` bytes total.
 * Used by writev/pwritev/pwritev64 to deliver a controlled short write
 * that consumes exactly the remaining budget. Caller frees the copy.
 */
static int trim_iov(
    const struct iovec *iov,
    int iovcnt,
    size_t allowed,
    struct iovec **trimmed_iov,
    int *trimmed_iovcnt
) {
    struct iovec *copy;
    int used = 0;

    if (iovcnt <= 0 || allowed == 0) {
        *trimmed_iov = NULL;
        *trimmed_iovcnt = 0;
        return 0;
    }

    copy = malloc((size_t)iovcnt * sizeof(*copy));
    if (copy == NULL)
        return -1;

    for (int i = 0; i < iovcnt && allowed > 0; i++) {
        size_t len = iov[i].iov_len;

        if (len > allowed)
            len = allowed;

        copy[used].iov_base = iov[i].iov_base;
        copy[used].iov_len = len;
        used++;
        allowed -= len;
    }

    *trimmed_iov = copy;
    *trimmed_iovcnt = used;
    return 0;
}

/*
 * fsync / fdatasync entry-check: once the byte budget is gone, sync
 * fails with ENOSPC. Real filesystems can defer ENOSPC to fsync if
 * delayed allocation is in play, so this models that behavior.
 */
static int fail_if_exhausted(void) {
    pthread_mutex_lock(&lock);
    if (total_written < fail_after) {
        pthread_mutex_unlock(&lock);
        return 0;
    }
    pthread_mutex_unlock(&lock);

    errno = ENOSPC;
    return -1;
}

static size_t current_fail_after(void) {
    size_t cap;

    pthread_mutex_lock(&lock);
    cap = fail_after;
    pthread_mutex_unlock(&lock);
    return cap;
}

static int positive_range_exceeds_budget(uintmax_t offset, uintmax_t len) {
    uintmax_t end;

    if (UINTMAX_MAX - offset < len)
        return 1;

    end = offset + len;
    return end > (uintmax_t)current_fail_after();
}

static int signed_range_exceeds_budget(off64_t offset, off64_t len) {
    if (offset < 0 || len < 0)
        return 0;

    return positive_range_exceeds_budget((uintmax_t)offset, (uintmax_t)len);
}

static int file_size_plus_len_exceeds_budget(int fd, off64_t len) {
    struct stat st;

    if (len <= 0)
        return 0;

    if (fstat(fd, &st) < 0)
        return 0;

    if (st.st_size < 0)
        return 0;

    return positive_range_exceeds_budget((uintmax_t)st.st_size, (uintmax_t)len);
}

static int fallocate_exceeds_budget(
    int fd,
    int mode,
    off64_t offset,
    off64_t len
) {
    if (offset < 0 || len <= 0)
        return 0;

    if ((mode & (FALLOC_FL_PUNCH_HOLE | FALLOC_FL_COLLAPSE_RANGE)) != 0)
        return 0;

    if ((mode & FALLOC_FL_INSERT_RANGE) != 0)
        return file_size_plus_len_exceeds_budget(fd, len);

    return signed_range_exceeds_budget(offset, len);
}

static int truncate_length_exceeds_budget(off64_t length) {
    if (length < 0)
        return 0;

    return positive_range_exceeds_budget((uintmax_t)length, 0);
}

/*
 * write/pwrite/pwrite64 follow an identical pattern:
 *   1. resolve the real symbol on first use,
 *   2. reserve at most `count` bytes — return -1/ENOSPC if exhausted,
 *   3. issue the real syscall on the trimmed length,
 *   4. refund any portion the kernel ultimately did not write.
 *
 * The writev/pwritev/pwritev64 variants additionally call trim_iov() to
 * truncate the iovec array to the reserved length.
 */
ssize_t write(int fd, const void *buf, size_t count) {
    size_t allowed = 0;
    ssize_t ret;

    pthread_once(&init_once, init_symbols);
    if (real_write == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (reserve_bytes(count, &allowed) < 0)
        return -1;

    ret = real_write(fd, buf, allowed);
    release_unused_reservation(allowed, ret);
    return ret;
}

ssize_t pwrite(int fd, const void *buf, size_t count, off_t offset) {
    size_t allowed = 0;
    ssize_t ret;

    pthread_once(&init_once, init_symbols);
    if (real_pwrite == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (reserve_bytes(count, &allowed) < 0)
        return -1;

    ret = real_pwrite(fd, buf, allowed, offset);
    release_unused_reservation(allowed, ret);
    return ret;
}

ssize_t pwrite64(int fd, const void *buf, size_t count, off64_t offset) {
    size_t allowed = 0;
    ssize_t ret;

    pthread_once(&init_once, init_symbols);
    if (real_pwrite64 == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (reserve_bytes(count, &allowed) < 0)
        return -1;

    ret = real_pwrite64(fd, buf, allowed, offset);
    release_unused_reservation(allowed, ret);
    return ret;
}

ssize_t writev(int fd, const struct iovec *iov, int iovcnt) {
    struct iovec *trimmed_iov = NULL;
    int trimmed_iovcnt = 0;
    size_t allowed = 0;
    size_t requested;
    ssize_t ret;

    pthread_once(&init_once, init_symbols);
    if (real_writev == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (iovcnt <= 0 || iov == NULL)
        return real_writev(fd, iov, iovcnt);

    requested = iov_total(iov, iovcnt);
    if (reserve_bytes(requested, &allowed) < 0)
        return -1;

    if (trim_iov(iov, iovcnt, allowed, &trimmed_iov, &trimmed_iovcnt) < 0) {
        release_unused_reservation(allowed, -1);
        return -1;
    }

    ret = real_writev(fd, trimmed_iov, trimmed_iovcnt);
    free(trimmed_iov);
    release_unused_reservation(allowed, ret);
    return ret;
}

ssize_t pwritev(int fd, const struct iovec *iov, int iovcnt, off_t offset) {
    struct iovec *trimmed_iov = NULL;
    int trimmed_iovcnt = 0;
    size_t allowed = 0;
    size_t requested;
    ssize_t ret;

    pthread_once(&init_once, init_symbols);
    if (real_pwritev == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (iovcnt <= 0 || iov == NULL)
        return real_pwritev(fd, iov, iovcnt, offset);

    requested = iov_total(iov, iovcnt);
    if (reserve_bytes(requested, &allowed) < 0)
        return -1;

    if (trim_iov(iov, iovcnt, allowed, &trimmed_iov, &trimmed_iovcnt) < 0) {
        release_unused_reservation(allowed, -1);
        return -1;
    }

    ret = real_pwritev(fd, trimmed_iov, trimmed_iovcnt, offset);
    free(trimmed_iov);
    release_unused_reservation(allowed, ret);
    return ret;
}

ssize_t pwritev64(int fd, const struct iovec *iov, int iovcnt, off64_t offset) {
    struct iovec *trimmed_iov = NULL;
    int trimmed_iovcnt = 0;
    size_t allowed = 0;
    size_t requested;
    ssize_t ret;

    pthread_once(&init_once, init_symbols);
    if (real_pwritev64 == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (iovcnt <= 0 || iov == NULL)
        return real_pwritev64(fd, iov, iovcnt, offset);

    requested = iov_total(iov, iovcnt);
    if (reserve_bytes(requested, &allowed) < 0)
        return -1;

    if (trim_iov(iov, iovcnt, allowed, &trimmed_iov, &trimmed_iovcnt) < 0) {
        release_unused_reservation(allowed, -1);
        return -1;
    }

    ret = real_pwritev64(fd, trimmed_iov, trimmed_iovcnt, offset);
    free(trimmed_iov);
    release_unused_reservation(allowed, ret);
    return ret;
}

/*
 * Defensive FORTIFY entry points. glibc does not currently route write()
 * through these symbols, but other libc/toolchain combinations may. Keep
 * the fortify size check on the caller's original request before applying
 * the shim's usual trimming behavior.
 */
ssize_t __write_chk(int fd, const void *buf, size_t count, size_t buflen) {
    if (count > buflen)
        abort();

    return write(fd, buf, count);
}

ssize_t __pwrite_chk(
    int fd,
    const void *buf,
    size_t count,
    off_t offset,
    size_t buflen
) {
    if (count > buflen)
        abort();

    return pwrite(fd, buf, count, offset);
}

ssize_t __pwrite64_chk(
    int fd,
    const void *buf,
    size_t count,
    off64_t offset,
    size_t buflen
) {
    if (count > buflen)
        abort();

    return pwrite64(fd, buf, count, offset);
}

static int write_all_budgeted(int fd, const char *buf, size_t len) {
    size_t done = 0;

    while (done < len) {
        ssize_t ret = write(fd, buf + done, len - done);

        if (ret < 0)
            return -1;

        if (ret == 0) {
            errno = EIO;
            return -1;
        }

        done += (size_t)ret;
    }

    return (int)done;
}

static int budgeted_vdprintf_chk(
    int fd,
    int flag,
    const char *format,
    va_list ap
) {
    char *formatted = NULL;
    va_list ap_copy;
    int len;
    int ret;

    pthread_once(&init_once, init_symbols);

    va_copy(ap_copy, ap);
    if (real___vasprintf_chk != NULL)
        len = real___vasprintf_chk(&formatted, flag, format, ap_copy);
    else
        len = vasprintf(&formatted, format, ap_copy);
    va_end(ap_copy);

    if (len < 0)
        return -1;

    ret = write_all_budgeted(fd, formatted, (size_t)len);
    free(formatted);
    return ret;
}

int vdprintf(int fd, const char *format, va_list ap) {
    return budgeted_vdprintf_chk(fd, 0, format, ap);
}

int dprintf(int fd, const char *format, ...) {
    va_list ap;
    int ret;

    va_start(ap, format);
    ret = budgeted_vdprintf_chk(fd, 0, format, ap);
    va_end(ap);
    return ret;
}

int __vdprintf_chk(int fd, int flag, const char *format, va_list ap) {
    return budgeted_vdprintf_chk(fd, flag, format, ap);
}

int __dprintf_chk(int fd, int flag, const char *format, ...) {
    va_list ap;
    int ret;

    va_start(ap, format);
    ret = budgeted_vdprintf_chk(fd, flag, format, ap);
    va_end(ap);
    return ret;
}

/*
 * fallocate / posix_fallocate do not have a short-success convention we can
 * use. Refuse allocation ranges that would extend past ENOSPC_AFTER and
 * otherwise delegate to the real libc entry point.
 */
int fallocate(int fd, int mode, off_t offset, off_t len) {
    pthread_once(&init_once, init_symbols);
    if (real_fallocate == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (fallocate_exceeds_budget(fd, mode, (off64_t)offset, (off64_t)len)) {
        errno = ENOSPC;
        return -1;
    }

    return real_fallocate(fd, mode, offset, len);
}

int fallocate64(int fd, int mode, off64_t offset, off64_t len) {
    pthread_once(&init_once, init_symbols);
    if (real_fallocate64 == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (fallocate_exceeds_budget(fd, mode, offset, len)) {
        errno = ENOSPC;
        return -1;
    }

    return real_fallocate64(fd, mode, offset, len);
}

int posix_fallocate(int fd, off_t offset, off_t len) {
    pthread_once(&init_once, init_symbols);
    if (real_posix_fallocate == NULL)
        return ENOSYS;

    if (fallocate_exceeds_budget(fd, 0, (off64_t)offset, (off64_t)len))
        return ENOSPC;

    return real_posix_fallocate(fd, offset, len);
}

int posix_fallocate64(int fd, off64_t offset, off64_t len) {
    pthread_once(&init_once, init_symbols);
    if (real_posix_fallocate64 == NULL)
        return ENOSYS;

    if (fallocate_exceeds_budget(fd, 0, offset, len))
        return ENOSPC;

    return real_posix_fallocate64(fd, offset, len);
}

int fsync(int fd) {
    pthread_once(&init_once, init_symbols);
    if (real_fsync == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (fail_if_exhausted() < 0)
        return -1;

    return real_fsync(fd);
}

int fdatasync(int fd) {
    pthread_once(&init_once, init_symbols);
    if (real_fdatasync == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (fail_if_exhausted() < 0)
        return -1;

    return real_fdatasync(fd);
}

/*
 * ftruncate / ftruncate64: refuse any extension past the byte budget.
 * Truncations *down* are still permitted — they don't consume space and
 * may be used by HDF5 cleanup paths. The check is on the absolute target
 * length, not on the delta, so it works for a freshly-opened fd too.
 */
int ftruncate(int fd, off_t length) {
    pthread_once(&init_once, init_symbols);
    if (real_ftruncate == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (truncate_length_exceeds_budget((off64_t)length)) {
        errno = ENOSPC;
        return -1;
    }

    return real_ftruncate(fd, length);
}

int ftruncate64(int fd, off64_t length) {
    pthread_once(&init_once, init_symbols);
    if (real_ftruncate64 == NULL) {
        errno = ENOSYS;
        return -1;
    }

    if (truncate_length_exceeds_budget(length)) {
        errno = ENOSPC;
        return -1;
    }

    return real_ftruncate64(fd, length);
}

/*
 * Test/control API. Default ELF visibility makes these symbols dlsym-able
 * from a test harness — declare them as extern in the test code.
 */
void enospc_shim_reset(size_t new_fail_after) {
    pthread_mutex_lock(&lock);
    total_written = 0;
    if (new_fail_after != 0)
        fail_after = new_fail_after;
    pthread_mutex_unlock(&lock);
}

size_t enospc_shim_get_total(void) {
    size_t total;

    pthread_mutex_lock(&lock);
    total = total_written;
    pthread_mutex_unlock(&lock);
    return total;
}

size_t enospc_shim_get_fail_after(void) {
    size_t cap;

    pthread_mutex_lock(&lock);
    cap = fail_after;
    pthread_mutex_unlock(&lock);
    return cap;
}