llmgw/internal/onboard/sshconfig.go

package onboard

import (
	"bufio"
	"bytes"
	"context"
	"net"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"
	"time"
)

// Where ssh reads its configuration, and arguments every `ssh -G` here is
// given; variables so a test can point them at a file of its own.
var (
	userSSHConfig   = func() string { return filepath.Join(sshHome(), ".ssh", "config") }
	systemSSHConfig = "/etc/ssh/ssh_config"
	sshExtraArgs    []string
)

// sshHome is the home ssh itself reads ~/.ssh from: the account's, from the
// password database, not $HOME, which sudo may leave as the caller's.
func sshHome() string {
	if u, err := user.Current(); err == nil && u.HomeDir != "" {
		return u.HomeDir
	}
	h, _ := os.UserHomeDir()
	return h
}

// configAlias is the ssh config's name for addr: the first concrete Host
// whose HostName is addr, or resolves to the same address. A machine is
// usually written in the config under its name, `Host gpu01` with
// `HostName 10.0.0.21` and the User, IdentityFile or ProxyJump it needs,
// and given to node add by its address, which that block does not match.
// Empty when addr is itself named on a Host line, or no block points at it.
func configAlias(ctx context.Context, addr string) string {
	aliases := map[string]bool{}
	var order []string
	seen := map[string]bool{}
	collect := func(path, base string) {
		for _, h := range hostNames(path, base, seen, 0) {
			if !aliases[h] {
				aliases[h] = true
				order = append(order, h)
			}
		}
	}
	collect(userSSHConfig(), filepath.Join(sshHome(), ".ssh"))
	if systemSSHConfig != "" {
		collect(systemSSHConfig, filepath.Dir(systemSSHConfig))
	}
	for h := range aliases {
		if strings.EqualFold(h, addr) {
			return ""
		}
	}
	if len(order) > 256 {
		order = order[:256]
	}
	// Names are looked up only when no HostName is the address as written,
	// and within one budget: a config of many hosts must not hold up a
	// connection by seconds each.
	var pending []string
	hostOf := map[string]string{}
	for _, a := range order {
		h := sshGet(ctx, a)["hostname"]
		if strings.EqualFold(h, addr) {
			return a
		}
		if h != "" {
			hostOf[a] = h
			pending = append(pending, a)
		}
	}
	if len(pending) == 0 {
		return ""
	}
	lctx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()
	want := lookup(lctx, addr)
	if len(want) == 0 {
		return ""
	}
	for _, a := range pending {
		for ip := range lookup(lctx, hostOf[a]) {
			if want[ip] {
				return a
			}
		}
	}
	return ""
}

// configUser is the user the ssh config logs in to host as, empty when it
// names none. ssh -G reports the local account when nothing set one, so
// that is taken as none: the config saying `User root` to a root caller is
// the one case read wrongly, and it gets the fallback user instead.
func configUser(ctx context.Context, host string) string {
	u := sshGet(ctx, host)["user"]
	if cur, err := user.Current(); err == nil && u == cur.Username {
		return ""
	}
	return u
}

// sshGet is ssh's own reading of its configuration for host, keys
// lowercased; empty when ssh cannot say.
func sshGet(ctx context.Context, host string) map[string]string {
	args := append(append([]string{"-G"}, sshExtraArgs...), "--", host)
	out, err := exec.CommandContext(ctx, "ssh", args...).Output()
	if err != nil {
		return nil
	}
	m := map[string]string{}
	sc := bufio.NewScanner(bytes.NewReader(out))
	for sc.Scan() {
		k, v, ok := strings.Cut(sc.Text(), " ")
		if k = strings.ToLower(k); ok {
			if _, dup := m[k]; !dup {
				m[k] = v
			}
		}
	}
	return m
}

func lookup(ctx context.Context, host string) map[string]bool {
	if ip := net.ParseIP(host); ip != nil {
		return map[string]bool{ip.String(): true}
	}
	addrs, err := net.DefaultResolver.LookupHost(ctx, host)
	if err != nil {
		return nil
	}
	m := map[string]bool{}
	for _, a := range addrs {
		if ip := net.ParseIP(a); ip != nil {
			m[ip.String()] = true
		}
	}
	return m
}

// hostNames are the concrete names on the Host lines of an ssh config file
// and the files it includes — no patterns, no negations. base is where a
// relative Include is looked for, as ssh does: ~/.ssh for the user's file,
// /etc/ssh for the system's.
func hostNames(path, base string, seen map[string]bool, depth int) []string {
	if depth > 16 || seen[path] {
		return nil
	}
	seen[path] = true
	data, err := os.ReadFile(path)
	if err != nil {
		return nil
	}
	var names []string
	for _, line := range strings.Split(string(data), "\n") {
		key, args := configLine(line)
		switch key {
		case "host":
			for _, a := range args {
				if a != "" && !strings.ContainsAny(a, "*?!") {
					names = append(names, a)
				}
			}
		case "include":
			for _, a := range args {
				if strings.HasPrefix(a, "~/") {
					a = filepath.Join(sshHome(), a[2:])
				} else if !filepath.IsAbs(a) {
					a = filepath.Join(base, a)
				}
				matches, _ := filepath.Glob(a)
				for _, m := range matches {
					names = append(names, hostNames(m, base, seen, depth+1)...)
				}
			}
		}
	}
	return names
}

// configLine splits one line of ssh config into its keyword, lowercased,
// and arguments: "Host a b", "Host=a b" and quoted arguments alike.
func configLine(line string) (string, []string) {
	line = strings.TrimSpace(line)
	if line == "" || strings.HasPrefix(line, "#") {
		return "", nil
	}
	i := strings.IndexAny(line, " \t=")
	if i < 0 {
		return strings.ToLower(line), nil
	}
	key := strings.ToLower(line[:i])
	rest := strings.TrimLeft(line[i:], " \t")
	rest = strings.TrimPrefix(rest, "=")
	var args []string
	for rest = strings.TrimSpace(rest); rest != ""; rest = strings.TrimSpace(rest) {
		if rest[0] == '"' {
			end := strings.IndexByte(rest[1:], '"')
			if end < 0 {
				args = append(args, rest[1:])
				break
			}
			args = append(args, rest[1:1+end])
			rest = rest[2+end:]
			continue
		}
		j := strings.IndexAny(rest, " \t")
		if j < 0 {
			args = append(args, rest)
			break
		}
		args = append(args, rest[:j])
		rest = rest[j:]
	}
	return key, args
}