// Copyright 2019 The gVisor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package sentrycmd

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"

	specs "github.com/opencontainers/runtime-spec/specs-go"
	"golang.org/x/sys/unix"
	"gvisor.dev/gvisor/pkg/abi/tpu"
	"gvisor.dev/gvisor/pkg/log"
	"gvisor.dev/gvisor/pkg/rdma"
	"gvisor.dev/gvisor/runsc/cmd/sandboxsetup"
	"gvisor.dev/gvisor/runsc/cmd/util"
	"gvisor.dev/gvisor/runsc/config"
	"gvisor.dev/gvisor/runsc/specutils"
)

// mountInChroot creates the destination mount point in the given chroot and
// mounts the source.
func mountInChroot(chroot, src, dst, typ string, flags uint32) error {
	chrootDst := filepath.Join(chroot, dst)
	log.Infof("Mounting %q at %q", src, chrootDst)

	if err := specutils.SafeSetupAndMount(src, chrootDst, typ, flags, "/proc"); err != nil {
		return fmt.Errorf("error mounting %q at %q: %v", src, chrootDst, err)
	}
	return nil
}

// bindMountSysfsInChrootReadonly bind-mounts the host sysfs directory src at
// dst inside the chroot read-only. mount(2) ignores attribute flags like
// MS_RDONLY when MS_BIND is set, so the bind must be followed by a remount
// that applies them. The remount also carries sysfs's standard attribute
// flags (nosuid,nodev,noexec), which the bind inherited and a remount may not
// clear.
func bindMountSysfsInChrootReadonly(chroot, src, dst string) error {
	if err := mountInChroot(chroot, src, dst, "bind", unix.MS_BIND); err != nil {
		return err
	}
	return specutils.SafeMount("", filepath.Join(chroot, dst), "",
		unix.MS_REMOUNT|unix.MS_BIND|unix.MS_RDONLY|unix.MS_NOSUID|unix.MS_NODEV|unix.MS_NOEXEC,
		"", "/proc")
}

// setupMinimalProcfs creates a minimal procfs-like tree at `${chroot}/proc`.
func setupMinimalProcfs(chroot string) error {
	// We can't always directly mount procfs because it may be obstructed
	// by submounts within it. See https://gvisor.dev/issue/10944.
	// All we really need from procfs is /proc/self and a few kernel
	// parameter files, which are typically not obstructed.
	// So we create a tmpfs at /proc and manually copy the kernel parameter
	// files into it. Then, to get /proc/self, we mount either a new
	// instance of procfs (if possible), or a recursive bind mount of the
	// procfs we do have access to (which still contains the obstructed
	// submounts but /proc/self is not obstructed), and we symlink
	// our /proc/self to the one in that mount.
	//
	// Why not try to mount the new procfs instance at /proc directly?
	// Because that would cause the set of files at /proc to differ
	// between the "new procfs instance" case and the "recursive bind
	// mount" case. Thus, this could introduce a bug whereby gVisor starts
	// to depend on a /proc file that is present in one case but not the
	// other, without decent test coverage to catch it.
	procRoot := filepath.Join(chroot, "/proc")
	if err := os.Mkdir(procRoot, 0755); err != nil {
		return fmt.Errorf("error creating /proc in chroot: %v", err)
	}
	if err := specutils.SafeMount("runsc-proc", procRoot, "tmpfs",
		unix.MS_NOSUID|unix.MS_NODEV|unix.MS_NOEXEC, "", "/proc"); err != nil {
		return fmt.Errorf("error mounting tmpfs in /proc: %v", err)
	}
	flags := uint32(unix.MS_NOSUID | unix.MS_NODEV | unix.MS_NOEXEC | unix.MS_RDONLY)
	procSubmountDir := "sandbox-proc"
	if newProcfsErr := mountInChroot(chroot, "proc", "/proc/"+procSubmountDir, "proc", flags); newProcfsErr != nil {
		log.Debugf("Unable to mount a new instance of the procfs file system at %q (%v); trying a recursive bind mount instead.", filepath.Join(procRoot, procSubmountDir), newProcfsErr)
		procSubmountDir = "host-proc"
		if bindErr := mountInChroot(chroot, "/proc", "/proc/"+procSubmountDir, "bind",
			unix.MS_BIND|unix.MS_REC|flags); bindErr != nil {
			return fmt.Errorf("error recursively bind-mounting proc at %q (%w) after also failing to mount a new procfs instance there (%v)", filepath.Join(procRoot, procSubmountDir), bindErr, newProcfsErr)
		}
		log.Debugf("Successfully mounted a recursive bind mount of procfs at %q; continuing.", filepath.Join(procRoot, procSubmountDir))
	}
	// Create needed directories.
	for _, d := range []string{
		"/proc/sys",
		"/proc/sys/kernel",
		"/proc/sys/vm",
	} {
		if err := os.Mkdir(filepath.Join(chroot, d), 0755); err != nil {
			return fmt.Errorf("error creating directory %q: %v", filepath.Join(chroot, d), err)
		}
	}
	// Copy needed files.
	for _, f := range []string{
		"/proc/sys/vm/mmap_min_addr",
		"/proc/sys/kernel/cap_last_cap",
	} {
		if err := sandboxsetup.CopyFile(filepath.Join(chroot, f), f); err != nil {
			return fmt.Errorf("failed to copy %q -> %q: %w", f, filepath.Join(chroot, f), err)
		}
	}
	// Create symlink for /proc/self.
	if err := os.Symlink(procSubmountDir+"/self", filepath.Join(procRoot, "self")); err != nil {
		return fmt.Errorf("error creating symlink %q -> %q: %w", filepath.Join(procRoot, "self"), procSubmountDir+"/self", err)
	}
	if err := os.Symlink(procSubmountDir+"/cpuinfo", filepath.Join(procRoot, "cpuinfo")); err != nil {
		return fmt.Errorf("error creating symlink %q -> %q: %w", filepath.Join(procRoot, "cpuinfo"), procSubmountDir+"/cpuinfo", err)
	}
	if err := os.Chmod(procRoot, 0o111); err != nil {
		return fmt.Errorf("error chmodding %q: %v", procRoot, err)
	}
	return nil
}

// setUpChroot creates an empty directory with runsc mounted at /runsc and proc
// mounted at /proc.
func setUpChroot(spec *specs.Spec, conf *config.Config) error {
	// We are a new mount namespace, so we can use /tmp as a directory to
	// construct a new root.
	chroot := os.TempDir()

	log.Infof("Setting up sandbox chroot in %q", chroot)

	// Convert all shared mounts into slave to be sure that nothing will be
	// propagated outside of our namespace.
	if err := specutils.SafeMount("", "/", "", unix.MS_SLAVE|unix.MS_REC, "", "/proc"); err != nil {
		return fmt.Errorf("error converting mounts: %v", err)
	}

	if err := specutils.SafeMount("runsc-root", chroot, "tmpfs", unix.MS_NOSUID|unix.MS_NODEV|unix.MS_NOEXEC, "", "/proc"); err != nil {
		return fmt.Errorf("error mounting tmpfs in chroot: %v", err)
	}

	if err := os.Mkdir(filepath.Join(chroot, "etc"), 0755); err != nil {
		return fmt.Errorf("error creating /etc in chroot: %v", err)
	}

	if err := sandboxsetup.CopyFile(filepath.Join(chroot, "etc/localtime"), "/etc/localtime"); err != nil {
		log.Warningf("Failed to copy /etc/localtime: %v. UTC timezone will be used.", err)
	}

	if err := setupMinimalProcfs(chroot); err != nil {
		return fmt.Errorf("error setting up minimal procfs in chroot %q: %v", chroot, err)
	}

	if err := tpuProxyUpdateChroot("/", chroot, spec, conf); err != nil {
		return fmt.Errorf("error configuring chroot for TPU devices: %w", err)
	}

	if err := rdmaSysfsUpdateChroot(chroot, spec, conf); err != nil {
		return fmt.Errorf("error configuring chroot for RDMA sysfs: %w", err)
	}

	if err := specutils.SafeMount("", chroot, "", unix.MS_REMOUNT|unix.MS_RDONLY|unix.MS_BIND, "", "/proc"); err != nil {
		return fmt.Errorf("error remounting chroot in read-only: %v", err)
	}

	return sandboxsetup.PivotRoot(chroot)
}

func tpuProxyUpdateChroot(hostRoot, chroot string, spec *specs.Spec, conf *config.Config) error {
	if !specutils.TPUProxyEnabled(spec, conf) {
		return nil
	}
	allowedDeviceIDs := map[uint64]struct{}{}
	paths, err := filepath.Glob(path.Join(hostRoot, "dev/vfio/*"))
	if err != nil {
		return fmt.Errorf("enumerating TPU device files: %w", err)
	}
	vfioDevicePath := path.Join(hostRoot, "dev/vfio/vfio")
	for _, devPath := range paths {
		if devPath == vfioDevicePath {
			continue
		}
		devNum := path.Base(devPath)
		iommuGroupPath := path.Join("/sys/kernel/iommu_groups", devNum)
		if err := bindMountSysfsInChrootReadonly(chroot, path.Join(hostRoot, iommuGroupPath), iommuGroupPath); err != nil {
			return fmt.Errorf("error mounting %q in chroot: %v", iommuGroupPath, err)
		}
		allowedDeviceIDs[tpu.TPUV5pDeviceID] = struct{}{}
		allowedDeviceIDs[tpu.TPUV5eDeviceID] = struct{}{}
		allowedDeviceIDs[tpu.TPUV6eDeviceID] = struct{}{}
		allowedDeviceIDs[tpu.TPU7xDeviceID] = struct{}{}
	}
	if len(allowedDeviceIDs) == 0 {
		paths, err = filepath.Glob(path.Join(hostRoot, "dev/accel*"))
		if err != nil {
			return fmt.Errorf("enumerating TPU device files: %w", err)
		}
		if len(paths) == 0 {
			return fmt.Errorf("could not find any TPU devices on the host")
		}
		allowedDeviceIDs[tpu.TPUV4DeviceID] = struct{}{}
		allowedDeviceIDs[tpu.TPUV4liteDeviceID] = struct{}{}
	}
	if len(allowedDeviceIDs) == 0 {
		return fmt.Errorf("no TPU devices found on the host")
	}
	sysDevicesGlob := path.Join(hostRoot, "/sys/devices/pci*")
	sysDevicesPaths, err := filepath.Glob(sysDevicesGlob)
	if err != nil {
		return fmt.Errorf("enumerating PCI device files: %w", err)
	}
	for _, sysDevicesPath := range sysDevicesPaths {
		if err := filepath.WalkDir(sysDevicesPath, func(path string, d os.DirEntry, err error) error {
			if d.Type().IsDir() && util.IsPCIDeviceDirTPU(path, allowedDeviceIDs) {
				chrootPath := strings.Replace(path, hostRoot, "/", 1)
				if err := bindMountSysfsInChrootReadonly(chroot, path, chrootPath); err != nil {
					return fmt.Errorf("error mounting %q in chroot: %v", path, err)
				}
			}
			return nil
		}); err != nil {
			return fmt.Errorf("walking %q: %w", sysDevicesPath, err)
		}
	}
	return err
}

// rdmaSysfsUpdateChroot collects the host's RDMA sysfs state into a JSON
// snapshot saved in the sandbox chroot, from which the sentry later builds the
// sandbox's virtual RDMA sysfs tree. Per-port dynamic files (GID tables, port
// state, counters) are not snapshotted; instead the host's per-ibdev "ports"
// subtrees are bind-mounted into the chroot read-only so the sentry can serve
// them live.
//
// Preconditions:
//   - Must run before pivot_root, while the host /sys is still reachable.
//   - The RDMA netdevs must still be in the host netns, so that they and their
//     /sys/class/net/<dev> attributes are visible in the host sysfs read here.
//
// The RDMA netdevs must therefore stay in the host netns until this collection
// runs; moving them into the sandbox netns should only be done afterwards.
//
// We deliberately do not support the reverse ordering - a sandbox netns
// pre-created and populated with the RDMA netdevs before `runsc create` - for
// two reasons:
//
//   - /sys/class/net entries are netns-tagged: a sysfs mount shows only the
//     netdevs of the netns that mounted it. Once a netdev is in the sandbox
//     netns it is gone from the host sysfs we read here, so its attributes
//     could only be collected via a sysfs mount tagged with the sandbox netns.
//   - The sandbox cannot make such a mount. Mounting a fresh sysfs requires the
//     current netns to be owned by the current userns (the kernel's
//     net_ns_current_may_mount check: net_ns->user_ns == current_user_ns).
//     Under directfs the sandbox runs in a new userns that does not own the
//     pre-created netns (that netns is owned by the host userns), so the mount
//     fails with EPERM.
//
// The per-ibdev "ports" bind mounts, by contrast, keep working after the
// netdevs move, because the InfiniBand class is not netns-tagged when the RDMA
// subsystem is in "rdma system netns shared" mode (a requirement for RDMA
// support here). GID tables and port state therefore stay readable through
// those mounts.
func rdmaSysfsUpdateChroot(chroot string, spec *specs.Spec, conf *config.Config) error {
	if !specutils.RDMAEnabled(spec, conf) {
		return nil
	}
	snap, err := rdma.Collect("/sys", specutils.UverbsDevicesInSpec(spec))
	if err != nil {
		return fmt.Errorf("collecting RDMA sysfs snapshot: %w", err)
	}
	if snap == nil {
		return nil
	}
	if err := snap.Save(filepath.Join(chroot, rdma.Path)); err != nil {
		return fmt.Errorf("saving RDMA sysfs snapshot: %w", err)
	}
	for i := range snap.Devices {
		// Bind-mount the host's per-ibdev "ports" subtree inside the chroot at
		// /sys/class/infiniband/<ibdev>/ports. On the host
		// /sys/class/infiniband/<ibdev> is a symlink into /sys/devices, and
		// the bind mount resolves it, so the real ports directory is what
		// lands in the chroot. In the sandbox the sentry builds the canonical
		// tree at /sys/devices/pci.../infiniband/<ibdev>/ (with
		// /sys/class/infiniband/<ibdev> as a symlink to it, mirroring the
		// host), and each live port file there records this
		// /sys/class/infiniband/<ibdev>/ports/... path as its host source; the
		// sentry opens that path to serve the file's contents live. So the
		// class path here need NOT match the sandbox's canonical (devices)
		// path since it only has to be where the sentry looks.
		portsPath := filepath.Join("/sys/class/infiniband", snap.Devices[i].IBDev, "ports")
		if err := bindMountSysfsInChrootReadonly(chroot, portsPath, portsPath); err != nil {
			return fmt.Errorf("bind-mounting %q read-only in chroot: %w", portsPath, err)
		}
	}
	return nil
}
