llmgw/internal/agent/supervisor/runas.go

package supervisor

import (
	"fmt"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"syscall"
)

// Engines run as their own user, not as the agent's root: an engine reads a
// model's files and answers requests, and a model's files are whatever a
// registry published. Root in that process would be root on the node for
// anything that can make an engine misbehave — a malformed file, a template,
// a flag nobody thought of.
//
// The user is DefaultRunAs, made by `llmgw node add`, in the video and render
// groups for the GPU devices. A node without it keeps running engines as
// root, and says so, so an existing fleet is not broken by an upgrade.

// DefaultRunAs is the user engines run as.
const DefaultRunAs = "llmgw-engine"

// runAs is how an engine process is started as another user: its
// credential, a working directory it owns, and the environment that points
// its caches there.
type runAs struct {
	name string
	cred *syscall.Credential
	home string
}

// lookupUser is user.Lookup, replaced by tests.
var lookupUser = user.Lookup

// resolveRunAs finds the user engines run as. nil means as this process —
// when it is not root, there is nobody to drop to; when the user does not
// exist, why is returned.
func resolveRunAs(name, workDir string) (*runAs, string) {
	if name == "" || name == "root" || os.Geteuid() != 0 {
		return nil, ""
	}
	u, err := lookupUser(name)
	if err != nil {
		return nil, fmt.Sprintf("there is no user %s, so engines run as root; `llmgw node add` makes it, "+
			"or: useradd --system --shell /usr/sbin/nologin %s && usermod -aG video,render %s", name, name, name)
	}
	r, err := credentialFor(u, workDir)
	if err != nil {
		return nil, fmt.Sprintf("engines run as root: %v", err)
	}
	return r, ""
}

func credentialFor(u *user.User, workDir string) (*runAs, error) {
	uid, err1 := strconv.ParseUint(u.Uid, 10, 32)
	gid, err2 := strconv.ParseUint(u.Gid, 10, 32)
	if err1 != nil || err2 != nil {
		return nil, fmt.Errorf("user %s has ids %s:%s", u.Username, u.Uid, u.Gid)
	}
	var groups []uint32
	if ids, err := u.GroupIds(); err == nil {
		for _, g := range ids {
			if n, err := strconv.ParseUint(g, 10, 32); err == nil {
				groups = append(groups, uint32(n))
			}
		}
	}
	// A home of its own under the work directory, which is root's: caches
	// (Triton, torch, Hugging Face) are written there and nowhere else.
	home := filepath.Join(workDir, "engine-home")
	if err := os.MkdirAll(home, 0o700); err != nil {
		return nil, fmt.Errorf("making %s: %w", home, err)
	}
	if err := os.Chown(home, int(uid), int(gid)); err != nil {
		return nil, fmt.Errorf("giving %s to %s: %w", home, u.Username, err)
	}
	// The work directory is 0750 root's; the engine has to pass through it
	// to its home.
	if err := os.Chmod(workDir, 0o751); err != nil {
		return nil, err
	}
	return &runAs{name: u.Username, home: home,
		cred: &syscall.Credential{Uid: uint32(uid), Gid: uint32(gid), Groups: groups}}, nil
}

// env is what is added to an engine's environment when it runs as r: the
// last value of a key wins in exec.
func (r *runAs) env() []string {
	return []string{"HOME=" + r.home, "XDG_CACHE_HOME=" + filepath.Join(r.home, ".cache"), "USER=" + r.name, "LOGNAME=" + r.name}
}

// AsEngineUser has cmd run as the engine user, as an engine would, with its
// home under workDir: for work the agent does that is not an engine but
// reads what an engine reads — training an adapter. It answers that home
// and a way to hand a path to the user, or why it runs as this process.
func AsEngineUser(cmd *exec.Cmd, name, workDir string) (home string, give func(path string) error, why string) {
	r, why := resolveRunAs(name, workDir)
	if r == nil {
		return "", func(string) error { return nil }, why
	}
	cmd.SysProcAttr = &syscall.SysProcAttr{Credential: r.cred, Setpgid: true}
	cmd.Env = append(cmd.Env, r.env()...)
	uid, gid := int(r.cred.Uid), int(r.cred.Gid)
	return r.home, func(path string) error {
		return filepath.WalkDir(path, func(p string, _ os.DirEntry, err error) error {
			if err != nil {
				return err
			}
			return os.Lchown(p, uid, gid)
		})
	}, ""
}