llmgw/internal/slurm/shared.go

package slurm

import (
	"fmt"
	"regexp"
	"slices"
	"strings"

	"golang.org/x/crypto/ssh"
)

// Shared storage: one path every machine of the Slurm cluster sees, from
// the site's NFS export, from an llmgw node that exports it, or mounted by
// the site itself.
const (
	// NFSPort is NFSv4's one port, which is all a client needs.
	NFSPort = 2049
	// ExportsPath is the export llmgw keeps on the machine that serves.
	ExportsPath = "/etc/exports.d/llmgw-slurm.exports"
	// MountUnitDir is where the mount unit is written.
	MountUnitDir = "/etc/systemd/system"
)

// Exports is the exports file for path, read-write to clients alone. Root
// on a client is nobody here: a client's root cannot take another's files.
func Exports(path string, clients []string) string {
	cs := slices.Clone(clients)
	slices.Sort(cs)
	cs = slices.Compact(cs)
	var b strings.Builder
	b.WriteString("# Written by llmgw (BACKLOG P): the Slurm cluster's shared files, to its machines alone.\n")
	if len(cs) == 0 {
		return b.String()
	}
	fmt.Fprintf(&b, "%q", path)
	for _, c := range cs {
		fmt.Fprintf(&b, " %s(rw,sync,no_subtree_check,root_squash)", c)
	}
	b.WriteString("\n")
	return b.String()
}

// MountUnitName is systemd's name for the mount unit of path, as
// systemd-escape --path makes it.
func MountUnitName(path string) string {
	p := strings.Trim(path, "/")
	var b strings.Builder
	for i := 0; i < len(p); i++ {
		c := p[i]
		switch {
		case c == '/':
			b.WriteByte('-')
		case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9', c == '_', c == '.' && i > 0:
			b.WriteByte(c)
		default:
			fmt.Fprintf(&b, `\x%02x`, c)
		}
	}
	return b.String() + ".mount"
}

// MountUnit mounts the export nfs at path. hard: a job waits for a server
// that is gone rather than reading errors; nconnect spreads a client's
// traffic over several connections. The version is the highest both
// sides speak.
func MountUnit(nfs, path string) string {
	return "# Written by llmgw (BACKLOG P): the Slurm cluster's shared files.\n" +
		"[Unit]\nDescription=llmgw: the Slurm cluster's shared files\n" +
		"After=network-online.target\nWants=network-online.target\n\n" +
		"[Mount]\nWhat=" + nfs + "\nWhere=" + path + "\nType=nfs\n" +
		"Options=hard,nconnect=4,_netdev\nTimeoutSec=60\n\n" +
		"[Install]\nWantedBy=remote-fs.target\n"
}

// Users: Linux users llmgw makes from its accounts, the same login and UID
// on every machine, logging in with the SSH keys given in the portal.
const (
	// AuthorizedKeysDir holds each user's keys, root's to write, so a
	// user's own home (on a file system shared with others) need hold none.
	AuthorizedKeysDir = "/etc/llmgw/ssh-keys"
	// SSHDropIn has sshd look there too, after the user's own files.
	SSHDropIn = "/etc/ssh/sshd_config.d/50-llmgw-slurm.conf"
	// MinUID is the lowest UID llmgw gives.
	MinUID = 200000
)

// SSHConf is SSHDropIn's settings.
const SSHConf = "# Written by llmgw (BACKLOG P): the Slurm users' keys, from llmgw's portal.\n" +
	"AuthorizedKeysFile .ssh/authorized_keys .ssh/authorized_keys2 " + AuthorizedKeysDir + "/%u\n"

// login is what useradd takes, as llmgw gives it.
var login = regexp.MustCompile(`^[a-z][a-z0-9_-]{0,31}$`)

// ValidLogin is whether a login is one llmgw gives.
func ValidLogin(s string) bool { return login.MatchString(s) }

// ParseSSHKeys reads OpenSSH public keys, one a line, as authorized_keys
// takes them, without options: what a person pastes into the portal.
// Blank lines and # comments are skipped; each key comes back as its type,
// its base64 and its comment.
func ParseSSHKeys(text string) ([]string, error) {
	var out []string
	for i, line := range strings.Split(text, "\n") {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		pk, comment, opts, _, err := ssh.ParseAuthorizedKey([]byte(line))
		if err != nil {
			return nil, fmt.Errorf("line %d is not an SSH public key (as ~/.ssh/id_ed25519.pub has it)", i+1)
		}
		if len(opts) > 0 {
			return nil, fmt.Errorf("line %d has options before its key; only the key is taken", i+1)
		}
		k := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(pk)))
		if comment = strings.Join(strings.Fields(comment), " "); comment != "" {
			k += " " + comment
		}
		if !slices.Contains(out, k) {
			out = append(out, k)
		}
	}
	if len(out) > 20 {
		return nil, fmt.Errorf("%d keys: 20 at most", len(out))
	}
	return out, nil
}