// 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 sys_test

import (
	"fmt"
	"os"
	"path"
	"testing"

	"github.com/google/go-cmp/cmp"
	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/rdma"
	"gvisor.dev/gvisor/pkg/sentry/fsimpl/sys"
	"gvisor.dev/gvisor/pkg/sentry/fsimpl/testutil"
	"gvisor.dev/gvisor/pkg/sentry/kernel"
	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
	"gvisor.dev/gvisor/pkg/sentry/vfs"
)

const (
	vfioDev = "vfio-dev"
)

func tryNewTestSystem(t *testing.T, idata *sys.InternalData) (*testutil.System, error) {
	k, err := testutil.Boot()
	if err != nil {
		t.Fatalf("Failed to create test kernel: %v", err)
	}
	ctx := k.SupervisorContext()
	creds := auth.CredentialsFromContext(ctx)
	k.VFS().MustRegisterFilesystemType(sys.Name, sys.FilesystemType{}, &vfs.RegisterFilesystemTypeOptions{
		AllowUserMount: true,
	})

	mountOpts := &vfs.MountOptions{
		GetFilesystemOptions: vfs.GetFilesystemOptions{
			InternalData: idata,
		},
	}

	mns, err := k.VFS().NewMountNamespace(ctx, creds, "", sys.Name, mountOpts, nil)
	if err != nil {
		return nil, err
	}
	return testutil.NewSystem(ctx, t, k.VFS(), mns), nil
}

func newTestSystemWithInternalData(t *testing.T, idata *sys.InternalData) *testutil.System {
	s, err := tryNewTestSystem(t, idata)
	if err != nil {
		t.Fatalf("Failed to create new mount namespace: %v", err)
	}
	return s
}

func newTestSystem(t *testing.T, pciTestDir string) *testutil.System {
	return newTestSystemWithInternalData(t, &sys.InternalData{
		EnableTPUProxyPaths: pciTestDir != "",
		TestSysfsPathPrefix: pciTestDir,
	})
}

func TestReadCPUFile(t *testing.T) {
	s := newTestSystem(t, "" /*pciTestDir*/)
	defer s.Destroy()
	k := kernel.KernelFromContext(s.Ctx)
	maxCPUCores := k.ApplicationCores()

	expected := fmt.Sprintf("0-%d\n", maxCPUCores-1)

	for _, fname := range []string{"online", "possible", "present"} {
		pop := s.PathOpAtRoot(fmt.Sprintf("devices/system/cpu/%s", fname))
		fd, err := s.VFS.OpenAt(s.Ctx, s.Creds, pop, &vfs.OpenOptions{})
		if err != nil {
			t.Fatalf("OpenAt(pop:%+v) = %+v failed: %v", pop, fd, err)
		}
		defer fd.DecRef(s.Ctx)
		content, err := s.ReadToEnd(fd)
		if err != nil {
			t.Fatalf("Read failed: %v", err)
		}
		if diff := cmp.Diff(expected, content); diff != "" {
			t.Fatalf("Read returned unexpected data:\n--- want\n+++ got\n%v", diff)
		}
	}
}

func TestSysRootContainsExpectedEntries(t *testing.T) {
	s := newTestSystem(t, "" /*pciTestDir*/)
	defer s.Destroy()
	pop := s.PathOpAtRoot("/")
	s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
		"block":    linux.DT_DIR,
		"bus":      linux.DT_DIR,
		"class":    linux.DT_DIR,
		"dev":      linux.DT_DIR,
		"devices":  linux.DT_DIR,
		"firmware": linux.DT_DIR,
		"fs":       linux.DT_DIR,
		"kernel":   linux.DT_DIR,
		"module":   linux.DT_DIR,
		"power":    linux.DT_DIR,
	})
}

func TestSysDevContainsBlockAndChar(t *testing.T) {
	s := newTestSystem(t, "" /*pciTestDir*/)
	defer s.Destroy()
	pop := s.PathOpAtRoot("/dev")
	s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
		"block": linux.DT_DIR,
		"char":  linux.DT_DIR,
	})
}

func TestCgroupMountpointExists(t *testing.T) {
	// Note: The mountpoint is only created if cgroups are available.
	s := newTestSystem(t, "" /*pciTestDir*/)
	defer s.Destroy()
	pop := s.PathOpAtRoot("/fs")
	s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
		"cgroup": linux.DT_DIR,
	})
	pop = s.PathOpAtRoot("/fs/cgroup")
	s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{ /*empty*/ })
}

// Check that sysfs creates the required PCI paths for V4 TPUs.
func TestEnableTPUProxyPathsV4(t *testing.T) {
	// Set up the fs tree that will be mirrored in the sentry.
	sysfsTestDir := t.TempDir()
	busPath := path.Join(sysfsTestDir, "sys", "bus", "pci", "devices")
	if err := os.MkdirAll(busPath, 0755); err != nil {
		t.Fatalf("Failed to create bus directory: %v", err)
	}
	classAccelPath := path.Join(sysfsTestDir, "sys", "class", "accel")
	if err := os.MkdirAll(classAccelPath, 0755); err != nil {
		t.Fatalf("Failed to create accel directory: %v", err)
	}
	for i, pciAddress := range []string{"0000:00:04.0", "0000:00:05.0"} {
		accelDev := fmt.Sprintf("accel%d", i)
		accelPath := path.Join(sysfsTestDir, "sys", "devices", "pci0000:00", pciAddress, "accel", accelDev)
		if err := os.MkdirAll(accelPath, 0755); err != nil {
			t.Fatalf("Failed to create accel directory: %v", err)
		}
		if err := os.Symlink(path.Join("..", "..", "..", pciAddress), path.Join(accelPath, pciAddress)); err != nil {
			t.Fatalf("Failed to symlink accel directory: %v", err)
		}
		if err := os.Symlink(path.Join("..", "..", "..", pciAddress), path.Join(accelPath, "device")); err != nil {
			t.Fatalf("Failed to symlink accel device directory: %v", err)
		}
		if _, err := os.Create(path.Join(accelPath, "chip_model")); err != nil {
			t.Fatalf("Failed to create chip_model: %v", err)
		}
		if _, err := os.Create(path.Join(accelPath, "device_owner")); err != nil {
			t.Fatalf("Failed to create device_owner: %v", err)
		}
		if _, err := os.Create(path.Join(accelPath, "pci_address")); err != nil {
			t.Fatalf("Failed to create pci_address: %v", err)
		}
		if err := os.Symlink(path.Join("..", "..", "..", "devices", "pci0000:00", pciAddress), path.Join(busPath, pciAddress)); err != nil {
			t.Fatalf("Failed to symlink bus directory: %v", err)
		}
		if err := os.Symlink(path.Join("..", "..", "devices", "pci0000:00", pciAddress, "accel", accelDev), path.Join(classAccelPath, accelDev)); err != nil {
			t.Fatalf("Failed to symlink accel directory: %v", err)
		}
	}

	s := newTestSystem(t, sysfsTestDir)
	defer s.Destroy()

	pop := s.PathOpAtRoot("/devices/pci0000:00")
	s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
		"0000:00:04.0": linux.DT_DIR,
		"0000:00:05.0": linux.DT_DIR,
	})
	pop = s.PathOpAtRoot("/devices/pci0000:00/0000:00:04.0/accel/accel0")
	s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
		"0000:00:04.0": linux.DT_LNK,
		"device":       linux.DT_LNK,
		"chip_model":   linux.DT_REG,
		"device_owner": linux.DT_REG,
		"pci_address":  linux.DT_REG,
	})
	pop = s.PathOpAtRoot("/bus/pci/devices")
	s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
		"0000:00:04.0": linux.DT_LNK,
		"0000:00:05.0": linux.DT_LNK,
	})
	pop = s.PathOpAtRoot("/class/accel")
	s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
		"accel0": linux.DT_LNK,
		"accel1": linux.DT_LNK,
	})
}

type PCIDeviceInfo struct {
	// IOMMU group.
	group             string
	pciPath           string
	pciAddress        string
	name              string
	nestedDeviceIndex int
}

func (dev PCIDeviceInfo) path() string {
	return path.Join(dev.pciPath, dev.pciAddress, vfioDev, dev.name)
}

func TestEnableTPUProxyPathsV5(t *testing.T) {
	// Set up the fs tree that will be mirrored in the sentry.
	sysfsTestDir := t.TempDir()
	pciPath0 := path.Join(sysfsTestDir, "sys", "devices", "pci0000:00")
	if err := os.MkdirAll(pciPath0, 0755); err != nil {
		t.Fatalf("Failed to create PCI directory: %v", err)
	}
	pciPath1 := path.Join(sysfsTestDir, "sys", "devices", "pci0000:10")
	if err := os.MkdirAll(pciPath1, 0755); err != nil {
		t.Fatalf("Failed to create PCI directory: %v", err)
	}
	busPath := path.Join(sysfsTestDir, "sys", "bus", "pci", "devices")
	if err := os.MkdirAll(busPath, 0755); err != nil {
		t.Fatalf("Failed to create bus directory: %v", err)
	}
	sysClassPath := path.Join(sysfsTestDir, "sys", "class", vfioDev)
	if err := os.MkdirAll(sysClassPath, 0755); err != nil {
		t.Fatalf("Failed to create class directory: %v", err)
	}

	devices := []PCIDeviceInfo{
		{
			group:             "0",
			pciPath:           pciPath0,
			pciAddress:        "0000:00:04.0",
			name:              "vfio0",
			nestedDeviceIndex: -1,
		},
		{
			group:             "1",
			pciPath:           pciPath0,
			pciAddress:        "0000:00:05.0",
			name:              "vfio1",
			nestedDeviceIndex: -1,
		},
		{
			group:             "2",
			pciPath:           pciPath1,
			pciAddress:        "0000:10:05.0",
			name:              "vfio2",
			nestedDeviceIndex: 3,
		},
		{
			group:             "3",
			pciPath:           pciPath1,
			pciAddress:        "0000:10:05.0/0000:03:00.1",
			name:              "vfio3",
			nestedDeviceIndex: -1,
		},
	}
	for _, device := range devices {
		devicePath := device.path()
		if err := os.MkdirAll(devicePath, 0755); err != nil {
			t.Fatalf("Failed to create PCI device directory: %v", err)
		}
		if err := os.Symlink(path.Join("..", "..", "..", device.pciAddress), path.Join(devicePath, "device")); err != nil {
			t.Fatalf("Failed to symlink device directory: %v", err)
		}
		if err := os.Symlink(path.Join("..", "..", "..", "devices", path.Base(device.pciPath), device.pciAddress), path.Join(busPath, path.Base(device.pciAddress))); err != nil {
			t.Fatalf("Failed to symlink bus directory: %v", err)
		}
		if err := os.Symlink(path.Join("..", "..", "devices", path.Base(device.pciPath), device.pciAddress, vfioDev, device.name), path.Join(sysClassPath, device.name)); err != nil {
			t.Fatalf("Failed to symlink class directory: %v", err)
		}
		iommuPath := path.Join(sysfsTestDir, "sys", "kernel", "iommu_groups", device.group, "devices")
		if err := os.MkdirAll(iommuPath, 0755); err != nil {
			t.Fatalf("Failed to create iommu_groups directory: %v", err)
		}
		if err := os.Symlink(path.Join("..", "..", "..", "..", "devices", path.Base(device.pciPath), device.pciAddress), path.Join(iommuPath, path.Base(device.pciAddress))); err != nil {
			t.Fatalf("Failed to symlink iommu_group devices directory: %v", err)
		}
		if err := os.Symlink(path.Join("..", "..", "..", "kernel", "iommu_groups", device.group), path.Join(device.pciPath, device.pciAddress, "iommu_group")); err != nil {
			t.Fatalf("Failed to symlink iommu_groups directory: %v", err)
		}
	}
	s := newTestSystem(t, sysfsTestDir)
	defer s.Destroy()

	for _, device := range devices {
		// Validate PCI device symlinks.
		pop := s.PathOpAtRoot(path.Join("devices", path.Base(device.pciPath), device.pciAddress))
		contents := map[string]testutil.DirentType{
			"iommu_group": linux.DT_LNK,
			vfioDev:       linux.DT_DIR,
		}
		if device.nestedDeviceIndex != -1 {
			deviceName := path.Base(devices[device.nestedDeviceIndex].pciAddress)
			contents[deviceName] = linux.DT_DIR
		}
		s.AssertAllDirentTypes(s.ListDirents(pop), contents)
		// Validate VFIO device symlinks.
		pop = s.PathOpAtRoot(path.Join("devices", path.Base(device.pciPath), device.pciAddress, vfioDev, device.name))
		s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
			"device": linux.DT_LNK,
		})
		// Validate $IOMMU_GROUP/devices.
		pop = s.PathOpAtRoot(path.Join("kernel", "iommu_groups", string(device.group), "devices"))
		s.AssertAllDirentTypes(s.ListDirents(pop), map[string]testutil.DirentType{
			path.Base(device.pciAddress): linux.DT_LNK,
		})
	}
}

// newRDMATestSnapshot returns a snapshot with a ConnectX NIC behind a bridge
// on domain 0000 and a GPU on an extended (VMD-style, 5-hex-digit) domain.
func newRDMATestSnapshot() *rdma.Snapshot {
	const nicLeaf = "devices/pci0000:07/0000:07:01.0/0000:0c:00.0"
	return &rdma.Snapshot{
		VerbsABIVersion: "6\n",
		PCINodes: []rdma.PCINode{
			{Path: "devices/pci0000:07"},
			{Path: "devices/pci0000:07/0000:07:01.0", Attrs: map[string]string{"class": "0x060400\n"}},
			{Path: nicLeaf, Attrs: map[string]string{
				"vendor": "0x15b3\n",
				"uevent": "PCI_SLOT_NAME=0000:0c:00.0\n",
				// Host-view CPU affinity: must be rewritten to the sandbox's
				// single-node view.
				"numa_node":     "1\n",
				"local_cpus":    "ffffffff,ffffff00,00000000\n",
				"local_cpulist": "56-111\n",
			}},
			{Path: "devices/pci10000:e0"},
			{Path: "devices/pci10000:e0/10000:e0:06.0", Attrs: map[string]string{
				"class":  "0x030200\n",
				"vendor": "0x10de\n",
				// A device on no NUMA node stays that way.
				"numa_node": "-1\n",
			}},
		},
		Devices: []rdma.Device{{
			Uverbs:     "uverbs0",
			IBDev:      "mlx5_0",
			LeafPCI:    nicLeaf,
			Dev:        "231:192\n",
			ABIVersion: "1\n",
			IBAttrs:    map[string]string{"node_guid": "0011:2233:4455:6677\n"},
			Ports: map[string]rdma.Port{
				"1": {
					StaticAttrs:  map[string]string{"link_layer": "Ethernet\n"},
					LiveAttrs:    []string{"state"},
					GIDNames:     []string{"0", "1"},
					CounterNames: []string{"port_rcv_data"},
				},
			},
			NetDevs: []rdma.NetDev{{Name: "eth1", Attrs: map[string]string{"mtu": "9000\n"}}},
		}},
		NUMA: &rdma.NUMA{
			// Two-node host (aggregate ranges "0-1"): the sandbox must
			// advertise a single node covering the sandbox's CPUs.
			Aggregate: map[string]string{"online": "0-1\n", "possible": "0-1\n"},
		},
	}
}

func readTestFile(s *testutil.System, p string) (string, error) {
	pop := s.PathOpAtRoot(p)
	fd, err := s.VFS.OpenAt(s.Ctx, s.Creds, pop, &vfs.OpenOptions{})
	if err != nil {
		return "", err
	}
	defer fd.DecRef(s.Ctx)
	return s.ReadToEnd(fd)
}

func TestRDMASysfs(t *testing.T) {
	s := newTestSystemWithInternalData(t, &sys.InternalData{RDMASysfs: newRDMATestSnapshot()})
	defer s.Destroy()

	const nicLeaf = "/devices/pci0000:07/0000:07:01.0/0000:0c:00.0"
	for p, want := range map[string]map[string]testutil.DirentType{
		"/devices/pci0000:07": {
			"0000:07:01.0": linux.DT_DIR,
			"pci_bus":      linux.DT_DIR,
		},
		"/devices/pci0000:07/0000:07:01.0": {
			"0000:0c:00.0": linux.DT_DIR,
			"pci_bus":      linux.DT_DIR,
			"subsystem":    linux.DT_LNK,
			"class":        linux.DT_REG,
		},
		nicLeaf: {
			"infiniband":       linux.DT_DIR,
			"infiniband_verbs": linux.DT_DIR,
			"net":              linux.DT_DIR,
			"subsystem":        linux.DT_LNK,
			"vendor":           linux.DT_REG,
			"uevent":           linux.DT_REG,
			"numa_node":        linux.DT_REG,
			"local_cpus":       linux.DT_REG,
			"local_cpulist":    linux.DT_REG,
		},
		nicLeaf + "/infiniband/mlx5_0": {
			"device":    linux.DT_LNK,
			"node_guid": linux.DT_REG,
			"ports":     linux.DT_DIR,
		},
		nicLeaf + "/infiniband/mlx5_0/ports/1": {
			"link_layer": linux.DT_REG,
			"state":      linux.DT_REG,
			"gids":       linux.DT_DIR,
			"gid_attrs":  linux.DT_DIR,
			"counters":   linux.DT_DIR,
		},
		nicLeaf + "/infiniband/mlx5_0/ports/1/gids": {
			"0": linux.DT_REG,
			"1": linux.DT_REG,
		},
		nicLeaf + "/infiniband_verbs/uverbs0": {
			"ibdev":       linux.DT_REG,
			"abi_version": linux.DT_REG,
			"dev":         linux.DT_REG,
			"device":      linux.DT_LNK,
		},
		nicLeaf + "/net/eth1": {
			"mtu":    linux.DT_REG,
			"device": linux.DT_LNK,
		},
		// Extended-domain (VMD-style) GPU root.
		"/devices/pci10000:e0": {
			"10000:e0:06.0": linux.DT_DIR,
			"pci_bus":       linux.DT_DIR,
		},
		"/class/infiniband":       {"mlx5_0": linux.DT_LNK},
		"/class/infiniband_verbs": {"uverbs0": linux.DT_LNK, "abi_version": linux.DT_REG},
		"/class/net":              {"eth1": linux.DT_LNK},
		"/class/pci_bus": {
			"0000:07":  linux.DT_LNK,
			"0000:0c":  linux.DT_LNK,
			"10000:e0": linux.DT_LNK,
		},
		"/bus/pci/devices": {
			"0000:07:01.0":  linux.DT_LNK,
			"0000:0c:00.0":  linux.DT_LNK,
			"10000:e0:06.0": linux.DT_LNK,
		},
		// The two host nodes collapse into a single sandbox node.
		"/devices/system/node": {"online": linux.DT_REG, "possible": linux.DT_REG, "node0": linux.DT_DIR},
		"/devices/system/node/node0": {
			"cpumap":   linux.DT_REG,
			"cpulist":  linux.DT_REG,
			"distance": linux.DT_REG,
		},
	} {
		pop := s.PathOpAtRoot(p)
		s.AssertAllDirentTypes(s.ListDirents(pop), want)
	}

	// The advertised NUMA view is a single node covering the sandbox's CPUs.
	wantCPUList := "0\n"
	if cores := kernel.KernelFromContext(s.Ctx).ApplicationCores(); cores > 1 {
		wantCPUList = fmt.Sprintf("0-%d\n", cores-1)
	}

	// Attribute contents, including via a /sys/class symlink (the path
	// consumers actually resolve).
	for p, want := range map[string]string{
		"/class/infiniband/mlx5_0/node_guid":                   "0011:2233:4455:6677\n",
		"/class/infiniband_verbs/abi_version":                  "6\n",
		nicLeaf + "/infiniband_verbs/uverbs0/ibdev":            "mlx5_0\n",
		"/class/infiniband/mlx5_0/ports/1/../../device/vendor": "0x15b3\n",
		"/devices/system/node/online":                          "0\n",
		"/devices/system/node/node0/cpulist":                   wantCPUList,
		"/devices/system/node/node0/distance":                  "10\n",
		nicLeaf + "/numa_node":                                 "0\n",
		nicLeaf + "/local_cpulist":                             wantCPUList,
		"/devices/pci10000:e0/10000:e0:06.0/numa_node":         "-1\n",
	} {
		got, err := readTestFile(s, p)
		if err != nil {
			t.Fatalf("reading %q: %v", p, err)
		}
		if got != want {
			t.Errorf("%q contains %q, want %q", p, got, want)
		}
	}
}
