/* SPDX-License-Identifier: LGPL-2.1-or-later */

#include <linux/fs.h>
#include <sys/file.h>
#include <sys/ioctl.h>

#include "sd-device.h"

#include "alloc-util.h"
#include "blkid-util.h"
#include "blockdev-util.h"
#include "device-private.h"
#include "device-util.h"
#include "errno-util.h"
#include "fd-util.h"
#include "log.h"
#include "reread-partition-table.h"
#include "set.h"
#include "string-util.h"

static int trigger_partitions(sd_device *dev, bool blkrrpart_success) {
        int ret = 0, r;

        assert(dev);

        /* search for partitions */
        _cleanup_(sd_device_enumerator_unrefp) sd_device_enumerator *e = NULL;
        r = partition_enumerator_new(dev, &e);
        if (r < 0)
                return log_device_debug_errno(dev, r, "Failed to initialize partition enumerator: %m");

        /* We have partitions and re-read the table, the kernel already sent out a "change"
         * event for the disk, and "remove/add" for all partitions. */
        if (blkrrpart_success && sd_device_enumerator_get_device_first(e))
                return 0;

        /* We have partitions but re-reading the partition table did not work, synthesize
         * "change" for the disk and all partitions. */
        r = sd_device_trigger(dev, SD_DEVICE_CHANGE);
        if (r < 0)
                RET_GATHER(ret, log_device_debug_errno(dev, r, "Failed to trigger 'change' uevent, proceeding: %m"));

        FOREACH_DEVICE(e, d) {
                r = sd_device_trigger(d, SD_DEVICE_CHANGE);
                if (r < 0)
                        RET_GATHER(ret, log_device_debug_errno(d, r, "Failed to trigger 'change' uevent, proceeding: %m"));
        }

        return ret;
}

static int fallback_ioctl(sd_device *d, int fd, RereadPartitionTableFlags flags) {
        int r;

        assert(d);
        assert(fd >= 0);

        r = RET_NERRNO(ioctl(fd, BLKRRPART, 0));
        if (r < 0)
                log_device_debug_errno(d, r, "Failed to reread partition table via BLKRRPART: %m");
        else
                log_device_debug(d, "Successfully reread partition table via BLKRRPART.");

        if (FLAGS_SET(flags, REREADPT_FORCE_UEVENT))
                RET_GATHER(r, trigger_partitions(d, r >= 0));

        return r;
}

#if HAVE_BLKID
static int process_partition(
                sd_device *d,
                int fd,
                blkid_partition pp,
                sd_device_enumerator *e,
                Set **partnos,
                RereadPartitionTableFlags flags,
                bool *changed) {

        int r;

        assert(d);
        assert(fd >= 0);
        assert(pp);
        assert(e);
        assert(partnos);
        assert(changed);

        const char *node;
        r = sd_device_get_devname(d, &node);
        if (r < 0)
                return log_device_debug_errno(d, r, "Failed to acquire device node path: %m");

        errno = 0;
        int nr = sym_blkid_partition_get_partno(pp);
        if (nr < 0)
                return log_debug_errno(errno_or_else(EIO), "Failed to read partition number of partition: %m");

        log_device_debug(d, "Processing partition %i...", nr);

        errno = 0;
        blkid_loff_t start = sym_blkid_partition_get_start(pp);
        if (start < 0)
                return log_debug_errno(errno_or_else(EIO), "Failed to read partition start offset of partition %i: %m", nr);
        assert((uint64_t) start < UINT64_MAX / 512U);

        errno = 0;
        blkid_loff_t size = sym_blkid_partition_get_size(pp);
        if (size < 0)
                return log_debug_errno(errno_or_else(EIO), "Failed to read partition size of partition %i: %m", nr);
        assert((uint64_t) size < UINT64_MAX / 512U);

        if (set_ensure_put(partnos, /* hash_ops= */ NULL, UINT_TO_PTR(nr)) < 0)
                return log_oom_debug();

        _cleanup_free_ char *subnode = NULL;
        r = partition_node_of(node, nr, &subnode);
        if (r < 0)
                return log_device_debug_errno(d, r, "Failed to determine partition node %i for '%s': %m", nr, node);

        _cleanup_(sd_device_unrefp) sd_device *partition = NULL;
        r = sd_device_new_from_devname(&partition, subnode);
        if (r < 0) {
                if (r != -ENODEV)
                        return log_device_debug_errno(d, r, "Failed to acquire device '%s': %m", subnode);
        } else {
                uint64_t start_kernel;
                r = device_get_sysattr_u64(partition, "start", &start_kernel);
                if (r < 0)
                        return log_device_debug_errno(partition, r, "Failed to get start of kernel partition device '%s': %m", subnode);

                uint64_t size_kernel;
                r = device_get_sysattr_u64(partition, "size", &size_kernel);
                if (r < 0)
                        return log_device_debug_errno(partition, r, "Failed to get size of kernel partition device '%s': %m", subnode);

                if (start_kernel == (uint64_t) start && size_kernel == (uint64_t) size) {
                        log_device_debug(partition, "Kernel partition device '%s' already matches partition table, not modifying.", subnode);

                        if (FLAGS_SET(flags, REREADPT_FORCE_UEVENT)) {
                                if (!*changed) {
                                        /* Make sure to synthesize a change event on the main device, before we issue the first one on a partition device */
                                        r = sd_device_trigger(d, SD_DEVICE_CHANGE);
                                        if (r < 0)
                                                return log_device_debug_errno(d, r, "Failed to issue 'change' uevent on device '%s': %m", node);

                                        log_device_debug(d, "Successfully issued 'change' uevent on device '%s'.", node);
                                        *changed = true;
                                }

                                r = sd_device_trigger(partition, SD_DEVICE_CHANGE);
                                if (r < 0)
                                        return log_device_debug_errno(partition, r, "Failed to issue 'change' uevent on partition '%s': %m", subnode);

                                log_device_debug(partition, "Successfully issued 'change' uevent on partition '%s'.", subnode);
                        }

                        return 0;
                }

                if (start_kernel == (uint64_t) start) {
                        /* If the start offsize doesn't change we can just resize the partition */
                        log_device_debug(partition, "Resizing partition %i...", nr);

                        r = block_device_resize_partition(fd, nr, (uint64_t) start * 512U, (uint64_t) size * 512U);
                        if (r < 0)
                                return log_device_debug_errno(partition, r, "Failed to resize kernel partition device '%s' to partition table values: %m", subnode);

                        log_device_debug(partition, "Successfully resized kernel partition device '%s' to match partition table.", subnode);
                        *changed = true;
                        return 1;
                }

                /* If the start offset changed we need to remove and recreate the partition */
                log_device_debug(partition, "Removing and recreating partition %i...", nr);

                /* NB: when logging below we use the parent device now, after all the partition device ceased
                 * existing by now, most likely. Let's explicitly get rid of the obsolete device object now,
                 * just to make a point. */
                partition = sd_device_unref(partition);

                r = block_device_remove_partition(fd, nr);
                if (r < 0)
                        return log_device_debug_errno(d, r, "Failed to remove kernel partition device '%s' in order to recreate it: %m", subnode);

                /* And now add it the partition anew */
                log_device_debug(d, "Successfully removed kernel partition device '%s' in order to recreate it.", subnode);
        }

        log_device_debug(d, "Adding partition %i...", nr);

        r = block_device_add_partition(fd, nr, (uint64_t) start * 512U, (uint64_t) size * 512U);
        if (r < 0)
                return log_device_debug_errno(d, r, "Failed to add kernel partition device %i to partition table values: %m", nr);

        log_device_debug(d, "Successfully added kernel partition device %i to match partition table.", nr);
        *changed = true;
        return 1;
}

static int remove_partitions(sd_device *d, int fd, sd_device_enumerator *e, Set *partnos, bool *changed) {
        int r;

        assert(d);
        assert(fd >= 0);
        assert(e);
        assert(changed);

        /* Removes all partitions of the specified device that we didn't find in the partition table (as
         * listed in the specified Set object) */

        int ret = 0;
        FOREACH_DEVICE(e, partition) {
                const char *devname;

                r = sd_device_get_devname(partition, &devname);
                if (r < 0)
                        return log_device_debug_errno(partition, r, "Failed to get name of partition: %m");

                unsigned nr;
                r = device_get_property_uint(partition, "PARTN", &nr);
                if (r < 0)
                        return log_device_debug_errno(partition, r, "Failed to read partition number property: %m");
                if (set_contains(partnos, UINT_TO_PTR(nr))) {
                        log_device_debug(partition, "Found kernel partition device %u in partition table, leaving around.", nr);
                        continue;
                }

                log_device_debug(partition, "Kernel knows partition %u which we didn't find, removing.", nr);

                r = block_device_remove_partition(fd, (int) nr);
                if (r < 0) /* NB: when logging we use the parent device below, after all the partition device ceased existing by now, most likely */
                        RET_GATHER(ret, log_device_debug_errno(d, r, "Failed to remove kernel partition device '%s' that vanished from partition table: %m", devname));
                else  {
                        log_device_debug(d, "Removed partition %u from kernel.", nr);
                        *changed = true;
                }
        }

        return ret;
}
#endif

static int reread_partition_table_full(sd_device *dev, int fd, RereadPartitionTableFlags flags) {
        int r;

        assert(dev);
        assert(fd >= 0);

        const char *p;
        r = sd_device_get_devname(dev, &p);
        if (r < 0)
                return log_device_debug_errno(dev, r, "Failed to get block device name: %m");

        _cleanup_close_ int lock_fd = -EBADF;
        if (FLAGS_SET(flags, REREADPT_BSD_LOCK)) {
                lock_fd = fd_reopen(fd, O_RDONLY|O_CLOEXEC|O_NOCTTY);
                if (lock_fd < 0)
                        return log_device_debug_errno(dev, lock_fd, "Failed to open lock fd for block device '%s': %m", p);

                if (flock(lock_fd, LOCK_EX|LOCK_NB) < 0) {
                        r = log_device_debug_errno(dev, errno, "Failed to take BSD lock on block device '%s': %m", p);

                        if (r == -EAGAIN && FLAGS_SET(flags, REREADPT_FORCE_UEVENT)) {
                                log_device_debug(dev, "Giving up rereading partition table of '%s'. Triggering change events for the device and its partitions.", p);
                                (void) trigger_partitions(dev, /* blkrrpart_success= */ false);
                        }

                        return r;
                }
        }

        r = blockdev_partscan_enabled(dev);
        if (r < 0)
                return log_device_debug_errno(dev, r, "Failed to test if block device '%s' knows partition scanning: %m", p);
        if (r == 0) {
                /* No partition scanning? Generate a uevent at least, if that's requested */
                if (FLAGS_SET(flags, REREADPT_FORCE_UEVENT)) {
                        r = sd_device_trigger(dev, SD_DEVICE_CHANGE);
                        if (r < 0)
                                return log_device_debug_errno(dev, r, "Failed to trigger 'change' uevent, proceeding: %m");

                        return 0;
                }

                return log_device_debug_errno(dev, SYNTHETIC_ERRNO(ENOTTY), "Block device '%s' does not support partition scanning.", p);
        }

#if HAVE_BLKID
        r = dlopen_libblkid(LOG_DEBUG);
        if (ERRNO_IS_NEG_NOT_SUPPORTED(r)) {
                log_device_debug(dev, "We don't have libblkid, falling back to BLKRRPART on '%s'.", p);
                return fallback_ioctl(dev, fd, flags);
        }
        if (r < 0)
                return log_device_debug_errno(dev, r, "Failed to load libblkid: %m");

        _cleanup_(blkid_free_probep) blkid_probe b = sym_blkid_new_probe();
        if (!b)
                return log_oom_debug();

        errno = 0;
        r = sym_blkid_probe_set_device(b, fd, /* off= */ 0, /* size= */ 0);
        if (r != 0)
                return log_device_debug_errno(dev, errno_or_else(ENOMEM), "Failed to open block device '%s': %m", p);

        (void) sym_blkid_probe_enable_partitions(b, 1);
        (void) sym_blkid_probe_set_partitions_flags(b, BLKID_PARTS_ENTRY_DETAILS);

        errno = 0;
        r = sym_blkid_do_safeprobe(b);
        if (r == _BLKID_SAFEPROBE_ERROR)
                return log_device_debug_errno(dev, errno_or_else(EIO), "Unable to probe for partition table of '%s': %m", p);
        if (IN_SET(r, _BLKID_SAFEPROBE_AMBIGUOUS, _BLKID_SAFEPROBE_NOT_FOUND)) {
                log_device_debug(dev, "Didn't find partition table on block device '%s', falling back to BLKRRPART.", p);
                return fallback_ioctl(dev, fd, flags);
        }

        assert(r == _BLKID_SAFEPROBE_FOUND);

        const char *pttype = NULL;
        (void) sym_blkid_probe_lookup_value(b, "PTTYPE", &pttype, NULL);
        if (!streq_ptr(pttype, "gpt")) {
                log_device_debug(dev, "Didn't find a GPT partition table on '%s', falling back to BLKRRPART.", p);
                return fallback_ioctl(dev, fd, flags);
        }

        errno = 0;
        blkid_partlist pl = sym_blkid_probe_get_partitions(b);
        if (!pl)
                return log_device_debug_errno(dev, errno_or_else(EIO), "Unable to read partition table of '%s': %m", p);

        errno = 0;
        int n_partitions = sym_blkid_partlist_numof_partitions(pl);
        if (n_partitions < 0)
                return log_device_debug_errno(dev, errno_or_else(EIO), "Unable to acquire number of entries in partition table of '%s': %m", p);

        _cleanup_(sd_device_enumerator_unrefp) sd_device_enumerator *e = NULL;
        r = partition_enumerator_new(dev, &e);
        if (r < 0)
                return log_device_debug_errno(dev, r, "Failed to enumerate kernel partition devices: %m");

        log_device_debug(dev, "Updating/adding kernel partition devices...");

        _cleanup_(set_freep) Set *found_partnos = NULL;
        bool changed = false;
        int ret = 0;
        for (int i = 0; i < n_partitions; i++) {
                errno = 0;
                blkid_partition pp = sym_blkid_partlist_get_partition(pl, i);
                if (!pp)
                        return log_device_debug_errno(dev, errno_or_else(EIO), "Unable to get partition data of partition %i of partition table of '%s': %m", i, p);

                RET_GATHER(ret, process_partition(dev, fd, pp, e, &found_partnos, flags, &changed));
        }

        /* Only delete unrecognized partitions if everything else worked */
        if (ret < 0)
                return ret;

        log_device_debug(dev, "Removing old kernel partition devices...");

        r = remove_partitions(dev, fd, e, found_partnos, &changed);
        if (r < 0)
                return r;

        if (changed)
                return 1;

        if (FLAGS_SET(flags, REREADPT_FORCE_UEVENT)) {
                /* No change? Then trigger an event manually if we were told to */
                r = sd_device_trigger(dev, SD_DEVICE_CHANGE);
                if (r < 0)
                        return log_device_debug_errno(dev, r, "Failed to issue 'change' uevent on device '%s': %m", p);
        }

        return 0;
#else
        log_device_debug(dev, "We don't have libblkid, falling back to BLKRRPART on '%s'.", p);
        return fallback_ioctl(dev, fd, flags);
#endif
}

int reread_partition_table(sd_device *dev, RereadPartitionTableFlags flags) {
        assert(dev);

        _cleanup_close_ int fd = sd_device_open(dev, O_RDONLY|O_CLOEXEC|O_NONBLOCK|O_NOCTTY);
        if (fd < 0)
                return log_debug_errno(fd, "Failed to open block device: %m");

        return reread_partition_table_full(dev, fd, flags);
}

int reread_partition_table_fd(int fd, RereadPartitionTableFlags flags) {
        int r;

        _cleanup_(sd_device_unrefp) sd_device *dev = NULL;
        r = block_device_new_from_fd(fd, /* flags= */ 0, &dev);
        if (r < 0)
                return log_debug_errno(r, "Failed to get block device object: %m");

        return reread_partition_table_full(dev, fd, flags);
}
