//go:build linux

package host

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strconv"
	"strings"

	"github.com/carlosframework/platform/internal/rootstate"
	"golang.org/x/sys/unix"
)

type generationService struct {
	group, executable, selectedPath string
	main                            int
	selected                        rootstate.FileIdentity
}

// generationServiceEnv is what the service lookup reads besides the task's
// own /proc entries: systemd's view of the unit (show), the unit files'
// ownership (trust), the main process's cgroup (under proc), and, for a
// task still running a renamed-over binary, the box's signed router images
// (binary, under opt) and the owner every trusted file must have. The
// production values are generationServiceReader's; tests replace the parts
// a non-root test cannot provide.
type generationServiceEnv struct {
	proc, opt string
	owner     uint32
	show      func(ctx context.Context, unit string) (map[string]string, error)
	trust     func(path string) error
	binary    func(ctx context.Context, opt, path string) (rootFloorBinary, error)
}

// generationServiceRefusal is a refusal of the service lookup. It names the
// check, the unit, the task's cgroup and its exe link, so that refusals
// that used to read as a bare "permission denied" can be told apart, and
// it wraps the cause, so errors.Is still sees os.ErrPermission or a task
// that exited.
type generationServiceRefusal struct {
	check, unit, group, exe string
	err                     error
}

func (r *generationServiceRefusal) Error() string {
	return fmt.Sprintf("%s check refused: unit %s, cgroup %s, exe %s: %v", r.check, r.unit, r.group, r.exe, r.err)
}

func (r *generationServiceRefusal) Unwrap() error { return r.err }

// generationServiceReader is the production service lookup: classify a
// task, and recheck the executable binding of a task it trusted
// (generationKernel.recheck).
func generationServiceReader(ctx context.Context) (func(string, []byte) (bool, error), func(string, bool) error) {
	return generationServiceEnv{proc: "/proc", opt: "/opt/carlos", show: generationServiceShow, trust: generationServiceTrust, binary: readRootFloorBinaryContext}.reader(ctx)
}

// generationExeID is what binds a task to the image it was proven to run:
// the device and inode of the descriptor opened through /proc/<pid>/exe,
// the kernel's file handle for it where the filesystem gives one, and its
// mount ID. rootstate.ObjectIdentity carries the same handle but also
// resolves the mount in root's own mount table, which a service in a
// private mount namespace does not share, so only its kernel facts are
// used here.
type generationExeID struct {
	device, inode, mount uint64
	handle               string
	handleType           int32
}

func generationExeIdentity(f *os.File) (generationExeID, error) {
	var st unix.Stat_t
	if err := unix.Fstat(int(f.Fd()), &st); err != nil {
		return generationExeID{}, err
	}
	var sx unix.Statx_t
	if err := unix.Statx(int(f.Fd()), "", unix.AT_EMPTY_PATH, unix.STATX_MNT_ID, &sx); err != nil {
		return generationExeID{}, err
	}
	id := generationExeID{device: uint64(st.Dev), inode: st.Ino, mount: sx.Mnt_id}
	handle, _, err := unix.NameToHandleAt(int(f.Fd()), "", unix.AT_EMPTY_PATH)
	if err == nil {
		id.handle, id.handleType = hex.EncodeToString(handle.Bytes()), handle.Type()
	} else if !errors.Is(err, unix.EOPNOTSUPP) {
		return generationExeID{}, err
	}
	return id, nil
}

// generationExeBinding is a task the lookup trusted because the image it
// ran was proven: that image's descriptor, kept open for the inspection so
// its inode cannot be reused, and its identity.
type generationExeBinding struct {
	file             *os.File
	id               generationExeID
	unit, group, exe string
}

func generationServiceShow(ctx context.Context, unit string) (map[string]string, error) {
	out, err := exec.CommandContext(ctx, "systemctl", "show", unit, "-p", "ControlGroup", "-p", "MainPID", "-p", "ExecStart", "-p", "FragmentPath", "-p", "DropInPaths").Output()
	if err != nil {
		return nil, err
	}
	props := map[string]string{}
	for _, line := range strings.Split(string(out), "\n") {
		if k, v, ok := strings.Cut(line, "="); ok {
			props[k] = v
		}
	}
	return props, nil
}

func generationServiceTrust(file string) error {
	file, err := filepath.EvalSymlinks(file)
	if err != nil {
		return err
	}
	f, err := rootstate.OpenTrustedFile(file)
	if err != nil {
		return err
	}
	f.Close()
	return nil
}

func (env generationServiceEnv) reader(ctx context.Context) (func(string, []byte) (bool, error), func(string, bool) error) {
	units := map[string]generationService{}
	bound := map[string]*generationExeBinding{}
	var proven []string
	var provenErr error
	proved := false
	// The digests of the current and the previous router binary, each
	// proven against a signed identity the box holds, read once per census
	// pass and only if a task's exe has been deleted.
	prove := func() ([]string, error) {
		if !proved {
			proved = true
			var errs []error
			for _, name := range []string{"carlos", "carlos.prev"} {
				b, err := env.binary(ctx, env.opt, filepath.Join(env.opt, name))
				if err != nil {
					if name == "carlos" || !errors.Is(err, os.ErrNotExist) {
						errs = append(errs, fmt.Errorf("%s: %w", name, err))
					}
					continue
				}
				proven = append(proven, b.Digest)
			}
			if len(proven) == 0 {
				provenErr = errors.Join(append(errs, os.ErrPermission)...)
			}
		}
		return proven, provenErr
	}
	classify := func(path string, status []byte) (bool, error) {
		var held *os.File
		defer func() {
			if held != nil {
				held.Close()
			}
		}()
		if old := bound[path]; old != nil {
			old.file.Close()
			delete(bound, path)
		}
		name, nameErr := os.Readlink(path + "/exe")
		exe := name
		if nameErr != nil {
			exe = "unreadable"
		}
		unit, group := "unknown", "unknown"
		refuse := func(check string, err error) error {
			return &generationServiceRefusal{check: check, unit: unit, group: group, exe: exe, err: err}
		}
		raw, err := generationRead(path + "/cgroup")
		if err != nil {
			return false, refuse("cgroup", err)
		}
		group = strings.TrimPrefix(strings.TrimSpace(string(raw)), "0::")
		if !strings.HasPrefix(group, "/") || strings.Contains(group, "\n") || strings.Contains(group, "/carlos-instances.slice/") {
			return false, refuse("cgroup", os.ErrPermission)
		}
		unit = filepath.Base(group)
		for _, part := range strings.Split(group, "/") {
			if strings.HasSuffix(part, ".service") {
				unit = part
				break
			}
		}
		if !strings.HasSuffix(unit, ".service") && group != "/init.scope" {
			return false, refuse("unit", os.ErrPermission)
		}
		s, ok := units[unit]
		if !ok {
			if group == "/init.scope" {
				s = generationService{group: group, main: 1}
			} else {
				props, err := env.show(ctx, unit)
				if err != nil {
					return false, refuse("systemctl show", err)
				}
				s.group, s.main = props["ControlGroup"], 0
				s.main, err = strconv.Atoi(props["MainPID"])
				if err != nil || s.main <= 0 || props["FragmentPath"] == "" {
					return false, refuse("main pid", os.ErrPermission)
				}
				for _, file := range append([]string{props["FragmentPath"]}, strings.Fields(props["DropInPaths"])...) {
					if err := env.trust(file); err != nil {
						return false, refuse("unit file "+file, err)
					}
				}
				_, tail, found := strings.Cut(props["ExecStart"], "path=")
				if !found {
					return false, refuse("exec start", os.ErrPermission)
				}
				s.executable, _, _ = strings.Cut(tail, " ;")
				if unit == "carlos-edge.service" && s.executable == "/opt/carlos/system/guard/carlos" {
					if !generationGuardEdgeStart(unit, props["ExecStart"]) {
						return false, refuse("guard start", os.ErrPermission)
					}
					if _, err := env.binary(ctx, env.opt, s.executable); err != nil {
						return false, refuse("guard binary", err)
					}
					selected, err := env.binary(ctx, env.opt, filepath.Join(env.opt, "carlos"))
					if err != nil {
						return false, refuse("selected binary", err)
					}
					s.selected, s.selectedPath = selected.File, filepath.Join(env.opt, "carlos")
				}
			}
			units[unit] = s
		}
		supervisor := generationServiceSupervisor(unit, group, s.group, status, s.main)
		if !generationServiceGroup(group, s.group) && !supervisor {
			return false, refuse("group", os.ErrPermission)
		}
		if nameErr != nil {
			return false, refuse("exe link", nameErr)
		}

		// Non-root services may execute only the program selected by their
		// root-owned unit. UID and a service-looking cgroup are insufficient.
		_, rootErr := generationRootTask(status)
		checked := rootErr != nil || s.selected != (rootstate.FileIdentity{})
		if program, deleted := strings.CutSuffix(name, " (deleted)"); deleted && checked {
			want := s.executable
			if s.selected != (rootstate.FileIdentity{}) {
				want = s.selectedPath
			}
			image, err := env.deletedProgram(ctx, path, program, want, prove)
			if err != nil {
				return false, refuse("deleted exe", err)
			}
			held = image
		} else {
			actual, err := generationServiceExecutable(path, name)
			if err != nil {
				return false, refuse("exe", fmt.Errorf("%s: %w", name, err))
			}
			if checked {
				want, e := generationIdentity(s.executable)
				if s.selected != (rootstate.FileIdentity{}) {
					want, e = s.selected, nil
				}
				if e != nil || want != actual {
					return false, refuse("selected program", os.ErrPermission)
				}
			}
		}
		mainGroup, err := generationRead(env.proc + "/" + strconv.Itoa(s.main) + "/cgroup")
		mainPath := strings.TrimPrefix(strings.TrimSpace(string(mainGroup)), "0::")
		if err != nil || (!generationServiceGroup(mainPath, s.group) && !(supervisor && mainPath == group)) {
			if err == nil {
				err = os.ErrPermission
			}
			return false, refuse("main cgroup", err)
		}
		if held != nil {
			id, err := generationExeIdentity(held)
			if err != nil {
				return false, refuse("exe binding", err)
			}
			bound[path] = &generationExeBinding{file: held, id: id, unit: unit, group: group, exe: exe}
			held = nil
		}
		return true, nil
	}
	return classify, func(path string, final bool) error { return env.recheck(bound, path, final) }
}

// recheck holds a trusted task to the image it was proven to run, at each
// later use of that trust and at the end of its inspection (final).
// /proc/<pid>/exe, opened again, must be the very inode that was hashed (device, inode, mount and file handle),
// still a root-owned regular file with no links left and no group or other
// write. A task that exec'd anything else since, the same PID with the same
// start time, refuses; a later census pass may classify it afresh. The
// bytes are not hashed again: the identity is the hashed descriptor's, which
// stays open, and the kernel refuses writes to an image being executed
// (ETXTBSY). final releases the binding whatever the result. A task the
// lookup trusted without a proven image (its live exe was the unit's
// selected program, as on main) has no binding and passes.
func (env generationServiceEnv) recheck(bound map[string]*generationExeBinding, path string, final bool) error {
	b := bound[path]
	if b == nil {
		return nil
	}
	if final {
		defer func() {
			b.file.Close()
			delete(bound, path)
		}()
	}
	refuse := func(err error) error {
		return &generationServiceRefusal{check: "exe binding", unit: b.unit, group: b.group, exe: b.exe, err: err}
	}
	f, err := os.Open(path + "/exe")
	if err != nil {
		return refuse(err)
	}
	defer f.Close()
	id, err := generationExeIdentity(f)
	if err != nil {
		return refuse(err)
	}
	var st unix.Stat_t
	if err := unix.Fstat(int(f.Fd()), &st); err != nil {
		return refuse(err)
	}
	if id != b.id || st.Uid != env.owner || st.Mode&unix.S_IFMT != unix.S_IFREG || st.Mode&0022 != 0 || st.Nlink != 0 {
		return refuse(fmt.Errorf("the task's image is no longer the one proven: %w", os.ErrPermission))
	}
	return nil
}

// deletedProgram accepts a non-root service task still running a carlos
// binary that an apply has since renamed over, whose exe link reads
// "<program> (deleted)". The program must be the one the unit selects and
// the box's router binary, <opt>/carlos. The task's own image, opened
// through /proc/<pid>/exe, must be a root-owned regular file with no links
// left that no group or other can write. Its bytes must hash to the router
// artifact of the signed identity of the current or the previous binary
// (<opt>/carlos, <opt>/carlos.prev), proven by readRootFloorBinary, the
// proof the guard makes before it execs one. Anything else refuses, and
// the census's reference checks still run on an accepted task.
//
// On success it returns the image's descriptor, open: the lookup keeps it
// to bind the task to that image for the rest of its inspection (recheck).
func (env generationServiceEnv) deletedProgram(ctx context.Context, task, program, want string, prove func() ([]string, error)) (*os.File, error) {
	if program != want || want != filepath.Join(env.opt, "carlos") {
		return nil, fmt.Errorf("%s is not this unit's router binary: %w", program, os.ErrPermission)
	}
	f, err := os.Open(task + "/exe")
	if err != nil {
		return nil, err
	}
	keep := false
	defer func() {
		if !keep {
			f.Close()
		}
	}()
	var st unix.Stat_t
	if err := unix.Fstat(int(f.Fd()), &st); err != nil {
		return nil, err
	}
	if st.Uid != env.owner || st.Mode&unix.S_IFMT != unix.S_IFREG || st.Mode&0022 != 0 || st.Nlink != 0 || st.Size > 512<<20 {
		return nil, fmt.Errorf("image is not an unlinked root-owned file: %w", os.ErrPermission)
	}
	sum := sha256.New()
	if n, err := io.Copy(sum, generationContextReader{ctx: ctx, r: io.LimitReader(f, 512<<20+1)}); err != nil {
		return nil, err
	} else if n > 512<<20 {
		return nil, os.ErrPermission
	}
	digests, err := prove()
	if err != nil {
		return nil, err
	}
	if !slices.Contains(digests, hex.EncodeToString(sum.Sum(nil))) {
		return nil, fmt.Errorf("image is neither the current nor the previous signed router build: %w", os.ErrPermission)
	}
	keep = true
	return f, nil
}

func generationGuardEdgeStart(unit, start string) bool {
	if unit != "carlos-edge.service" || strings.Count(start, "{") != 1 || strings.Count(start, "}") != 1 {
		return false
	}
	fields := strings.Split(start, " ; ")
	return len(fields) >= 4 && fields[0] == "{ path=/opt/carlos/system/guard/carlos" &&
		fields[1] == "argv[]=/opt/carlos/system/guard/carlos system guard-exec edge" && fields[2] == "ignore_errors=no"
}

func generationServiceGroup(group, base string) bool {
	if group == base && base != "" {
		return true
	}
	if filepath.Clean(group) != group || base == "" || !strings.HasPrefix(group, base+"/") {
		return false
	}
	for p := group; ; p = filepath.Dir(p) {
		var st unix.Stat_t
		if unix.Lstat(filepath.Join("/sys/fs/cgroup", p), &st) != nil || st.Uid != 0 || st.Mode&0022 != 0 || st.Mode&unix.S_IFMT != unix.S_IFDIR {
			return false
		}
		if p == base {
			return true
		}
	}
}

func generationServiceExecutable(task, name string) (rootstate.FileIdentity, error) {
	var zero rootstate.FileIdentity
	if !filepath.IsAbs(name) || filepath.Clean(name) != name || name == "/" {
		return zero, os.ErrPermission
	}
	root, err := os.Open(task + "/root")
	if err != nil {
		return zero, err
	}
	defer root.Close()
	fd, err := unix.FcntlInt(root.Fd(), unix.F_DUPFD_CLOEXEC, 0)
	if err != nil {
		return zero, err
	}
	defer func() { unix.Close(fd) }()
	parts := strings.Split(strings.TrimPrefix(name, "/"), "/")
	for i := 0; ; i++ {
		var st unix.Stat_t
		if err := unix.Fstat(fd, &st); err != nil {
			return zero, err
		}
		kind := uint32(unix.S_IFDIR)
		if i == len(parts) {
			kind = unix.S_IFREG
		}
		if st.Uid != 0 || st.Mode&unix.S_IFMT != kind || st.Mode&0022 != 0 || (kind == unix.S_IFREG && st.Nlink != 1) {
			return zero, os.ErrPermission
		}
		if i == len(parts) {
			id := rootstate.FileIdentity{Device: uint64(st.Dev), Inode: st.Ino}
			actual, err := generationIdentity(task + "/exe")
			if err != nil || actual != id {
				return zero, os.ErrPermission
			}
			return id, nil
		}
		next, err := unix.Openat(fd, parts[i], unix.O_RDONLY|unix.O_NONBLOCK|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0)
		if err != nil {
			return zero, err
		}
		unix.Close(fd)
		fd = next
	}
}

func generationServiceSupervisor(unit, group, base string, status []byte, main int) bool {
	if unit != "carlos-edge.service" || base == "" || filepath.Clean(base) != base || group != base+"/supervise" || main <= 0 {
		return false
	}
	tgid := ""
	for _, line := range strings.Split(string(status), "\n") {
		if value, ok := strings.CutPrefix(line, "Tgid:"); ok {
			if tgid != "" {
				return false
			}
			tgid = strings.TrimSpace(value)
		}
	}
	return tgid == strconv.Itoa(main)
}