llmgw/internal/slurm/pam.go

package slurm

import (
	"errors"
	"fmt"
	"path"
	"regexp"
	"strconv"
	"strings"
)

// Logging in to a node: llmgw's users reach the controller, the machine
// they submit from, whenever they like, and a node only while a job of
// theirs runs there, their session put in that job (pam_slurm_adopt).
// Without it a user could ssh to a node and use its cards outside Slurm,
// on a node llmgw serves models from too.
const (
	// PAMSSHD is sshd's PAM stack, where the lines go.
	PAMSSHD = "/etc/pam.d/sshd"
	// AdoptModule is the module's file name.
	AdoptModule = "pam_slurm_adopt.so"

	pamBegin   = "# llmgw (BACKLOG P): llmgw's Slurm users reach this node only while a job of theirs runs here."
	pamSession = "# llmgw (BACKLOG P): pam_systemd skipped for llmgw's Slurm users, so their session stays in the job pam_slurm_adopt put it in."
	pamWas     = "# llmgw, in place of: "
	pamEnd     = "# llmgw end"
)

// AdoptPackages are the packages with AdoptModule, by family.
var AdoptPackages = map[string][]string{"deb": {"libpam-slurm-adopt"}, "rpm": {"slurm-pam_slurm"}}

// ModuleDirs are where PAM's modules are, on Debian's and Red Hat's.
var ModuleDirs = []string{
	"/usr/lib/x86_64-linux-gnu/security", "/lib/x86_64-linux-gnu/security",
	"/usr/lib/aarch64-linux-gnu/security", "/lib/aarch64-linux-gnu/security",
	"/usr/lib64/security", "/lib64/security", "/usr/lib/security", "/lib/security",
}

// PAMAdopt is sshd's PAM stack with llmgw's lines in it (on) or taken out
// and the stack as it was. Only the users listed in the file list, one a
// line, are held to their jobs; anyone else, root and the site's
// administrators, logs in as before.
//
// The account lines go after the last account line, so a user the
// machine's own rules refuse is still refused. pam_systemd, which would
// move the session out of the job into a scope of its own, is skipped for
// the listed users: the file sshd includes it from (read through read) is
// put in sshd's stack in its place, with the one line in front of
// pam_systemd. contained is whether that was done; when it cannot be done
// safely, the users are still held to their jobs, their sessions not to
// the job's cards and cores.
func PAMAdopt(stack, list string, read func(name string) (string, error), on bool) (out string, contained bool, err error) {
	lines, err := pamStrip(stack)
	if err != nil {
		return "", false, err
	}
	if !on {
		return join(lines), false, nil
	}
	last := -1
	for i, l := range lines {
		f := strings.Fields(l)
		if len(f) > 0 && (f[0] == "account" || (f[0] == "@include" && len(f) > 1 && strings.Contains(f[1], "account"))) {
			last = i
		}
	}
	if last < 0 {
		return "", false, errors.New(PAMSSHD + " has no account lines to put llmgw's after")
	}
	block := []string{
		pamBegin,
		// Listed: on to the next line; not listed, or no list: past it.
		"account [success=ignore default=1] pam_listfile.so item=user sense=allow file=" + list + " onerr=fail",
		"account required " + AdoptModule,
		pamEnd,
	}
	lines = append(lines[:last+1:last+1], append(block, lines[last+1:]...)...)
	lines, contained = pamSkipSystemd(lines, list, read)
	return join(lines), contained, nil
}

func join(lines []string) string { return strings.Join(lines, "\n") + "\n" }

// pamStrip takes llmgw's blocks out, putting back the lines they stand in
// place of.
func pamStrip(stack string) ([]string, error) {
	var lines []string
	in := false
	for _, l := range strings.Split(strings.TrimSuffix(stack, "\n"), "\n") {
		switch {
		case !in && (l == pamBegin || l == pamSession):
			in = true
		case in && l == pamEnd:
			in = false
		case in && strings.HasPrefix(l, pamWas):
			lines = append(lines, strings.TrimPrefix(l, pamWas))
		case !in:
			lines = append(lines, l)
		}
	}
	if in {
		return nil, errors.New("llmgw's lines in " + PAMSSHD + " have no end; put them right by hand")
	}
	return lines, nil
}

// pamLine is a PAM line's type, control and module, the control with its
// brackets.
var pamLine = regexp.MustCompile(`^\s*(-?[a-z]+)\s+(\[[^\]]*\]|\S+)\s+(\S+)`)

// pamJump is a jump of N in a bracketed control.
var pamJump = regexp.MustCompile(`=(\d+)`)

func isSession(l string) bool {
	m := pamLine.FindStringSubmatch(l)
	return m != nil && strings.TrimPrefix(m[1], "-") == "session"
}

// pamSkipSystemd puts the guard in front of pam_systemd, in the session
// file sshd includes it from or in sshd's own lines; false, and lines as
// they were, when it is in neither or cannot be done without changing
// what any other jump in the stack lands on.
func pamSkipSystemd(lines []string, list string, read func(string) (string, error)) ([]string, bool) {
	guard := "session [success=1 default=ignore] pam_listfile.so item=user sense=allow file=" + list + " onerr=fail"
	hasSystemd := func(l string) bool {
		m := pamLine.FindStringSubmatch(l)
		return m != nil && strings.TrimPrefix(m[1], "-") == "session" && path.Base(m[3]) == "pam_systemd.so"
	}
	for i, l := range lines {
		f := strings.Fields(l)
		var body []string
		switch {
		case hasSystemd(l):
			body = []string{l}
		case len(f) == 2 && f[0] == "@include" && strings.Contains(f[1], "session"),
			len(f) == 3 && f[0] == "session" && f[1] == "include":
			name := f[len(f)-1]
			if strings.Contains(name, "/") {
				continue
			}
			text, err := read(name)
			if err != nil {
				return lines, false
			}
			for _, il := range strings.Split(text, "\n") {
				t := strings.TrimSpace(il)
				if t == "" || strings.HasPrefix(t, "#") {
					continue
				}
				if f[0] == "session" && !isSession(t) {
					continue // include takes only the lines of its own type
				}
				body = append(body, t)
			}
		default:
			continue
		}
		at := -1
		for j, b := range body {
			if hasSystemd(b) {
				at = j
			}
			if g := strings.Fields(b); len(g) > 0 && (g[0] == "@include" || strings.Contains(b, "include") || strings.Contains(b, "substack")) {
				return lines, false // nested: not followed
			}
		}
		if at < 0 {
			continue
		}
		// A jump in sshd's own session lines may cross what is put in.
		for j, o := range lines {
			if j != i && isSession(o) && pamJump.MatchString(pamControl(o)) {
				return lines, false
			}
		}
		body, ok := pamInsert(body, at, guard)
		if !ok {
			return lines, false
		}
		block := append([]string{pamSession, pamWas + l}, body...)
		block = append(block, pamEnd)
		return append(lines[:i:i], append(block, lines[i+1:]...)...), true
	}
	return lines, false
}

func pamControl(l string) string {
	if m := pamLine.FindStringSubmatch(l); m != nil {
		return m[2]
	}
	return ""
}

// pamInsert puts guard in front of body[at], a jump that crosses it made
// one longer so it lands where it did.
func pamInsert(body []string, at int, guard string) ([]string, bool) {
	// Each session line's place among the session lines.
	idx := make([]int, len(body))
	n := 0
	for j, b := range body {
		idx[j] = n
		if isSession(b) {
			n++
		}
	}
	out := make([]string, 0, len(body)+1)
	for j, b := range body {
		if j == at {
			out = append(out, guard)
		}
		ctl := pamControl(b)
		if isSession(b) && j < at && strings.HasPrefix(ctl, "[") {
			from, target := idx[j], idx[at]
			fixed := pamJump.ReplaceAllStringFunc(ctl, func(s string) string {
				k, err := strconv.Atoi(s[1:])
				if err != nil {
					return s
				}
				if from < target && target <= from+k {
					k++
				}
				return fmt.Sprintf("=%d", k)
			})
			b = strings.Replace(b, ctl, fixed, 1)
		}
		out = append(out, b)
	}
	return out, true
}