/*
 * Copyright (C) 2018-2024 Canonical Ltd
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <dirent.h>
#include <errno.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

#include "config.h"

#include "../libsnap-confine-private/cleanup-funcs.h"
#include "../libsnap-confine-private/infofile.h"
#include "../libsnap-confine-private/mountinfo.h"
#include "../libsnap-confine-private/string-utils.h"
#include "../libsnap-confine-private/utils.h"

static sc_mountinfo_entry *find_dir_mountinfo(sc_mountinfo *mounts, const char *mnt_dir) {
    sc_mountinfo_entry *cur, *root = NULL;
    for (cur = sc_first_mountinfo_entry(mounts); cur != NULL; cur = sc_next_mountinfo_entry(cur)) {
        // Look for the mount info entry. We take the last one, which
        // would be the last mount on top of mnt_dir.
        if (sc_streq(mnt_dir, cur->mount_dir)) {
            root = cur;
        }
    }
    return root;
}

// Create a mount unit in normal_dir that is performed at early stages for
// "what" in directory "where".
// WARNING we need to escape special characters in "where" to create the unit
// name. We should do the same as systemd-escape(1), but for simplicity we just
// replace slashes with dashes, which is fine for the moment as this is called
// currently for mountpoints /usr/lib/{modules,firmware} only.
static int create_early_mount(const char *normal_dir, const char *what, const char *where) {
    // Replace directory separators with dashes to build the unit name.
    char *unit_name SC_CLEANUP(sc_cleanup_string) = NULL;
    // (... + 1) to remove the initial '/'
    unit_name = sc_strdup(where + 1);
    for (char *p = unit_name; (p = strchr(p, '/')) != NULL; *p = '-');

    // Construct the file name for a new systemd mount unit.
    char fname[PATH_MAX + 1] = {0};
    sc_must_snprintf(fname, sizeof fname, "%s/%s.mount", normal_dir, unit_name);

    // Open the mount unit and write the contents.
    FILE *f SC_CLEANUP(sc_cleanup_file) = NULL;
    f = fopen(fname, "w");
    if (!f) {
        fprintf(stderr, "cannot write to %s: %m\n", fname);
        return 1;
    }
    fprintf(f, "[Unit]\n");
    fprintf(f, "Description=Early mount of kernel drivers tree\n");
    fprintf(f, "DefaultDependencies=no\n");
    fprintf(f, "After=systemd-remount-fs.service\n");
    fprintf(f, "Before=sysinit.target\n");
    fprintf(f, "Before=systemd-udevd.service systemd-modules-load.service\n");
    fprintf(f, "Before=umount.target\n");
    fprintf(f, "Conflicts=umount.target\n");
    fprintf(f, "\n");
    fprintf(f, "[Mount]\n");
    fprintf(f, "What=%s\n", what);
    fprintf(f, "Where=%s\n", where);
    fprintf(f, "Options=bind,shared\n");

    // Wanted by sysinit.target.wants - create folders if needed and symlink

    char wants_d[PATH_MAX + 1] = {0};
    sc_must_snprintf(wants_d, sizeof wants_d, "%s/sysinit.target.wants", normal_dir);
    if (mkdir(wants_d, 0755) != 0 && errno != EEXIST) {
        fprintf(stderr, "cannot create %s directory: %m\n", wants_d);
        return 1;
    }

    char target[PATH_MAX + 1] = {0};
    char lnpath[PATH_MAX + 1] = {0};
    sc_must_snprintf(target, sizeof target, "../%s.mount", unit_name);
    sc_must_snprintf(lnpath, sizeof lnpath, "%s/%s.mount", wants_d, unit_name);
    if (symlink(target, lnpath) != 0) {
        fprintf(stderr, "cannot create symlink %s: %m\n", lnpath);
        return 1;
    }

    return 0;
}

#define MAJOR_LOOP_DEV 7
#define SNAPD_DRIVERS_TREE_DIR "/var/lib/snapd/kernel"
#define FIRMWARE_DIR "firmware"
#define MODULES_DIR "modules"
#define FIRMWARE_MNTPOINT "/usr/lib/" FIRMWARE_DIR
#define MODULES_MNTPOINT "/usr/lib/" MODULES_DIR

static int ensure_kernel_drivers_mounts(const char *normal_dir) {
    const char *const kern_mnt_dir = "/run/mnt/kernel";
    // Find mount information
    sc_mountinfo *mounts SC_CLEANUP(sc_cleanup_mountinfo) = NULL;
    mounts = sc_parse_mountinfo("/proc/1/mountinfo");
    if (!mounts) {
        fprintf(stderr, "cannot open or parse /proc/1/mountinfo\n");
        return 1;
    }
    // Create mount units only if not already present (which would be the
    // case for an old initramfs) - otherwise systemd-fstab-generator
    // complains, and older initramfs won't come in a kernel snap with
    // support for components anyway.
    const char *const kern_mntpts[] = {FIRMWARE_MNTPOINT, MODULES_MNTPOINT};
    for (size_t i = 0; i < SC_ARRAY_SIZE(kern_mntpts); ++i) {
        sc_mountinfo_entry *minfo = find_dir_mountinfo(mounts, kern_mntpts[i]);
        // If the mounts already exist (old initramfs), do not create them -
        // note that we additionally check for SNAPD_DRIVERS_TREE_DIR in the
        // mount source to make sure the units created here are still
        // generated on "systemctl daemon-reload".
        if (minfo && strstr(minfo->root, SNAPD_DRIVERS_TREE_DIR) == NULL) {
            return 0;
        }
    }

    // Find active kernel name and revision by looking at what was
    // mounted in /run/mnt/kernel by snap-bootstrap.

    sc_mountinfo_entry *kern_minfo = find_dir_mountinfo(mounts, kern_mnt_dir);
    if (!kern_minfo) {
        // This is not Ubuntu Core / hybrid, do nothing and do not fail
        return 0;
    }
    // Mount source should be a snap
    if (!sc_streq(kern_minfo->fs_type, "squashfs")) {
        fprintf(stderr, "unexpected fs type (%s) for %s\n", kern_minfo->fs_type, kern_mnt_dir);
        return 1;
    }
    // We expect a loop device as source
    if (kern_minfo->dev_major != MAJOR_LOOP_DEV) {
        fprintf(stderr, "mount source %s for %s is not a loop device\n", kern_minfo->mount_source, kern_mnt_dir);
        return 1;
    }
    // Find out backing file
    char fname[PATH_MAX + 1] = {0};
    sc_must_snprintf(fname, sizeof fname, "/sys/dev/block/%u:%u/loop/backing_file", kern_minfo->dev_major,
                     kern_minfo->dev_minor);
    FILE *f SC_CLEANUP(sc_cleanup_file) = NULL;
    f = fopen(fname, "r");
    if (!f) {
        fprintf(stderr, "cannot open %s: %m\n", fname);
        return 1;
    }
    char snap_path[PATH_MAX + 1] = {0};
    if (fgets(snap_path, sizeof snap_path, f) == NULL) {
        fprintf(stderr, "while reading %s: %m\n", fname);
        return 1;
    }
    // Now parse the snap path
    size_t i;
    for (i = strlen(snap_path); i > 0 && snap_path[--i] != '/';);
    char *snap_fname = snap_path + i + 1;

    // snap_fname is expected to contain "<name>_<rev>.snap\n" - fgets includes
    // that new line at the end, but anyway we ignore what comes after the dot.
    char *saveptr = NULL;
    char *snap_name = strtok_r(snap_fname, "_", &saveptr);
    if (snap_name == NULL) {
        fprintf(stderr, "snap name not found in loop backing file\n");
        return 1;
    }
    char *snap_rev = strtok_r(NULL, ".", &saveptr);
    if (snap_rev == NULL) {
        fprintf(stderr, "snap revision not found in loop backing file\n");
        return 1;
    }

    int res;
    char what[PATH_MAX + 1] = {0};
    sc_must_snprintf(what, sizeof what, SNAPD_DRIVERS_TREE_DIR "/%s/%s/lib/" MODULES_DIR, snap_name, snap_rev);
    res = create_early_mount(normal_dir, what, MODULES_MNTPOINT);
    if (res != 0) {
        return res;
    }
    sc_must_snprintf(what, sizeof what, SNAPD_DRIVERS_TREE_DIR "/%s/%s/lib/" FIRMWARE_DIR, snap_name, snap_rev);
    return create_early_mount(normal_dir, what, FIRMWARE_MNTPOINT);
}

static int ensure_root_fs_shared(const char *normal_dir) {
    // Load /proc/1/mountinfo so that we can inspect the root filesystem.
    sc_mountinfo *mounts SC_CLEANUP(sc_cleanup_mountinfo) = NULL;
    mounts = sc_parse_mountinfo("/proc/1/mountinfo");
    if (!mounts) {
        fprintf(stderr, "cannot open or parse /proc/1/mountinfo\n");
        return 1;
    }
    sc_mountinfo_entry *root = find_dir_mountinfo(mounts, "/");
    if (!root) {
        fprintf(stderr, "cannot find mountinfo entry of the root filesystem\n");
        return 1;
    }
    // Check if the root file-system is mounted with shared option.
    if (strstr(root->optional_fields, "shared:") != NULL) {
        // The workaround is not needed, everything is good as-is.
        return 0;
    }
    // Construct the file name for a new systemd mount unit.
    char fname[PATH_MAX + 1] = {0};
    sc_must_snprintf(fname, sizeof fname, "%s/" SNAP_MOUNT_DIR_SYSTEMD_UNIT ".mount", normal_dir);

    // Open the mount unit and write the contents.
    FILE *f SC_CLEANUP(sc_cleanup_file) = NULL;
    f = fopen(fname, "wt");
    if (!f) {
        fprintf(stderr, "cannot open %s: %m\n", fname);
        return 1;
    }
    fprintf(f,
            "# Ensure that snap mount directory is mounted \"shared\" "
            "so snaps can be refreshed correctly (LP: #1668759).\n");
    fprintf(f, "[Unit]\n");
    fprintf(f,
            "Description=Ensure that the snap directory "
            "shares mount events.\n");
    fprintf(f, "[Mount]\n");
    fprintf(f, "What=" STATIC_SNAP_MOUNT_DIR "\n");
    fprintf(f, "Where=" STATIC_SNAP_MOUNT_DIR "\n");
    fprintf(f, "Type=none\n");
    fprintf(f, "Options=bind,shared\n");

    /* We do not need to create symlinks from any target since
     * this generated mount will automically be added to implicit
     * dependencies of sub mount units through
     * `RequiresMountsFor`.
     */

    return 0;
}

static bool file_exists(const char *path) {
    struct stat buf;
    // Not using lstat to automatically resolve symbolic links,
    // including handling, as an error, dangling symbolic links.
    return stat(path, &buf) == 0 && (buf.st_mode & S_IFMT) == S_IFREG;
}

// PATH may not be set (the case on 16.04), in which case this is the fallback
// for looking up squashfuse / snapfuse executable.
// Based on what systemd uses when compiled for systems with "unmerged /usr"
// (see man systemd.exec).
static const char *const path_fallback = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin";

static bool executable_exists(const char *name) {
    char *path = getenv("PATH");
    char *path_copy SC_CLEANUP(sc_cleanup_string) = NULL;
    if (path == NULL) {
        path_copy = sc_strdup(path_fallback);
    } else {
        path_copy = sc_strdup(path);
    }

    char *ptr = NULL;
    char *token = strtok_r(path_copy, ":", &ptr);
    char fname[PATH_MAX + 1] = {0};
    while (token) {
        sc_must_snprintf(fname, sizeof fname, "%s/%s", token, name);
        if (access(fname, X_OK) == 0) {
            return true;
        }
        token = strtok_r(NULL, ":", &ptr);
    }
    return false;
}

static bool is_snap_try_snap_unit(const char *units_dir, const char *mount_unit_name) {
    char fname[PATH_MAX + 1] = {0};
    sc_must_snprintf(fname, sizeof fname, "%s/%s", units_dir, mount_unit_name);
    FILE *f SC_CLEANUP(sc_cleanup_file) = NULL;
    f = fopen(fname, "r");
    if (!f) {
        // not really expected
        fprintf(stderr, "cannot open mount unit %s: %m\n", fname);
        return false;
    }

    char *what SC_CLEANUP(sc_cleanup_string) = NULL;
    sc_error *err = NULL;
    if (sc_infofile_get_ini_section_key(f, "Mount", "What", &what, &err) < 0) {
        fprintf(stderr, "cannot read mount unit %s: %s\n", fname, sc_error_msg(err));
        sc_cleanup_error(&err);
        return false;
    }

    struct stat st;
    // if What points to a directory, then it's a snap try unit.
    return stat(what, &st) == 0 && (st.st_mode & S_IFMT) == S_IFDIR;
}

static int ensure_fusesquashfs_inside_container(const char *normal_dir) {
    // check if we are running inside a container, systemd
    // provides this file all the way back to trusty if run in a
    // container
    if (!file_exists("/run/systemd/container")) {
        return 0;
    }

    const char *fstype;
    if (executable_exists("squashfuse")) {
        fstype = "fuse.squashfuse";
    } else if (executable_exists("snapfuse")) {
        fstype = "fuse.snapfuse";
    } else {
        fprintf(stderr, "cannot find squashfuse or snapfuse executable\n");
        return 2;
    }

    DIR *units_dir SC_CLEANUP(sc_cleanup_closedir) = NULL;
    units_dir = opendir("/etc/systemd/system");
    if (units_dir == NULL) {
        // nothing to do
        return 0;
    }

    char fname[PATH_MAX + 1] = {0};

    struct dirent *ent;
    while ((ent = readdir(units_dir))) {
        // find snap mount units, i.e:
        // snap-somename.mount or var-lib-snapd-snap-somename.mount
        if (!sc_endswith(ent->d_name, ".mount")) {
            continue;
        }
        if (!(sc_startswith(ent->d_name, "snap-") || sc_startswith(ent->d_name, "var-lib-snapd-snap-"))) {
            continue;
        }
        if (is_snap_try_snap_unit("/etc/systemd/system", ent->d_name)) {
            continue;
        }
        sc_must_snprintf(fname, sizeof fname, "%s/%s.d", normal_dir, ent->d_name);
        if (mkdir(fname, 0755) != 0) {
            if (errno != EEXIST) {
                fprintf(stderr, "cannot create %s directory: %m\n", fname);
                return 2;
            }
        }

        sc_must_snprintf(fname, sizeof fname, "%s/%s.d/container.conf", normal_dir, ent->d_name);

        FILE *f SC_CLEANUP(sc_cleanup_file) = NULL;
        f = fopen(fname, "w");
        if (!f) {
            fprintf(stderr, "cannot open %s: %m\n", fname);
            return 2;
        }
        fprintf(f, "[Mount]\nType=%s\nOptions=nodev,ro,x-gdu.hide,x-gvfs-hide,allow_other\nLazyUnmount=yes\n", fstype);
    }

    return 0;
}

int main(int argc, char **argv) {
    if (argc != 4) {
        printf("usage: snapd-generator normal-dir early-dir late-dir\n");
        return 1;
    }
    const char *normal_dir = argv[1];
    // For reference, but we don't use those variables here.
    // const char *early_dir = argv[2];
    // const char *late_dir = argv[3];

    int status = 0;
    status = ensure_root_fs_shared(normal_dir);
    status |= ensure_fusesquashfs_inside_container(normal_dir);
    status |= ensure_kernel_drivers_mounts(normal_dir);

    return status;
}