//go:build linux && rootrunner

package host

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"syscall"
	"testing"

	"golang.org/x/sys/unix"
)

// consoleUpgradeFixture is a protected console directory installed by
// install-console with build A, and source files for builds A to D.
type consoleUpgradeFixture struct {
	t      *testing.T
	d, dir string
	src    map[string]string // build name to source path
	body   map[string][]byte
	sha    map[string]string
}

func newConsoleUpgradeFixture(t *testing.T) *consoleUpgradeFixture {
	t.Helper()
	d := activationRootFixture(t)
	f := &consoleUpgradeFixture{t: t, d: d, dir: filepath.Join(d, "opt", "system", "console"), src: map[string]string{}, body: map[string][]byte{}, sha: map[string]string{}}
	f.check(os.Chmod(d, 0o755))
	f.check(os.MkdirAll(filepath.Join(d, "src"), 0o755))
	for _, name := range []string{"A", "B", "C", "D"} {
		f.body[name] = []byte("console build " + name)
		f.sha[name] = artifactSHA(f.body[name])
		f.src[name] = filepath.Join(d, "src", name)
		f.check(os.WriteFile(f.src[name], f.body[name], 0o755))
	}
	if got, err := installProtectedConsole(f.dir, f.src["A"], f.sha["A"], false); err != nil || got.Already {
		t.Fatal("first install", got, err)
	}
	f.want("A", "A")
	return f
}

func (f *consoleUpgradeFixture) check(err error) {
	f.t.Helper()
	if err != nil {
		f.t.Fatal(err)
	}
}

func (f *consoleUpgradeFixture) path(name string) string { return filepath.Join(f.dir, name) }

// slot requires name to be a root:root 0755 one-link regular file holding
// build b.
func (f *consoleUpgradeFixture) slot(name, b string) {
	f.t.Helper()
	got, err := os.ReadFile(f.path(name))
	st, serr := os.Lstat(f.path(name))
	if err != nil || serr != nil || !bytes.Equal(got, f.body[b]) {
		f.t.Fatalf("%s: want build %s, got %q %v %v", name, b, got, err, serr)
	}
	sys := st.Sys().(*syscall.Stat_t)
	if st.Mode() != 0o755 || sys.Uid != 0 || sys.Gid != 0 || sys.Nlink != 1 {
		f.t.Fatalf("%s mode %v uid %d gid %d links %d", name, st.Mode(), sys.Uid, sys.Gid, sys.Nlink)
	}
}

// entries lists the directory's names.
func (f *consoleUpgradeFixture) entries() string {
	f.t.Helper()
	list, err := os.ReadDir(f.dir)
	f.check(err)
	var names []string
	for _, e := range list {
		names = append(names, e.Name())
	}
	sort.Strings(names)
	return strings.Join(names, ",")
}

// want requires console to hold current, console.prev to hold prev and
// nothing else in the directory.
func (f *consoleUpgradeFixture) want(current, prev string) {
	f.t.Helper()
	f.slot("console", current)
	f.slot("console.prev", prev)
	if got := f.entries(); got != "console,console.prev" {
		f.t.Fatal("directory holds", got)
	}
}

// upgrade runs the upgrade to build to, from the sha of build from.
func (f *consoleUpgradeFixture) upgrade(to, from string, dry bool) (consoleInstalled, error) {
	return upgradeProtectedConsole(f.dir, f.src[to], f.sha[to], f.sha[from], dry)
}

// refused requires the upgrade to build to from build from to refuse, dry
// and not, and the directory listing to be unchanged by either.
func (f *consoleUpgradeFixture) refused(what, to, from string) {
	f.t.Helper()
	before := f.entries()
	for _, dry := range []bool{true, false} {
		if got, err := f.upgrade(to, from, dry); err == nil {
			f.t.Fatal("accepted", what, dry, got)
		}
		if after := f.entries(); after != before {
			f.t.Fatal(what, "changed the directory", before, after)
		}
	}
}

// TestRootActivationConsoleUpgrade is install-console --upgrade as real
// root: the swap, its refusals, crash recovery, a leftover temp file, the
// Already rerun and the dry run.
func TestRootActivationConsoleUpgrade(t *testing.T) {
	f := newConsoleUpgradeFixture(t)

	// A wrong --from refuses and changes nothing; so does a malformed one.
	f.refused("wrong from", "B", "C")
	f.want("A", "A")
	for _, from := range []string{strings.ToUpper(f.sha["A"]), f.sha["A"][:63], ""} {
		for _, dry := range []bool{true, false} {
			if _, err := upgradeProtectedConsole(f.dir, f.src["B"], f.sha["B"], from, dry); err == nil {
				t.Fatalf("accepted from %q", from)
			}
		}
	}
	// A source that does not hash to the target refuses.
	if _, err := upgradeProtectedConsole(f.dir, f.src["C"], f.sha["B"], f.sha["A"], false); err == nil {
		t.Fatal("accepted another source")
	}
	f.want("A", "A")

	// The dry run checks and reports, and changes nothing.
	got, err := f.upgrade("B", "A", true)
	if err != nil || got.Already || got.From != f.sha["A"] || got.PreviousSHA != f.sha["A"] || got.Path != f.path("console") || got.Previous != f.path("console.prev") {
		t.Fatal("dry run", got, err)
	}
	f.want("A", "A")

	// Upgrade A to B: prev is A, current is B, no temp file left.
	got, err = f.upgrade("B", "A", false)
	if err != nil || got.Already || got.From != f.sha["A"] || got.PreviousSHA != f.sha["A"] {
		t.Fatal("upgrade", got, err)
	}
	f.want("B", "A")

	// A rerun is Already and changes nothing, dry or not.
	for _, dry := range []bool{true, false} {
		got, err = f.upgrade("B", "A", dry)
		if err != nil || !got.Already || got.PreviousSHA != f.sha["A"] {
			t.Fatal("rerun", dry, got, err)
		}
		f.want("B", "A")
	}

	// A malformed --from refuses even when the console is already current.
	for _, from := range []string{strings.ToUpper(f.sha["A"]), f.sha["A"][:63], ""} {
		for _, dry := range []bool{true, false} {
			if _, err := upgradeProtectedConsole(f.dir, f.src["B"], f.sha["B"], from, dry); err == nil {
				t.Fatalf("Already with from %q", from)
			}
		}
	}

	// Without --upgrade an occupied slot still refuses, as before.
	for _, b := range []string{"B", "C"} {
		for _, dry := range []bool{true, false} {
			if _, err := installProtectedConsole(f.dir, f.src[b], f.sha[b], dry); err == nil {
				t.Fatal("first install over an installed console", b, dry)
			}
		}
	}
	f.want("B", "A")

	// The current console's shape: each refuses, nothing changes.
	cur := f.path("console")
	shape := func(what string, setup, undo func()) {
		t.Helper()
		setup()
		f.refused(what, "C", "B")
		undo()
		f.want("B", "A")
	}
	shape("current symlink", func() { f.check(os.Rename(cur, cur+".real")); f.check(os.Symlink(cur+".real", cur)) }, func() { f.check(os.Remove(cur)); f.check(os.Rename(cur+".real", cur)) })
	shape("current not root's", func() { f.check(os.Chown(cur, 1000, 1000)) }, func() { f.check(os.Chown(cur, 0, 0)) })
	shape("current owner not root", func() { f.check(os.Chown(cur, 1000, 0)) }, func() { f.check(os.Chown(cur, 0, 0)) })
	shape("current group not root", func() { f.check(os.Chown(cur, 0, 1000)) }, func() { f.check(os.Chown(cur, 0, 0)) })
	shape("current two links", func() { f.check(os.Link(cur, cur+".link")) }, func() { f.check(os.Remove(cur + ".link")) })
	shape("current group writable", func() { f.check(os.Chmod(cur, 0o775)) }, func() { f.check(os.Chmod(cur, 0o755)) })
	shape("current not executable by others", func() { f.check(os.Chmod(cur, 0o750)) }, func() { f.check(os.Chmod(cur, 0o755)) })
	shape("current directory", func() { f.check(os.Rename(cur, cur+".real")); f.check(os.Mkdir(cur, 0o755)) }, func() { f.check(os.Remove(cur)); f.check(os.Rename(cur+".real", cur)) })

	// console.prev's shape: a symlink, a directory or the wrong owner each
	// refuses.
	prev := f.path("console.prev")
	shape("prev symlink", func() { f.check(os.Rename(prev, prev+".real")); f.check(os.Symlink(prev+".real", prev)) }, func() { f.check(os.Remove(prev)); f.check(os.Rename(prev+".real", prev)) })
	shape("prev directory", func() { f.check(os.Rename(prev, prev+".real")); f.check(os.Mkdir(prev, 0o755)) }, func() { f.check(os.Remove(prev)); f.check(os.Rename(prev+".real", prev)) })
	shape("prev not root's", func() { f.check(os.Chown(prev, 1000, 1000)) }, func() { f.check(os.Chown(prev, 0, 0)) })
	shape("prev owner not root", func() { f.check(os.Chown(prev, 1000, 0)) }, func() { f.check(os.Chown(prev, 0, 0)) })
	shape("prev two links", func() { f.check(os.Link(prev, prev+".link")) }, func() { f.check(os.Remove(prev + ".link")) })
	shape("prev group writable", func() { f.check(os.Chmod(prev, 0o775)) }, func() { f.check(os.Chmod(prev, 0o755)) })

	// console.prev's bytes are not checked: it is about to be replaced.
	f.check(os.WriteFile(prev+".x", []byte("an older build"), 0o755))
	f.check(os.Rename(prev+".x", prev))
	if got, err = f.upgrade("C", "B", false); err != nil || got.Already || got.PreviousSHA != f.sha["B"] {
		t.Fatal("upgrade over any prev", got, err)
	}
	f.want("C", "B")

	// A run cut off after renaming console to console.prev: console is
	// absent, console.prev holds the old build and console.next the new.
	f.check(os.WriteFile(f.path(consoleUpgradeTemp), f.body["D"], 0o755))
	f.check(os.Rename(cur, prev))
	if got := f.entries(); got != "console.next,console.prev" {
		t.Fatal("crash state", got)
	}
	f.refused("crash state, wrong from", "D", "B")
	f.refused("crash state, from the new build", "D", "D")
	if got, err = f.upgrade("D", "C", true); err != nil || got.Already || got.PreviousSHA != f.sha["C"] {
		t.Fatal("crash state dry run", got, err)
	}
	if got := f.entries(); got != "console.next,console.prev" {
		t.Fatal("crash state dry run changed", got)
	}
	if got, err = f.upgrade("D", "C", false); err != nil || got.Already || got.From != f.sha["C"] || got.PreviousSHA != f.sha["C"] {
		t.Fatal("crash recovery", got, err)
	}
	f.want("D", "C")

	// With both slots gone there is nothing to upgrade from.
	f.check(os.Rename(cur, cur+".keep"))
	f.check(os.Rename(prev, prev+".keep"))
	f.refused("no console", "A", "D")
	f.refused("no console, from prev", "A", "C")
	f.check(os.Rename(cur+".keep", cur))
	f.check(os.Rename(prev+".keep", prev))

	// A run cut off before its renames leaves a partial console.next: the
	// dry run leaves it, a rerun removes it and completes.
	f.check(os.WriteFile(f.path(consoleUpgradeTemp), []byte("partial"), 0o600))
	if _, err = f.upgrade("A", "D", true); err != nil {
		t.Fatal("dry run over a leftover", err)
	}
	if b, _ := os.ReadFile(f.path(consoleUpgradeTemp)); string(b) != "partial" {
		t.Fatal("dry run touched the leftover")
	}
	if got, err = f.upgrade("A", "D", false); err != nil || got.Already || got.PreviousSHA != f.sha["D"] {
		t.Fatal("upgrade over a leftover", got, err)
	}
	f.want("A", "D")

	// A console.next that is not a regular file refuses and is left.
	f.check(os.Mkdir(f.path(consoleUpgradeTemp), 0o700))
	f.refused("console.next directory", "B", "A")
	f.check(os.Remove(f.path(consoleUpgradeTemp)))
	f.check(os.Symlink(f.src["B"], f.path(consoleUpgradeTemp)))
	f.refused("console.next symlink", "B", "A")
	f.check(os.Remove(f.path(consoleUpgradeTemp)))
	f.want("A", "D")

	// An upgrade never creates the directory.
	missing := filepath.Join(f.d, "opt", "system", "other")
	for _, dry := range []bool{true, false} {
		if _, err := upgradeProtectedConsole(missing, f.src["B"], f.sha["B"], f.sha["A"], dry); err == nil {
			t.Fatal("upgraded a missing directory", dry)
		}
	}
	if _, err := os.Lstat(missing); !os.IsNotExist(err) {
		t.Fatal("upgrade created the directory", err)
	}

	// A group-writable directory refuses, as for the first install.
	f.check(os.Chmod(f.dir, 0o775))
	f.refused("group-writable directory", "B", "A")
	f.check(os.Chmod(f.dir, 0o755))
	f.want("A", "D")
}

// TestRootActivationConsoleUpgradeIdentity runs the upgrade through
// runInstallConsole: the staged identity must be the published one and
// name the running build, or nothing changes.
func TestRootActivationConsoleUpgradeIdentity(t *testing.T) {
	f := newConsoleUpgradeFixture(t)
	w := newArtifactWorld(t)
	identity := w.run(t, w.identityPlan(t)).Output
	idPath := filepath.Join(f.d, "identity.json")
	f.check(os.WriteFile(idPath, identity, 0o644))
	f.body["I"] = []byte("console build")
	f.sha["I"] = artifactSHA(f.body["I"])
	f.src["I"] = filepath.Join(f.d, "src", "I")
	f.check(os.WriteFile(f.src["I"], f.body["I"], 0o755))

	for _, build := range []string{"v0.28.1", "v0.28.3", ""} {
		for _, dry := range []bool{true, false} {
			if _, _, got, err := runInstallConsole(t.Context(), w.store, w.pub, idPath, f.src["I"], f.dir, build, "linux-arm64", f.sha["A"], dry); err == nil {
				t.Fatal("accepted build", build, got)
			}
		}
	}
	for _, dry := range []bool{true, false} {
		if _, _, _, err := runInstallConsole(t.Context(), w.store, w.pub, idPath, f.src["I"], f.dir, "v0.28.2", "linux-amd64", f.sha["A"], dry); err == nil {
			t.Fatal("accepted another platform")
		}
	}
	f.want("A", "A")

	tag, sha, got, err := runInstallConsole(t.Context(), w.store, w.pub, idPath, f.src["I"], f.dir, "v0.28.2", "linux-arm64", f.sha["A"], true)
	if err != nil || tag != "v0.28.2" || sha != f.sha["I"] || got.Already || got.PreviousSHA != f.sha["A"] {
		t.Fatal("dry run", tag, sha, got, err)
	}
	f.want("A", "A")
	if _, _, got, err = runInstallConsole(t.Context(), w.store, w.pub, idPath, f.src["I"], f.dir, "v0.28.2", "linux-arm64", f.sha["A"], false); err != nil || got.Already {
		t.Fatal("upgrade", got, err)
	}
	f.want("I", "A")
	if _, _, got, err = runInstallConsole(t.Context(), w.store, w.pub, idPath, f.src["I"], f.dir, "v0.28.2", "linux-arm64", f.sha["A"], false); err != nil || !got.Already {
		t.Fatal("rerun", got, err)
	}
	// Without --upgrade the installed console is still refused.
	if _, _, _, err = runInstallConsole(t.Context(), w.store, w.pub, idPath, f.src["I"], f.dir, "v0.28.2", "linux-arm64", "", false); err == nil {
		t.Fatal("first install over an upgraded console")
	}
	f.want("I", "A")
}

// TestRootActivationConsoleUpgradeSourceAlias: an upgrade never takes its
// source from its own slots or temp file, whether named by a path inside
// the directory or reached through a bind mount of it (same device and
// inode, another path). Each refuses with nothing changed.
func TestRootActivationConsoleUpgradeSourceAlias(t *testing.T) {
	f := newConsoleUpgradeFixture(t)
	// console = A, console.prev = B: the target B is in console.prev.
	prev := f.path("console.prev")
	f.check(os.WriteFile(prev+".x", f.body["B"], 0o755))
	f.check(os.Rename(prev+".x", prev))
	alias := filepath.Join(f.d, "alias")
	f.check(os.Mkdir(alias, 0o755))
	f.check(unix.Mount(f.dir, alias, "", unix.MS_BIND, ""))
	t.Cleanup(func() { unix.Unmount(alias, unix.MNT_DETACH) })

	refuse := func(what, source, to, from string) {
		t.Helper()
		before := f.entries()
		for _, dry := range []bool{true, false} {
			if got, err := upgradeProtectedConsole(f.dir, source, f.sha[to], f.sha[from], dry); err == nil {
				t.Fatal("accepted", what, dry, got)
			}
			if after := f.entries(); after != before {
				t.Fatal(what, "changed the directory", before, after)
			}
		}
	}
	refuse("console.prev by path", prev, "B", "A")
	refuse("console.prev through a bind mount", filepath.Join(alias, "console.prev"), "B", "A")
	refuse("console by path", f.path("console"), "A", "B")
	refuse("console through a bind mount", filepath.Join(alias, "console"), "A", "B")
	// Another file inside the directory, not a slot, is refused by path.
	f.check(os.WriteFile(f.path("console.src"), f.body["C"], 0o755))
	refuse("a file in the directory", f.path("console.src"), "C", "A")
	f.check(os.Remove(f.path("console.src")))

	// A leftover console.next holding the target.
	f.check(os.WriteFile(f.path(consoleUpgradeTemp), f.body["C"], 0o755))
	refuse("console.next by path", f.path(consoleUpgradeTemp), "C", "A")
	refuse("console.next through a bind mount", filepath.Join(alias, consoleUpgradeTemp), "C", "A")
	if b, _ := os.ReadFile(f.path(consoleUpgradeTemp)); !bytes.Equal(b, f.body["C"]) {
		t.Fatal("console.next changed")
	}
	f.check(os.Remove(f.path(consoleUpgradeTemp)))
	f.slot("console", "A")
	f.slot("console.prev", "B")

	// The same target from a source outside the directory upgrades.
	if got, err := f.upgrade("B", "A", false); err != nil || got.Already {
		t.Fatal("upgrade from an outside source", got, err)
	}
	f.want("B", "A")
}

// TestRootActivationConsoleUpgradeAlreadySyncs: a run cut off after its
// second rename, before the directory fsync, leaves console = target; the
// rerun is Already and fsyncs the directory first, and a failed fsync is
// the rerun's error. A dry rerun does not sync.
func TestRootActivationConsoleUpgradeAlreadySyncs(t *testing.T) {
	f := newConsoleUpgradeFixture(t)
	var dirSt unix.Stat_t
	f.check(unix.Stat(f.dir, &dirSt))
	var calls int
	var fail error
	saved := consoleUpgradeSync
	t.Cleanup(func() { consoleUpgradeSync = saved })
	consoleUpgradeSync = func(fd int) error {
		var st unix.Stat_t
		if err := unix.Fstat(fd, &st); err != nil || st.Dev != dirSt.Dev || st.Ino != dirSt.Ino {
			t.Fatal("synced something other than the directory", err)
		}
		calls++
		if fail != nil {
			return fail
		}
		return saved(fd)
	}

	// The state just before the final fsync.
	f.check(os.WriteFile(f.path(consoleUpgradeTemp), f.body["B"], 0o755))
	f.check(os.Rename(f.path("console"), f.path("console.prev")))
	f.check(os.Rename(f.path(consoleUpgradeTemp), f.path("console")))
	f.want("B", "A")

	if got, err := f.upgrade("B", "A", true); err != nil || !got.Already || calls != 0 {
		t.Fatal("dry rerun", got, err, calls)
	}
	if got, err := f.upgrade("B", "A", false); err != nil || !got.Already || calls != 1 {
		t.Fatal("rerun", got, err, calls)
	}
	fail = syscall.EIO
	if got, err := f.upgrade("B", "A", false); err != syscall.EIO || calls != 2 {
		t.Fatal("rerun with a failing fsync", got, err, calls)
	}
	fail = nil
	f.want("B", "A")

	// A full upgrade syncs once, after its renames.
	if got, err := f.upgrade("C", "B", false); err != nil || got.Already || calls != 3 {
		t.Fatal("upgrade", got, err, calls)
	}
	f.want("C", "B")
}

// TestRootActivationConsoleUpgradeOversizedSlot: a slot larger than
// consoleMaxBytes is refused, not hashed by its prefix, so --from naming
// the prefix's digest cannot match it.
func TestRootActivationConsoleUpgradeOversizedSlot(t *testing.T) {
	f := newConsoleUpgradeFixture(t)
	cur := f.path("console")
	big := cur + ".big"
	f.check(os.WriteFile(big, nil, 0o755))
	f.check(os.Truncate(big, consoleMaxBytes+2)) // sparse: zeros
	f.check(os.Rename(big, cur))
	prefix := zeroDigest(t, consoleMaxBytes+1) // what a LimitReader(max+1) hashes
	f.sha["prefix"] = prefix
	f.refused("oversized console", "B", "prefix")
	if _, err := os.Lstat(cur); err != nil {
		t.Fatal(err)
	}
	// As console.prev in the crash state it is refused too.
	f.check(os.Rename(cur, f.path("console.prev")))
	f.refused("oversized console.prev", "B", "prefix")
}

// zeroDigest is the sha256 of n zero bytes.
func zeroDigest(t *testing.T, n int64) string {
	t.Helper()
	h := sha256.New()
	buf := make([]byte, 1<<20)
	for n > 0 {
		k := int64(len(buf))
		if n < k {
			k = n
		}
		h.Write(buf[:k])
		n -= k
	}
	return hex.EncodeToString(h.Sum(nil))
}