llmgw/internal/agent/slurmnode/users.go

package slurmnode

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"slices"
	"strconv"
	"strings"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/slurm"
)

// usersMarker lists, a line each, the logins llmgw made here: the ones it
// may change and take away. A user that was here before is never touched.
func (k *Keeper) usersMarker() string { return k.marker + "-users" }

// users keeps the Linux users want has, the same login and UID as on every
// other machine, their keys where sshd reads them, and their homes when
// this machine makes them; and takes away those llmgw made that are gone.
// Answers how many it has. A user that cannot be made is said, the others
// are made anyway.
func (k *Keeper) users(ctx context.Context, want *llmgwv1.Slurm) (int32, error) {
	made := k.madeUsers()
	if !want.GetManageUsers() {
		return 0, k.removeUsers(ctx, made, nil)
	}
	var errs []error
	var keep []string
	n := int32(0)
	for _, u := range want.GetUsers() {
		err := k.user(ctx, u, slices.Contains(made, u.GetLogin()), want.GetMakeHomes())
		if errors.Is(err, errNotMade) {
			errs = append(errs, err)
			continue
		}
		keep = append(keep, u.GetLogin())
		if err != nil {
			errs = append(errs, err)
			continue
		}
		n++
	}
	if err := k.keepMade(keep); err != nil {
		return n, err
	}
	if err := k.removeUsers(ctx, made, keep); err != nil {
		errs = append(errs, err)
	}
	if err := k.sshConf(ctx, len(keep) > 0); err != nil {
		errs = append(errs, err)
	}
	// The controller is where they log in and submit from; a node only
	// while a job of theirs runs there.
	if err := k.nodeLogin(ctx, len(keep) > 0 && want.GetNode() && !want.GetController()); err != nil {
		errs = append(errs, err)
	}
	return n, errors.Join(errs...)
}

// nodeLogin holds the users llmgw made here to their jobs when they log in
// (pam_slurm_adopt), or no longer. Only they are: root and the site's
// administrators log in as before, so a mistake here never locks out the
// people who would put it right.
func (k *Keeper) nodeLogin(ctx context.Context, on bool) error {
	old, err := k.host.ReadFile(slurm.PAMSSHD)
	if err != nil {
		if !on {
			return nil
		}
		return fmt.Errorf("reading %s: %w", slurm.PAMSSHD, err)
	}
	if on {
		if err := k.adoptModule(ctx); err != nil {
			return err
		}
	}
	read := func(name string) (string, error) {
		b, err := k.host.ReadFile(filepath.Join(filepath.Dir(slurm.PAMSSHD), name))
		return string(b), err
	}
	stack, contained, err := slurm.PAMAdopt(string(old), k.usersMarker(), read, on)
	if err != nil {
		return err
	}
	if stack == string(old) {
		return nil
	}
	if on {
		k.log.Info("llmgw's Slurm users log in to this node only while a job of theirs runs here",
			"in_the_job", contained)
		if !contained {
			k.log.Warn("pam_systemd could not be skipped for them safely here: their ssh sessions are not held " +
				"to their job's cores and cards, only to having one")
		}
	} else {
		k.log.Info("llmgw's Slurm users are no longer held to their jobs when they log in here")
	}
	// Read by sshd for each login: nothing to reload.
	_, err = k.write(slurm.PAMSSHD, []byte(stack), 0o644, "")
	return err
}

// adoptModule is there, or is installed.
func (k *Keeper) adoptModule(ctx context.Context) error {
	if k.installed[slurm.AdoptModule] {
		return nil
	}
	find := func() bool {
		for _, d := range slurm.ModuleDirs {
			if _, err := k.host.Run(ctx, nil, "test", "-f", filepath.Join(d, slurm.AdoptModule)); err == nil {
				return true
			}
		}
		return false
	}
	if !find() {
		if k.inst == nil {
			return fmt.Errorf("%s missing, and this agent installs nothing", slurm.AdoptModule)
		}
		k.log.Info("installing pam_slurm_adopt, which holds llmgw's Slurm users to their jobs here")
		if err := k.inst.InstallPackages(ctx, slurm.AdoptPackages); err != nil {
			return err
		}
		if !find() {
			return fmt.Errorf("%s is still missing after installing %s", slurm.AdoptModule,
				strings.Join(slurm.AdoptPackages["deb"], ", "))
		}
	}
	k.installed[slurm.AdoptModule] = true
	return nil
}

// errNotMade is a user llmgw could not make here, nor took for its own.
var errNotMade = errors.New("not made")

type notMade struct{ msg string }

func (e notMade) Error() string        { return e.msg }
func (e notMade) Is(target error) bool { return target == errNotMade }

// user keeps one user; ours is whether llmgw made it here before.
func (k *Keeper) user(ctx context.Context, u *llmgwv1.SlurmUser, ours, makeHome bool) error {
	login, uid, home := u.GetLogin(), u.GetUid(), u.GetHome()
	if !slurm.ValidLogin(login) || uid < slurm.MinUID || !filepath.IsAbs(home) || strings.Contains(home, "..") {
		return notMade{fmt.Sprintf("the leader sent a user %q, UID %d, home %q", login, uid, home)}
	}
	id := strconv.FormatUint(uint64(uid), 10)
	gecos := strings.Map(func(r rune) rune {
		if r == ':' || r == ',' || r == '\n' || r < ' ' {
			return -1
		}
		return r
	}, u.GetName())
	byName, okName := k.passwd(ctx, login)
	byID, okID := k.passwd(ctx, id)
	switch {
	case okName && byName[2] != id:
		return notMade{fmt.Sprintf("%s is already a user here with UID %s, not llmgw's %s", login, byName[2], id)}
	case !okName && okID:
		return notMade{fmt.Sprintf("UID %s is already %s's here, so %s is not made", id, byID[0], login)}
	case !okName:
		if _, ok := k.group(ctx, login); !ok {
			if out, err := k.host.Run(ctx, nil, "groupadd", "-g", id, login); err != nil {
				return notMade{fmt.Sprintf("groupadd %s: %v: %s", login, err, lastLine(out))}
			}
		}
		k.log.Info("a Slurm user is made here, as llmgw's accounts say", "login", login, "uid", uid)
		if out, err := k.host.Run(ctx, nil, "useradd", "-u", id, "-g", id, "-M", "-d", home, "-s", "/bin/bash",
			"-c", gecos, login); err != nil {
			return notMade{fmt.Sprintf("useradd %s: %v: %s", login, err, lastLine(out))}
		}
	case ours && (byName[4] != gecos || byName[5] != home):
		if out, err := k.host.Run(ctx, nil, "usermod", "-c", gecos, "-d", home, login); err != nil {
			return fmt.Errorf("usermod %s: %v: %s", login, err, lastLine(out))
		}
	}
	keys := strings.Join(u.GetSshKeys(), "\n")
	file := filepath.Join(slurm.AuthorizedKeysDir, login)
	if keys == "" {
		_ = k.host.Remove(file)
	} else if _, err := k.write(file, []byte(keys+"\n"), 0o644, ""); err != nil {
		return err
	}
	if makeHome {
		return k.home(ctx, login, id, home)
	}
	return nil
}

// home makes a user's home directory, from /etc/skel, when it is not there.
func (k *Keeper) home(ctx context.Context, login, id, home string) error {
	if _, err := k.host.Run(ctx, nil, "test", "-d", home); err == nil {
		return nil
	}
	if err := k.host.MkdirAll(filepath.Dir(home), 0o755, ""); err != nil {
		return err
	}
	k.log.Info("a Slurm user's home is made", "login", login, "home", home)
	for _, c := range [][]string{
		{"install", "-d", "-m", "0700", "-o", id, "-g", id, home},
		{"cp", "-rT", "/etc/skel", home},
		{"chown", "-R", id + ":" + id, home},
	} {
		if out, err := k.host.Run(ctx, nil, c[0], c[1:]...); err != nil {
			return fmt.Errorf("making %s's home: %s: %v: %s", login, c[0], err, lastLine(out))
		}
	}
	return nil
}

// removeUsers takes away the users llmgw made here that are not in keep.
// Their homes and files stay; their UID is never given to another.
func (k *Keeper) removeUsers(ctx context.Context, made, keep []string) error {
	var errs []error
	left := slices.Clone(keep)
	for _, login := range made {
		if slices.Contains(keep, login) {
			continue
		}
		k.log.Info("a Slurm user llmgw made is taken away; its home stays", "login", login)
		_ = k.host.Remove(filepath.Join(slurm.AuthorizedKeysDir, login))
		if _, ok := k.passwd(ctx, login); ok {
			if out, err := k.host.Run(ctx, nil, "userdel", login); err != nil {
				// Still logged in, or running something: tried again.
				errs = append(errs, fmt.Errorf("userdel %s: %v: %s", login, err, lastLine(out)))
				left = append(left, login)
				continue
			}
		}
	}
	if !slices.Equal(left, made) {
		if err := k.keepMade(left); err != nil {
			errs = append(errs, err)
		}
	}
	if len(left) == 0 {
		if err := k.sshConf(ctx, false); err != nil {
			errs = append(errs, err)
		}
		if err := k.nodeLogin(ctx, false); err != nil {
			errs = append(errs, err)
		}
	}
	return errors.Join(errs...)
}

func (k *Keeper) madeUsers() []string {
	b, err := k.host.ReadFile(k.usersMarker())
	if err != nil {
		return nil
	}
	return strings.Fields(string(b))
}

func (k *Keeper) keepMade(logins []string) error {
	s := slices.Clone(logins)
	slices.Sort(s)
	s = slices.Compact(s)
	if len(s) == 0 {
		return k.host.Remove(k.usersMarker())
	}
	_, err := k.write(k.usersMarker(), []byte(strings.Join(s, "\n")+"\n"), 0o644, "")
	return err
}

// sshConf has sshd read the users' keys from AuthorizedKeysDir, or no
// longer.
func (k *Keeper) sshConf(ctx context.Context, on bool) error {
	if !on {
		if _, err := k.host.ReadFile(slurm.SSHDropIn); err != nil {
			return nil
		}
		_ = k.host.Remove(slurm.SSHDropIn)
		return k.reloadSSH(ctx)
	}
	changed, err := k.write(slurm.SSHDropIn, []byte(slurm.SSHConf), 0o644, "")
	if err != nil || !changed {
		return err
	}
	if out, err := k.host.Run(ctx, nil, "sshd", "-t"); err != nil {
		_ = k.host.Remove(slurm.SSHDropIn)
		return fmt.Errorf("sshd refuses llmgw's setting for the users' keys, so it is not kept: %v: %s", err, lastLine(out))
	}
	k.log.Info("sshd reads the Slurm users' keys from llmgw too", "dir", slurm.AuthorizedKeysDir)
	return k.reloadSSH(ctx)
}

// reloadSSH has sshd read its settings again, under Debian's name for it
// or Red Hat's; one that is not running reads them when it starts.
func (k *Keeper) reloadSSH(ctx context.Context) error {
	for _, u := range []string{"ssh", "sshd"} {
		if k.active(ctx, u) {
			return k.systemctl(ctx, "reload", u)
		}
	}
	return nil
}

// passwd is getent's line for a user, by login or UID, in its fields.
func (k *Keeper) passwd(ctx context.Context, key string) ([]string, bool) {
	out, err := k.host.Run(ctx, nil, "getent", "passwd", key)
	f := strings.Split(strings.TrimSpace(string(out)), ":")
	if err != nil || len(f) < 7 {
		return nil, false
	}
	return f, true
}

func (k *Keeper) group(ctx context.Context, key string) ([]string, bool) {
	out, err := k.host.Run(ctx, nil, "getent", "group", key)
	f := strings.Split(strings.TrimSpace(string(out)), ":")
	if err != nil || len(f) < 3 {
		return nil, false
	}
	return f, true
}