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

#include "sd-json.h"

#include "errno-util.h"
#include "fd-util.h"
#include "fdset.h"
#include "json-util.h"
#include "log.h"
#include "luo.h"
#include "luo-util.h"
#include "manager.h"
#include "serialize.h"
#include "service.h"
#include "unit.h"
#include "unit-name.h"

static int luo_read_mapping(int session_fd, sd_json_variant **ret) {
        int r;

        assert(session_fd >= 0);
        assert(ret);

        _cleanup_close_ int mapping_fd = luo_session_retrieve_fd(session_fd, LUO_MAPPING_INDEX);
        if (mapping_fd < 0)
                return log_warning_errno(mapping_fd, "Failed to retrieve LUO mapping fd (fd_index 0): %m");

        _cleanup_(sd_json_variant_unrefp) sd_json_variant *v = NULL;
        r = sd_json_parse_fd(
                        "luo-mapping",
                        mapping_fd,
                        SD_JSON_PARSE_MUST_BE_OBJECT|SD_JSON_PARSE_REOPEN_FD,
                        &v,
                        /* reterr_line= */ NULL,
                        /* reterr_column= */ NULL);
        if (r < 0)
                return log_warning_errno(r, "Failed to parse LUO mapping JSON: %m");

        *ret = TAKE_PTR(v);
        return 0;
}

static void luo_session_finishp(int *fd) {
        assert(fd);

        if (*fd >= 0)
                (void) luo_session_finish(*fd);
        safe_close(*fd);
}

int manager_luo_restore_fd_stores(Manager *m) {
        _cleanup_(sd_json_variant_unrefp) sd_json_variant *mapping = NULL;
        _cleanup_close_ int device_fd = -EBADF;
        _cleanup_(luo_session_finishp) int session_fd = -EBADF;
        const char *unit_id;
        sd_json_variant *unit_json;
        int r, n_total = 0;

        assert(m);

        if (MANAGER_IS_USER(m))
                return 0;

        device_fd = luo_open_device();
        if (ERRNO_IS_NEG_DEVICE_ABSENT(device_fd)) {
                log_debug_errno(device_fd, "No /dev/liveupdate device found, skipping LUO fd store restoration.");
                return 0;
        }
        if (device_fd < 0)
                return log_warning_errno(device_fd, "Failed to open /dev/liveupdate: %m");

        session_fd = luo_retrieve_session(device_fd, LUO_SESSION_NAME);
        if (session_fd == -ENOENT) {
                log_debug("No LUO session '%s' found, skipping fd store restoration.", LUO_SESSION_NAME);
                return 0;
        }
        if (session_fd < 0)
                return log_warning_errno(session_fd, "Failed to retrieve LUO session '%s': %m", LUO_SESSION_NAME);

        log_debug("Found LUO session '%s', restoring fd stores.", LUO_SESSION_NAME);

        r = luo_read_mapping(session_fd, &mapping);
        if (r < 0)
                return r;

        struct {
                uint64_t version;
                sd_json_variant *units;
        } q = {};

        static const sd_json_dispatch_field mapping_dispatch_table[] = {
                { "version", _SD_JSON_VARIANT_TYPE_INVALID, sd_json_dispatch_uint64,        voffsetof(q, version), SD_JSON_MANDATORY },
                { "units",   SD_JSON_VARIANT_OBJECT,        sd_json_dispatch_variant_noref, voffsetof(q, units),   0                 },
                {}
        };

        r = sd_json_dispatch(mapping, mapping_dispatch_table, SD_JSON_ALLOW_EXTENSIONS|SD_JSON_LOG, &q);
        if (r < 0)
                return r;

        if (q.version > LUO_PROTOCOL_VERSION) {
                log_warning("LUO mapping has unsupported version %" PRIu64 ", skipping state restoration.", q.version);
                return 0;
        }

        /* Retrieve all fds from the session and dispatch each to the named unit, eagerly loading the
         * unit if necessary. */
        JSON_VARIANT_OBJECT_FOREACH(unit_id, unit_json, q.units) {
                sd_json_variant *entry, *fds_json;

                if (!unit_name_is_valid(unit_id, UNIT_NAME_ANY)) {
                        log_warning("Invalid unit name '%s' in LUO mapping, skipping.", unit_id);
                        continue;
                }

                fds_json = sd_json_variant_by_key(unit_json, "fdstore");
                if (!sd_json_variant_is_array(fds_json)) {
                        log_warning("LUO mapping fdstore for unit '%s' is not a JSON array, skipping.", unit_id);
                        continue;
                }

                JSON_VARIANT_ARRAY_FOREACH(entry, fds_json) {
                        struct {
                                const char *type;
                                const char *name;
                                uint64_t token;
                                const char *session_name;
                        } p = {
                                .token = UINT64_MAX,
                        };

                        static const sd_json_dispatch_field dispatch_table[] = {
                                { "type",        SD_JSON_VARIANT_STRING,        sd_json_dispatch_const_string, voffsetof(p, type),         SD_JSON_MANDATORY },
                                { "name",        SD_JSON_VARIANT_STRING,        sd_json_dispatch_const_string, voffsetof(p, name),         SD_JSON_MANDATORY },
                                { "token",       _SD_JSON_VARIANT_TYPE_INVALID, sd_json_dispatch_uint64,       voffsetof(p, token),        0                 },
                                { "sessionName", SD_JSON_VARIANT_STRING,        sd_json_dispatch_const_string, voffsetof(p, session_name), 0                 },
                                {}
                        };

                        _cleanup_close_ int fd = -EBADF;

                        r = sd_json_dispatch(entry, dispatch_table, SD_JSON_ALLOW_EXTENSIONS|SD_JSON_LOG|SD_JSON_WARNING, &p);
                        if (r < 0)
                                continue;

                        if (streq(p.type, "fd")) {
                                if (p.token == UINT64_MAX) {
                                        log_warning("LUO mapping for unit '%s' fd '%s': missing 'token' field.", unit_id, p.name);
                                        continue;
                                }
                                if (p.token == LUO_MAPPING_INDEX) {
                                        log_warning("LUO mapping for unit '%s' fd '%s': token %" PRIu64 " is reserved for the mapping memfd.", unit_id, p.name, p.token);
                                        continue;
                                }

                                fd = luo_session_retrieve_fd(session_fd, p.token);
                                if (fd < 0) {
                                        log_warning_errno(fd, "Failed to retrieve LUO fd for unit '%s' name '%s' token %" PRIu64 ": %m",
                                                          unit_id, p.name, p.token);
                                        continue;
                                }
                        } else if (streq(p.type, "luo_session")) {
                                if (!p.session_name) {
                                        log_warning("LUO mapping for unit '%s' fd '%s': missing sessionName.", unit_id, p.name);
                                        continue;
                                }

                                fd = luo_retrieve_session(device_fd, p.session_name);
                                if (fd < 0) {
                                        log_warning_errno(fd, "Failed to retrieve LUO session '%s' for unit '%s' name '%s': %m",
                                                          p.session_name, unit_id, p.name);
                                        continue;
                                }

                                log_debug("Retrieved LUO session '%s' for unit fd store '%s'.", p.session_name, p.name);
                        } else {
                                log_warning("LUO mapping for unit '%s' fd '%s': unknown type '%s', skipping.",
                                            unit_id, p.name, p.type);
                                continue;
                        }

                        r = manager_dispatch_external_fd_to_unit(m, unit_id, p.name, /* index= */ 0, TAKE_FD(fd), "LUO");
                        if (r > 0)
                                n_total++;
                        /* On error fd is already consumed by manager_dispatch_external_fd_to_unit. */
                }
        }

        if (n_total > 0)
                log_debug("Restored %d fd(s) total from LUO session.", n_total);

        return n_total;
}

int manager_luo_serialize_fd_stores(Manager *m, FILE **ret_f, FDSet **ret_fds) {
        _cleanup_(sd_json_variant_unrefp) sd_json_variant *root = NULL, *units = NULL;
        _cleanup_fclose_ FILE *f = NULL;
        _cleanup_fdset_free_ FDSet *fds = NULL;
        Unit *u;
        int r, n_serialized = 0;

        assert(m);
        assert(ret_f);
        assert(ret_fds);

        if (MANAGER_IS_USER(m)) {
                *ret_f = NULL;
                *ret_fds = NULL;
                return 0;
        }

        fds = fdset_new();
        if (!fds)
                return log_oom();

        /* Build a JSON object: { "version": 1,
         *                        "units": { "unit_id": { "fdstore": [ { "type": "fd", "name": "...", "fd": N },
         *                                                             { "type": "luo_session", "name": "...", "fd": N, "sessionName": "..." } ] }, ... } }
         * This is passed to systemd-shutdown which will create a LUO session and preserve the fds. */
        HASHMAP_FOREACH(u, m->units) {
                _cleanup_(sd_json_variant_unrefp) sd_json_variant *entries = NULL;
                Service *s;

                if (u->type != UNIT_SERVICE)
                        continue;

                s = SERVICE(u);

                if (!IN_SET(s->fd_store_preserve_mode, EXEC_PRESERVE_YES, EXEC_PRESERVE_ON_SUCCESS))
                        continue;

                if (!s->fd_store)
                        continue;

                LIST_FOREACH(fd_store, fs, s->fd_store) {
                        _cleanup_free_ char *session_name = NULL;
                        int copy;

                        /* Check if this fd is itself a LUO session, as those cannot be nested and need
                         * special handling */
                        r = fd_get_luo_session_name(fs->fd, &session_name);
                        if (r < 0 && r != -EMEDIUMTYPE) {
                                log_warning_errno(r, "Failed to check if fd '%s' of unit '%s' is a LUO session, skipping: %m",
                                                  fs->fdname, u->id);
                                continue;
                        }

                        /* Ensure nobody tries to hijack our session, as we will create this later before
                         * kexec */
                        if (streq_ptr(session_name, LUO_SESSION_NAME)) {
                                log_warning("Skipping fd '%s' of unit '%s' for LUO serialization, as the session name '%s' infringes systemd's namespace.",
                                            fs->fdname, u->id, session_name);
                                continue;
                        }

                        copy = fdset_put_dup(fds, fs->fd);
                        if (copy < 0)
                                return log_error_errno(copy, "Failed to duplicate fd for LUO serialization: %m");

                        r = sd_json_variant_append_arraybo(
                                        &entries,
                                        SD_JSON_BUILD_PAIR_STRING("type", session_name ? "luo_session" : "fd"),
                                        SD_JSON_BUILD_PAIR_STRING("name", fs->fdname),
                                        SD_JSON_BUILD_PAIR_INTEGER("fd", copy),
                                        JSON_BUILD_PAIR_STRING_NON_EMPTY("sessionName", session_name));
                        if (r < 0)
                                return log_error_errno(r, "Failed to build JSON for LUO serialization: %m");

                        n_serialized++;
                }

                if (!entries)
                        continue;

                _cleanup_(sd_json_variant_unrefp) sd_json_variant *unit_json = NULL;
                r = sd_json_buildo(
                                &unit_json,
                                SD_JSON_BUILD_PAIR_VARIANT("fdstore", entries));
                if (r < 0)
                        return log_error_errno(r, "Failed to build LUO unit object: %m");

                r = sd_json_variant_set_field(&units, u->id, unit_json);
                if (r < 0)
                        return log_error_errno(r, "Failed to add unit to LUO serialization JSON: %m");
        }

        if (n_serialized == 0) {
                log_debug("No fd store entries to serialize for LUO.");
                *ret_f = NULL;
                *ret_fds = NULL;
                return 0;
        }

        r = sd_json_buildo(
                        &root,
                        SD_JSON_BUILD_PAIR_UNSIGNED("version", LUO_PROTOCOL_VERSION),
                        SD_JSON_BUILD_PAIR_CONDITION(!!units, "units", SD_JSON_BUILD_VARIANT(units)));
        if (r < 0)
                return log_error_errno(r, "Failed to build LUO serialization JSON: %m");

        r = open_serialization_file("luo-fd-store", &f);
        if (r < 0)
                return log_error_errno(r, "Failed to create LUO serialization file: %m");

        r = sd_json_variant_dump(root, /* flags= */ 0, f, /* prefix= */ NULL);
        if (r < 0)
                return log_error_errno(r, "Failed to dump LUO serialization JSON: %m");

        r = finish_serialization_file(f);
        if (r < 0)
                return log_error_errno(r, "Failed to finish LUO serialization file: %m");

        r = fd_cloexec(fileno(f), false);
        if (r < 0)
                return log_error_errno(r, "Failed to disable O_CLOEXEC for LUO serialization: %m");

        r = fdset_cloexec(fds, false);
        if (r < 0)
                return log_error_errno(r, "Failed to disable O_CLOEXEC for LUO serialization fds: %m");

        log_info("Serialized %d fd store entries for LUO.", n_serialized);

        *ret_f = TAKE_PTR(f);
        *ret_fds = TAKE_PTR(fds);
        return n_serialized;
}
