From bc42f5583c7e852b948e1a7ea420c55d980a884c Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Fri, 28 Aug 2026 14:07:52 +0000 Subject: [PATCH] [shim] Don't revoke host SSH access of co-located tasks The shim appends `cfg.HostSshKeys` to the host user's `authorized_keys` when a task starts and removes every entry matching those fingerprints when it ends. That key is the user's own SSH key, shared by all of their runs unless overridden with `dstack apply --ssh-identity`, so on an instance that holds more than one of their jobs at a time -- a fleet with `blocks` -- the first job to finish revokes the SSH access of the jobs still running. The same unconditional removal also drops a key the user added by hand, if a run happened to reuse it. The per-task add/remove deltas are replaced with a reconcile of the entries the shim owns, that is, the ones carrying the `# added by dstack-shim` marker shipped in #4206. For every host user, those entries become exactly the keys of the stored tasks that are not cleaned up yet; everything else in the file -- keys added by the user, keys added by the server when provisioning an SSH fleet, comments -- is kept verbatim. Refcounting is implicit in set membership, so there is no new state to keep or persist: `HostSshKeys` and `HostSshUser` are already a part of the task config persisted since #4220. * Honoring the marker is safe as of 0.21.3, the first release carrying both the marker and the task state file, so no shim version writes one without the other. A task restored without a state file has no config, and therefore no keys in the reconciled set, but it was started by an earlier shim, whose entries are unmarked and never touched. * `cleanupLocked()` commits `cleanedUp` before releasing the resources, so that the finishing task is already out of the reconciled set by the time its keys are reconciled. If the shim stops running in between, the state file, written last, still reports the task as not cleaned up, and the cleanup runs again, idempotently, after the restart. * The file is written to a temporary file in the same dir and renamed over the old one, preserving its mode and setting its owner, so that a failed write cannot leave the user without their keys. `~/.ssh` and the file itself are created if missing, as the server does when provisioning an SSH fleet, instead of failing on a missing file. The `.bak` copy is dropped: the rename is atomic and the entries the shim does not own are copied through untouched, so it protected nothing, while leaving a root-owned copy of the user's keys in `~/.ssh`. * A single reconcile on start drops the entries of the tasks that are gone and collapses the duplicates left by the earlier versions, which appended a key once per task using it. Fixes: https://github.com/dstackai/dstack/issues/4174 Co-authored-by: Claude Opus 5 (1M context) --- runner/internal/shim/authorized_keys.go | 285 ++++++++------- runner/internal/shim/authorized_keys_test.go | 366 ++++++++++--------- runner/internal/shim/docker.go | 89 ++++- runner/internal/shim/docker_test.go | 117 ++++++ runner/internal/shim/process_test.go | 10 + runner/internal/shim/state.go | 2 +- 6 files changed, 561 insertions(+), 308 deletions(-) diff --git a/runner/internal/shim/authorized_keys.go b/runner/internal/shim/authorized_keys.go index a34fcb142..5bed6b475 100644 --- a/runner/internal/shim/authorized_keys.go +++ b/runner/internal/shim/authorized_keys.go @@ -6,11 +6,10 @@ import ( "context" "errors" "fmt" - "io" "os" "os/user" "path/filepath" - "slices" + "strconv" "strings" "golang.org/x/crypto/ssh" @@ -23,16 +22,24 @@ import ( // provisioning SSH fleets. sshd ignores everything after the key blob, so the marker // has no effect on authentication; it only records that the entry is ours. // -// Nothing honors the marker yet -- keys are still removed by fingerprint regardless of -// the comment. It is written now so that, once removal does honor it, entries added by -// earlier shim versions are already marked. Until then, a task can only be started and -// finalized by the same shim process, so there is no cross-version handoff to protect. +// The marker is what makes Reconcile() safe: only the entries carrying it are rewritten, +// so a key added by the user, or by the server at fleet provisioning time, is never +// touched. Entries added by a shim version that did not write the marker are kept +// forever; leaking a key we cannot prove we added beats revoking it. // // The marker must be matched as an exact suffix: a substring match on `# added by -// dstack` would also claim the keys the server adds at fleet provisioning time, which -// the shim must never remove. +// dstack` would also claim the keys the server adds at fleet provisioning time. const publicKeyMarker = "# added by dstack-shim" +const ( + sshDirName = ".ssh" + authorizedKeysFileName = "authorized_keys" + // Modes of the dir and the file when the shim creates them; an existing + // authorized_keys file keeps its own mode + sshDirMode = 0o700 + authorizedKeysFileMode = 0o600 +) + func PublicKeyFingerprint(key string) (string, error) { pk, _, _, _, err := ssh.ParseAuthorizedKey([]byte(key)) if err != nil { @@ -42,37 +49,6 @@ func PublicKeyFingerprint(key string) (string, error) { return keyFingerprint, nil } -func IsPublicKeysEqual(left string, right string) bool { - leftFingerprint, err := PublicKeyFingerprint(left) - if err != nil { - return false - } - - rightFingerprint, err := PublicKeyFingerprint(right) - if err != nil { - return false - } - - return leftFingerprint == rightFingerprint -} - -func RemovePublicKeys(fileKeys []string, keysToRemove []string) []string { - newKeys := slices.DeleteFunc(fileKeys, func(fileKey string) bool { - delete := slices.ContainsFunc(keysToRemove, func(removeKey string) bool { - return IsPublicKeysEqual(fileKey, removeKey) - }) - return delete - }) - return newKeys -} - -func AppendPublicKeys(fileKeys []string, keysToAppend []string) []string { - newKeys := []string{} - newKeys = append(newKeys, fileKeys...) - newKeys = append(newKeys, keysToAppend...) - return newKeys -} - // canonicalizePublicKey validates a public key received from the server and returns the // authorized_keys line to write for it, marked with publicKeyMarker. // @@ -100,17 +76,59 @@ func canonicalizePublicKey(publicKey string) (string, error) { return keyLine + " " + strings.Join(commentFields, " "), nil } +// isShimEntry reports whether the authorized_keys line has been added by the shim, that +// is, whether it carries publicKeyMarker. canonicalizePublicKey() always puts the marker +// last, therefore matching it as a suffix cannot claim an entry that merely mentions it. +func isShimEntry(line string) bool { + return strings.HasSuffix(strings.TrimRight(line, " \t\r"), " "+publicKeyMarker) +} + type AuthorizedKeys struct { user string lookup func(username string) (*user.User, error) } -// AppendPublicKeys appends the keys to the user's authorized_keys file, marking them -// with publicKeyMarker. Invalid entries are skipped, so that one bad key does not keep -// the rest out of the file. -// Duplicates are not detected: a key already present in the file is appended once more. -func (ak AuthorizedKeys) AppendPublicKeys(ctx context.Context, publicKeys []string) error { +// Reconcile makes the shim-owned part of the user's authorized_keys file exactly the +// given set of keys: the entries carrying publicKeyMarker are replaced with the entries +// for publicKeys. Everything else -- keys added by the user, keys added by the server +// when provisioning an SSH fleet, comments, blank lines -- is kept verbatim and in order. +// +// The keys of all the tasks that still need them are passed on every call, so that a task +// releasing a key shared with another task does not revoke it, see dstackai/dstack#4174. +// Duplicates collapse into a single entry, which is what makes the file depend on the set +// of keys in use and not on the number of tasks using them. +// A key that is also present as an entry the shim does not own is still written as an +// entry of its own: the two have different owners, and sharing one would let a manual +// edit revoke the access of a running task. +// +// Invalid keys are skipped, so that one bad key does not keep the rest out of the file. +func (ak AuthorizedKeys) Reconcile(ctx context.Context, publicKeys []string) error { + usr, err := ak.lookup(ak.user) + if err != nil { + return fmt.Errorf("lookup user %s: %w", ak.user, err) + } + path := authorizedKeysPath(usr.HomeDir) + + lines, mode, err := readAuthorizedKeys(path) + if err != nil { + return err + } + kept := make([]string, 0, len(lines)+len(publicKeys)) + for _, line := range lines { + if !isShimEntry(line) { + kept = append(kept, line) + } + } + kept = append(kept, shimEntries(ctx, publicKeys)...) + + return writeAuthorizedKeys(path, kept, mode, usr) +} + +// shimEntries returns the marked authorized_keys lines to write for the keys, skipping +// the invalid ones and collapsing the duplicates +func shimEntries(ctx context.Context, publicKeys []string) []string { lines := make([]string, 0, len(publicKeys)) + seen := make(map[string]struct{}, len(publicKeys)) for _, publicKey := range publicKeys { line, err := canonicalizePublicKey(publicKey) if err != nil { @@ -121,44 +139,117 @@ func (ak AuthorizedKeys) AppendPublicKeys(ctx context.Context, publicKeys []stri ) continue } + // Matched by fingerprint and not by line, so that the same key submitted with + // different comments does not yield two entries + fingerprint, err := PublicKeyFingerprint(line) + if err != nil { + // canonicalizePublicKey() has already parsed the key, so this cannot happen + log.Error(ctx, "failed to fingerprint canonicalized key", "err", err) + continue + } + if _, ok := seen[fingerprint]; ok { + continue + } + seen[fingerprint] = struct{}{} lines = append(lines, line) } - if len(lines) == 0 { - return nil - } - return ak.transformAuthorizedKeys(AppendPublicKeys, lines) + return lines } -// RemovePublicKeys removes the keys from the user's authorized_keys file, matching by -// fingerprint and ignoring the comment. That is, publicKeyMarker is not honored yet, so -// entries the shim did not add are removed as well, and every matching entry is removed -// even if another task still relies on the key. See dstackai/dstack#4174. -func (ak AuthorizedKeys) RemovePublicKeys(publicKeys []string) error { - return ak.transformAuthorizedKeys(RemovePublicKeys, publicKeys) -} +// readAuthorizedKeys returns the lines of the file and its mode. A missing file is +// reported as an empty one with the mode to create it with: sshd treats it the same way, +// and the shim creates it as the server does when provisioning an SSH fleet. +func readAuthorizedKeys(path string) ([]string, os.FileMode, error) { + file, err := os.Open(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil, authorizedKeysFileMode, nil + } + return nil, 0, fmt.Errorf("open authorized keys: %w", err) + } + defer file.Close() -func (ak AuthorizedKeys) read(r io.Reader) ([]string, error) { - lines := []string{} - scanner := bufio.NewScanner(r) + info, err := file.Stat() + if err != nil { + return nil, 0, fmt.Errorf("stat authorized keys: %w", err) + } + var lines []string + scanner := bufio.NewScanner(file) for scanner.Scan() { - text := scanner.Text() - lines = append(lines, text) + lines = append(lines, scanner.Text()) } if err := scanner.Err(); err != nil { - return []string{}, fmt.Errorf("scan authorized keys: %w", err) + return nil, 0, fmt.Errorf("scan authorized keys: %w", err) } - return lines, nil + return lines, info.Mode().Perm(), nil } -func (ak AuthorizedKeys) write(w io.Writer, lines []string) error { - wr := bufio.NewWriter(w) - for _, line := range lines { - _, err := fmt.Fprintln(wr, line) +// writeAuthorizedKeys replaces the file with the given lines, creating it and its dir if +// they do not exist. The content is written to a temporary file in the same dir and +// renamed over the old one, so that a failed or interrupted write cannot leave the user +// with a partial file, that is, without some of their keys. +func writeAuthorizedKeys(path string, lines []string, mode os.FileMode, usr *user.User) (err error) { + uid, gid, err := userIDs(usr) + if err != nil { + return err + } + dir := filepath.Dir(path) + if _, statErr := os.Stat(dir); errors.Is(statErr, os.ErrNotExist) { + if err := os.MkdirAll(dir, sshDirMode); err != nil { + return fmt.Errorf("create %s: %w", dir, err) + } + // The shim normally runs as root, therefore a dir it creates must be given + // to the user. An existing dir is left alone, as it is not ours to fix + if err := os.Chown(dir, uid, gid); err != nil { + return fmt.Errorf("chown %s: %w", dir, err) + } + } + + var content []byte + if len(lines) > 0 { + // The last line is terminated as well, so that appending to the file by hand + // cannot join a new entry to the last one + content = []byte(strings.Join(lines, "\n") + "\n") + } + tempPath := path + ".tmp" + defer func() { if err != nil { - return fmt.Errorf("write line: %w", err) + // sshd does not read the temporary file, so leaving it behind would only + // clutter the dir with a half-written copy of the user's keys + if removeErr := os.Remove(tempPath); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) { + err = errors.Join(err, fmt.Errorf("remove %s: %w", tempPath, removeErr)) + } } + }() + if err := writeFileSync(tempPath, content, mode); err != nil { + return fmt.Errorf("write authorized keys: %w", err) + } + // The mode passed to writeFileSync() only applies to a file it creates, and is + // masked by the umask, therefore it is set explicitly + if err := os.Chmod(tempPath, mode); err != nil { + return fmt.Errorf("chmod %s: %w", tempPath, err) + } + if err := os.Chown(tempPath, uid, gid); err != nil { + return fmt.Errorf("chown %s: %w", tempPath, err) + } + if err := os.Rename(tempPath, path); err != nil { + return fmt.Errorf("rename %s: %w", tempPath, err) + } + return nil +} + +// userIDs returns the numeric ids of the user. os/user reports them as strings, since +// not all platforms have numeric ids, but Linux, the only platform the shim runs on, does +func userIDs(usr *user.User) (int, int, error) { + uid, err := strconv.Atoi(usr.Uid) + if err != nil { + return 0, 0, fmt.Errorf("parse uid %q of user %s: %w", usr.Uid, usr.Username, err) + } + gid, err := strconv.Atoi(usr.Gid) + if err != nil { + return 0, 0, fmt.Errorf("parse gid %q of user %s: %w", usr.Gid, usr.Username, err) } - return wr.Flush() + return uid, gid, nil } func (ak AuthorizedKeys) GetHomeDirectory() (string, error) { @@ -174,61 +265,9 @@ func (ak AuthorizedKeys) GetAuthorizedKeysPath() (string, error) { if err != nil { return "", err } - return filepath.Join(homeDir, ".ssh", "authorized_keys"), nil + return authorizedKeysPath(homeDir), nil } -func (ak AuthorizedKeys) transformAuthorizedKeys(transform func([]string, []string) []string, publicKeys []string) error { - authorizedKeysPath, err := ak.GetAuthorizedKeysPath() - if err != nil { - return fmt.Errorf("get authorized keys path: %w", err) - } - - info, err := os.Stat(authorizedKeysPath) - if err != nil { - return fmt.Errorf("stat authorized keys: %w", err) - } - fileMode := info.Mode().Perm() - - authorizedKeysFile, err := os.OpenFile(authorizedKeysPath, os.O_RDWR, fileMode) - if err != nil { - return fmt.Errorf("open authorized keys: %w", err) - } - defer authorizedKeysFile.Close() - - lines, err := ak.read(authorizedKeysFile) - if err != nil { - return fmt.Errorf("read authorized keys: %w", err) - } - - // write backup - authorizedKeysPath, err = ak.GetAuthorizedKeysPath() - if err != nil { - return fmt.Errorf("get authorized keys path: %w", err) - } - - authorizedKeysPathBackup := authorizedKeysPath + ".bak" - authorizedKeysBackup, err := os.OpenFile(authorizedKeysPathBackup, os.O_RDWR|os.O_CREATE|os.O_TRUNC, fileMode) - if err != nil { - return fmt.Errorf("open authorized keys backup: %w", err) - } - defer authorizedKeysBackup.Close() - if err := ak.write(authorizedKeysBackup, lines); err != nil { - return fmt.Errorf("write authorized keys backup: %w", err) - } - - // transform lines - newLines := transform(lines, publicKeys) - - // write authorized_keys - if err := authorizedKeysFile.Truncate(0); err != nil { - return fmt.Errorf("truncate authorized keys: %w", err) - } - if _, err := authorizedKeysFile.Seek(0, 0); err != nil { - return fmt.Errorf("seek authorized keys: %w", err) - } - if err := ak.write(authorizedKeysFile, newLines); err != nil { - return fmt.Errorf("write authorized keys: %w", err) - } - - return nil +func authorizedKeysPath(homeDir string) string { + return filepath.Join(homeDir, sshDirName, authorizedKeysFileName) } diff --git a/runner/internal/shim/authorized_keys_test.go b/runner/internal/shim/authorized_keys_test.go index c2150f184..e2ab99b32 100644 --- a/runner/internal/shim/authorized_keys_test.go +++ b/runner/internal/shim/authorized_keys_test.go @@ -1,16 +1,23 @@ package shim import ( - "context" "fmt" "os" "os/user" - "path" + "path/filepath" + "strconv" "testing" "github.com/stretchr/testify/require" ) +// Two valid keys, used to build authorized_keys entries. Short on purpose: the blob is +// noise in a test that is about the lines around it +const ( + testKeyOne = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGYzO2yHhoIzYHnGH5CT/hpTNGRHvJHkKQlXqPZ0Uxwj" + testKeyTwo = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAILuLmyPGV/gcatBaZFxRPKGQVJ4vBjuqEsHIkKGrGZKS" +) + func TestPublicKeyFingerprint(t *testing.T) { key := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" expectedFingerprint := "SHA256:9HymzYAtJKNh8gKufl3EVoRSauL4E7Mbmuzqlcvii50" @@ -26,75 +33,6 @@ func TestPublicKeyFingerprintError(t *testing.T) { require.Empty(t, fingerprint) } -func TestIsPublicKeysEqual(t *testing.T) { - keyLeft := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - keyRight := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - - result := IsPublicKeysEqual(keyLeft, keyRight) - require.True(t, result) -} - -func TestIsPublicKeysEqualBrokenKey(t *testing.T) { - keyLeft := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - keyRight := "ssh-rsa AAAAP66um5MadfhB5dSnEM= thebits@barracuda" - - resultFwd := IsPublicKeysEqual(keyLeft, keyRight) - require.False(t, resultFwd) - - resultBck := IsPublicKeysEqual(keyRight, keyLeft) - require.False(t, resultBck) -} - -func TestIsPublicKeysNotEqual(t *testing.T) { - keyLeft := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - keyRight := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCfAFHwfyMKPFbKq+D/vYNaXjqer4uV5+zvlrPY2bvkdRT4GiH4hm2s1Z7+fUEYQBNfw5O9SgxGotqyguUJbuVUc2BCNdD8HC3PxKtEev35ga4G3jjyuVeHcL2T9pn+F8IW1o3SpDGATAHJyFtArPYz31Hwg6PiuggPNdPLMSzZNrwNVuPwT1uDMKFqAh+1ryIVi7389fjZ7aBR9F06VIPpWIVVKqSVD+NbHtwWqCw8AsprJE3bPwVW09OJeQX8GXryKasaX4t4HMXmO/UI8tprnyf05dAl7NQOPY9Iut5PgfzEVY/T0M1RSnZi7i+1x7WBWX3aMM/Hv+NUeX2YtuAN" - - result := IsPublicKeysEqual(keyLeft, keyRight) - require.False(t, result) -} - -func TestRemovePublicKeys(t *testing.T) { - keyLeft := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - keyRight := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCfAFHwfyMKPFbKq+D/vYNaXjqer4uV5+zvlrPY2bvkdRT4GiH4hm2s1Z7+fUEYQBNfw5O9SgxGotqyguUJbuVUc2BCNdD8HC3PxKtEev35ga4G3jjyuVeHcL2T9pn+F8IW1o3SpDGATAHJyFtArPYz31Hwg6PiuggPNdPLMSzZNrwNVuPwT1uDMKFqAh+1ryIVi7389fjZ7aBR9F06VIPpWIVVKqSVD+NbHtwWqCw8AsprJE3bPwVW09OJeQX8GXryKasaX4t4HMXmO/UI8tprnyf05dAl7NQOPY9Iut5PgfzEVY/T0M1RSnZi7i+1x7WBWX3aMM/Hv+NUeX2YtuAN" - - keys := []string{keyLeft, keyRight} - newKeys := RemovePublicKeys(keys, []string{keyRight}) - - require.Len(t, newKeys, 1) - require.Equal(t, newKeys, []string{keyLeft}) -} - -func TestRemovePublicKeysRemoveAll(t *testing.T) { - keyLeft := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - keyRight := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCfAFHwfyMKPFbKq+D/vYNaXjqer4uV5+zvlrPY2bvkdRT4GiH4hm2s1Z7+fUEYQBNfw5O9SgxGotqyguUJbuVUc2BCNdD8HC3PxKtEev35ga4G3jjyuVeHcL2T9pn+F8IW1o3SpDGATAHJyFtArPYz31Hwg6PiuggPNdPLMSzZNrwNVuPwT1uDMKFqAh+1ryIVi7389fjZ7aBR9F06VIPpWIVVKqSVD+NbHtwWqCw8AsprJE3bPwVW09OJeQX8GXryKasaX4t4HMXmO/UI8tprnyf05dAl7NQOPY9Iut5PgfzEVY/T0M1RSnZi7i+1x7WBWX3aMM/Hv+NUeX2YtuAN" - - keys := []string{keyLeft, keyRight} - newKeys := RemovePublicKeys(keys, []string{keyRight, keyLeft}) - - require.Empty(t, newKeys) -} - -func TestRemovePublicKeysRemoveNotContained(t *testing.T) { - keyLeft := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - keyRight := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCfAFHwfyMKPFbKq+D/vYNaXjqer4uV5+zvlrPY2bvkdRT4GiH4hm2s1Z7+fUEYQBNfw5O9SgxGotqyguUJbuVUc2BCNdD8HC3PxKtEev35ga4G3jjyuVeHcL2T9pn+F8IW1o3SpDGATAHJyFtArPYz31Hwg6PiuggPNdPLMSzZNrwNVuPwT1uDMKFqAh+1ryIVi7389fjZ7aBR9F06VIPpWIVVKqSVD+NbHtwWqCw8AsprJE3bPwVW09OJeQX8GXryKasaX4t4HMXmO/UI8tprnyf05dAl7NQOPY9Iut5PgfzEVY/T0M1RSnZi7i+1x7WBWX3aMM/Hv+NUeX2YtuAN" - - keys := []string{keyLeft, keyRight} - newKeys := RemovePublicKeys(keys, []string{"# line with comment"}) - - require.Equal(t, keys, newKeys) -} - -func TestAppendPublicKeys(t *testing.T) { - keyLeft := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - keyRight := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCfAFHwfyMKPFbKq+D/vYNaXjqer4uV5+zvlrPY2bvkdRT4GiH4hm2s1Z7+fUEYQBNfw5O9SgxGotqyguUJbuVUc2BCNdD8HC3PxKtEev35ga4G3jjyuVeHcL2T9pn+F8IW1o3SpDGATAHJyFtArPYz31Hwg6PiuggPNdPLMSzZNrwNVuPwT1uDMKFqAh+1ryIVi7389fjZ7aBR9F06VIPpWIVVKqSVD+NbHtwWqCw8AsprJE3bPwVW09OJeQX8GXryKasaX4t4HMXmO/UI8tprnyf05dAl7NQOPY9Iut5PgfzEVY/T0M1RSnZi7i+1x7WBWX3aMM/Hv+NUeX2YtuAN" - comment := "# line with coment" - - keys := []string{keyLeft, keyRight} - newKeys := AppendPublicKeys(keys, []string{comment}) - - require.Equal(t, []string{keyLeft, keyRight, comment}, newKeys) -} - func mockUserLookup(username string) (*user.User, error) { if username == "test_user" { return &user.User{ @@ -143,6 +81,7 @@ func TestGetAuthorizedKeysPath(t *testing.T) { filePath, err := ak.GetAuthorizedKeysPath() if tc.isError { require.Error(t, err) + require.Equal(t, tc.expected, filePath) } else { require.NoError(t, err) require.Equal(t, tc.expected, filePath) @@ -150,53 +89,8 @@ func TestGetAuthorizedKeysPath(t *testing.T) { } } -func TestAppendKey(t *testing.T) { - ak := AuthorizedKeys{user: "test_user", lookup: mockUserLookup} - filePath, err := ak.GetAuthorizedKeysPath() - require.NoError(t, err) - - err = os.MkdirAll(path.Dir(filePath), os.ModePerm) - require.NoError(t, err) - - key := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - err = os.WriteFile(filePath, []byte(key), os.ModePerm) - require.NoError(t, err) - - newKey := "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGYzO2yHhoIzYHnGH5CT/hpTNGRHvJHkKQlXqPZ0Uxwj user@host" - err = ak.AppendPublicKeys(context.Background(), []string{newKey}) - require.NoError(t, err) - - b, err := os.ReadFile(filePath) - require.NoError(t, err) - require.Contains(t, string(b), key) - require.Contains(t, string(b), newKey+" # added by dstack-shim") -} - -func TestAppendKeySkipsInvalid(t *testing.T) { - ak := AuthorizedKeys{user: "test_user", lookup: mockUserLookup} - filePath, err := ak.GetAuthorizedKeysPath() - require.NoError(t, err) - - err = os.MkdirAll(path.Dir(filePath), os.ModePerm) - require.NoError(t, err) - err = os.WriteFile(filePath, []byte{}, os.ModePerm) - require.NoError(t, err) - - first := "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGYzO2yHhoIzYHnGH5CT/hpTNGRHvJHkKQlXqPZ0Uxwj first" - second := "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAILuLmyPGV/gcatBaZFxRPKGQVJ4vBjuqEsHIkKGrGZKS second" - - // authorized_keys is line-based, therefore an entry must hold exactly one key - err = ak.AppendPublicKeys(context.Background(), []string{first + "\n" + second, first}) - require.NoError(t, err) - - b, err := os.ReadFile(filePath) - require.NoError(t, err) - require.NotContains(t, string(b), second) - require.Equal(t, first+" # added by dstack-shim\n", string(b)) -} - func TestCanonicalizePublicKey(t *testing.T) { - const blob = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGYzO2yHhoIzYHnGH5CT/hpTNGRHvJHkKQlXqPZ0Uxwj" + const blob = testKeyOne testCases := []struct { name string @@ -261,76 +155,208 @@ func TestCanonicalizePublicKey(t *testing.T) { } } -func TestRemoveKey(t *testing.T) { - ak := AuthorizedKeys{user: "test_user", lookup: mockUserLookup} - filePath, err := ak.GetAuthorizedKeysPath() - require.NoError(t, err) +func TestIsShimEntry(t *testing.T) { + testCases := []struct { + name string + line string + expected bool + }{ + { + name: "shim entry", + line: testKeyOne + " user@host # added by dstack-shim", + expected: true, + }, + { + name: "shim entry without comment", + line: testKeyOne + " # added by dstack-shim", + expected: true, + }, + { + name: "trailing whitespace", + line: testKeyOne + " # added by dstack-shim \r", + expected: true, + }, + { + // the marker the server writes when provisioning an SSH fleet is a prefix of + // ours, and the keys it marks must never be touched by the shim + name: "server entry", + line: testKeyOne + " user@host # added by dstack", + expected: false, + }, + { + name: "user entry", + line: testKeyOne + " user@host", + expected: false, + }, + { + name: "marker not at the end", + line: testKeyOne + " # added by dstack-shim user@host", + expected: false, + }, + { + name: "blank line", + line: "", + expected: false, + }, + } - err = os.MkdirAll(path.Dir(filePath), os.ModePerm) - require.NoError(t, err) + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + require.Equal(t, tc.expected, isShimEntry(tc.line)) + }) + } +} - key := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - err = os.WriteFile(filePath, []byte(key), os.ModePerm) - require.NoError(t, err) +// newTestAuthorizedKeys returns an AuthorizedKeys for a user whose home is a temp dir, +// along with the path of their authorized_keys file. The mocked lookup reports the ids of +// the current process, so that setting the ownership of the file succeeds without root +func newTestAuthorizedKeys(t *testing.T) (AuthorizedKeys, string) { + t.Helper() + usr := &user.User{ + Username: "test_user", + HomeDir: t.TempDir(), + Uid: strconv.Itoa(os.Getuid()), + Gid: strconv.Itoa(os.Getgid()), + } + ak := AuthorizedKeys{ + user: usr.Username, + lookup: func(username string) (*user.User, error) { + if username != usr.Username { + return nil, fmt.Errorf("user not found") + } + return usr, nil + }, + } + return ak, authorizedKeysPath(usr.HomeDir) +} - err = ak.RemovePublicKeys([]string{key}) - require.NoError(t, err) +func writeTestAuthorizedKeys(t *testing.T, path string, content string) { + t.Helper() + require.NoError(t, os.MkdirAll(filepath.Dir(path), sshDirMode)) + require.NoError(t, os.WriteFile(path, []byte(content), authorizedKeysFileMode)) +} - b, err := os.ReadFile(filePath) - require.NoError(t, err) - require.Empty(t, string(b)) +func TestReconcile(t *testing.T) { + const ( + marker = " # added by dstack-shim" + // an entry added by the user by hand, which the shim must never remove + manual = testKeyOne + " added-by-hand" + // an entry added by the server when provisioning an SSH fleet + server = testKeyTwo + " dstack # added by dstack" + ) - back, err := os.ReadFile(filePath + ".bak") - require.NoError(t, err) - require.Contains(t, string(back), key) -} + testCases := []struct { + name string + content string + keys []string + expected string + }{ + { + name: "adds an entry", + content: "", + keys: []string{testKeyOne + " user@host"}, + expected: testKeyOne + " user@host" + marker + "\n", + }, + { + name: "removes an entry that is no longer in use", + content: testKeyOne + marker + "\n", + keys: nil, + expected: "", + }, + { + name: "keeps an entry that is still in use", + content: testKeyOne + " user@host" + marker + "\n", + keys: []string{testKeyOne + " user@host"}, + expected: testKeyOne + " user@host" + marker + "\n", + }, + { + // the whole point of dstackai/dstack#4174: two co-located tasks share a key, + // and the one that finishes first must not revoke it for the other + name: "collapses a key used more than once into one entry", + content: "", + keys: []string{testKeyOne + " first", testKeyOne + " second"}, + expected: testKeyOne + " first" + marker + "\n", + }, + { + name: "keeps the entries the shim does not own", + content: "# a comment\n\n" + manual + "\n" + server + "\n", + keys: nil, + expected: "# a comment\n\n" + manual + "\n" + server + "\n", + }, + { + name: "keeps a key added by hand and used by a task as two entries", + content: manual + "\n", + keys: []string{testKeyOne + " user@host"}, + expected: manual + "\n" + testKeyOne + " user@host" + marker + "\n", + }, + { + name: "removes only the entries the shim owns", + content: manual + "\n" + testKeyTwo + marker + "\n" + server + "\n", + keys: nil, + expected: manual + "\n" + server + "\n", + }, + { + name: "skips an invalid key without dropping the valid ones", + content: "", + keys: []string{"not a key", testKeyOne + "\n" + testKeyTwo, testKeyTwo}, + expected: testKeyTwo + marker + "\n", + }, + { + name: "terminates the last line of a file without a trailing newline", + content: manual, + keys: nil, + expected: manual + "\n", + }, + } -func TestRemoveTwoKey(t *testing.T) { - ak := AuthorizedKeys{user: "test_user", lookup: mockUserLookup} - filePath, err := ak.GetAuthorizedKeysPath() - require.NoError(t, err) + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ak, path := newTestAuthorizedKeys(t) + writeTestAuthorizedKeys(t, path, tc.content) - err = os.MkdirAll(path.Dir(filePath), os.ModePerm) - require.NoError(t, err) + require.NoError(t, ak.Reconcile(t.Context(), tc.keys)) - first := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - second := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCfAFHwfyMKPFbKq+D/vYNaXjqer4uV5+zvlrPY2bvkdRT4GiH4hm2s1Z7+fUEYQBNfw5O9SgxGotqyguUJbuVUc2BCNdD8HC3PxKtEev35ga4G3jjyuVeHcL2T9pn+F8IW1o3SpDGATAHJyFtArPYz31Hwg6PiuggPNdPLMSzZNrwNVuPwT1uDMKFqAh+1ryIVi7389fjZ7aBR9F06VIPpWIVVKqSVD+NbHtwWqCw8AsprJE3bPwVW09OJeQX8GXryKasaX4t4HMXmO/UI8tprnyf05dAl7NQOPY9Iut5PgfzEVY/T0M1RSnZi7i+1x7WBWX3aMM/Hv+NUeX2YtuAN" - third := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDIAGg0prDVeane6xLvMPBKQHxNUpt4q/hmuAAxjOUW0GWMPS2qE3l8YkmWeK80nKvio4M/IYWe67HIVeibdvKPoFJTtgm93WeJT9KD6h7MCschAf78mAIBhzUMK+9UYl5pE2jpfqc0SXkUsXDxMVN+ST9lN7fXUVsCPXO6qJG+0hLA3vs5r0aY1Td72vI4h45DhwjdpYkY1KTNJwfSwyvZpoN9n85JjaqXsjLG/NhieDBKu0VJE1a44aWuFwmULmpDZcUcWtk074pPMMvuh/Go5gbTaIf1gsniBKNLrfTeGjIHE/Hu9o1G3GGpq6CDqOjb0ykukWZbD2qfV0gERwIR dstack" + content, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, tc.expected, string(content)) + }) + } +} - err = os.WriteFile(filePath, []byte(first+"\n"+second+"\n"+third), os.ModePerm) - require.NoError(t, err) +func TestReconcileCreatesMissingFile(t *testing.T) { + ak, path := newTestAuthorizedKeys(t) - err = ak.RemovePublicKeys([]string{first, third}) - require.NoError(t, err) + require.NoError(t, ak.Reconcile(t.Context(), []string{testKeyOne})) - b, err := os.ReadFile(filePath) + content, err := os.ReadFile(path) require.NoError(t, err) - require.NotContains(t, string(b), first) - require.NotContains(t, string(b), third) - require.Contains(t, string(b), second) -} - -func TestAppendTwoKey(t *testing.T) { - ak := AuthorizedKeys{user: "test_user", lookup: mockUserLookup} - filePath, err := ak.GetAuthorizedKeysPath() + require.Equal(t, testKeyOne+" # added by dstack-shim\n", string(content)) + fileInfo, err := os.Stat(path) require.NoError(t, err) - - err = os.MkdirAll(path.Dir(filePath), os.ModePerm) + require.Equal(t, os.FileMode(authorizedKeysFileMode), fileInfo.Mode().Perm()) + dirInfo, err := os.Stat(filepath.Dir(path)) require.NoError(t, err) + require.Equal(t, os.FileMode(sshDirMode), dirInfo.Mode().Perm()) +} - first := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCdqa9VimGtCppxtz6T0kXfA6csnRlGS0zmTNvH2XCIYYbNFcymjL1SpFXfYQvXrnoK7nR+4dHP66um5Mi4OWHC1pB4t2OPYNnEYuYJ/VFpPv0/ykGAijV+IZjh6wS5r1o/EfiG8kMlv2TGhDb/jjsJXl9zb3i0urTrG0Sk6iw7F7QL/pXUe1cKuhdxOUzw/ddNZ5fBCikAr2cYfI0kiqe4U/pRSV5mPNAuQvBFK+K7UDdKfKIf4YxTFjXFbcgD7XUC5nInhIdSvGFYLdHSuafwWz8Q5ds/EyAPCyMU2wsA+AIP5XpdIraJLDTQT1J4PjcYwecNibWU2rkobl9FDVcflZq+0s0HbmJRlB4uExTNRZP7ykMKp9MtJsQGB6uA41KYNsvV5a+7SX39syNDHGTB13gHQHmYEHgSmHIcyEE2tEh7Zb6OAFCsytUKzBl51FIS3V70ve9kqJUcldBEkGJh6PeFOvYQZ95Gl2Uob0ujKCVDrzMylepnadfhB5dSnEM= thebits@barracuda" - second := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQCfAFHwfyMKPFbKq+D/vYNaXjqer4uV5+zvlrPY2bvkdRT4GiH4hm2s1Z7+fUEYQBNfw5O9SgxGotqyguUJbuVUc2BCNdD8HC3PxKtEev35ga4G3jjyuVeHcL2T9pn+F8IW1o3SpDGATAHJyFtArPYz31Hwg6PiuggPNdPLMSzZNrwNVuPwT1uDMKFqAh+1ryIVi7389fjZ7aBR9F06VIPpWIVVKqSVD+NbHtwWqCw8AsprJE3bPwVW09OJeQX8GXryKasaX4t4HMXmO/UI8tprnyf05dAl7NQOPY9Iut5PgfzEVY/T0M1RSnZi7i+1x7WBWX3aMM/Hv+NUeX2YtuAN" - third := "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDIAGg0prDVeane6xLvMPBKQHxNUpt4q/hmuAAxjOUW0GWMPS2qE3l8YkmWeK80nKvio4M/IYWe67HIVeibdvKPoFJTtgm93WeJT9KD6h7MCschAf78mAIBhzUMK+9UYl5pE2jpfqc0SXkUsXDxMVN+ST9lN7fXUVsCPXO6qJG+0hLA3vs5r0aY1Td72vI4h45DhwjdpYkY1KTNJwfSwyvZpoN9n85JjaqXsjLG/NhieDBKu0VJE1a44aWuFwmULmpDZcUcWtk074pPMMvuh/Go5gbTaIf1gsniBKNLrfTeGjIHE/Hu9o1G3GGpq6CDqOjb0ykukWZbD2qfV0gERwIR dstack" +func TestReconcileKeepsFileMode(t *testing.T) { + ak, path := newTestAuthorizedKeys(t) + writeTestAuthorizedKeys(t, path, "") + require.NoError(t, os.Chmod(path, 0o644)) - err = os.WriteFile(filePath, []byte(first), os.ModePerm) - require.NoError(t, err) + require.NoError(t, ak.Reconcile(t.Context(), []string{testKeyOne})) - err = ak.AppendPublicKeys(context.Background(), []string{second, third}) + fileInfo, err := os.Stat(path) require.NoError(t, err) + require.Equal(t, os.FileMode(0o644), fileInfo.Mode().Perm()) + // sshd does not read the temporary file, so leaving it behind would only make the + // next call write to a stale one + _, err = os.Stat(path + ".tmp") + require.ErrorIs(t, err, os.ErrNotExist) +} - b, err := os.ReadFile(filePath) - require.NoError(t, err) - require.Contains(t, string(b), first) - require.Contains(t, string(b), second) - require.Contains(t, string(b), third) +func TestReconcileUnknownUser(t *testing.T) { + ak := AuthorizedKeys{user: "no_such_user", lookup: mockUserLookup} + + require.Error(t, ak.Reconcile(t.Context(), []string{testKeyOne})) } diff --git a/runner/internal/shim/docker.go b/runner/internal/shim/docker.go index 818d98154..dd402e558 100644 --- a/runner/internal/shim/docker.go +++ b/runner/internal/shim/docker.go @@ -159,6 +159,11 @@ type DockerRunner struct { tasks TaskStorage // stateMu serializes task state file updates, see saveTaskState() stateMu sync.Mutex + // authorizedKeysMu serializes authorized_keys updates, covering the whole + // read-modify-write cycle, see reconcileHostSshKeys() + authorizedKeysMu sync.Mutex + // userLookup resolves a host user to its home dir and ids. Overridden in tests + userLookup func(username string) (*user.User, error) } func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*DockerRunner, error) { @@ -201,6 +206,7 @@ func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*Docke gpuVendor: gpuVendor, gpuLock: gpuLock, tasks: NewTaskStorage(), + userLookup: user.Lookup, } // The task dirs are scanned once: the tasks whose dirs are claimed by a container @@ -212,6 +218,12 @@ func NewDockerRunner(ctx context.Context, dockerParams DockerParameters) (*Docke // Must be called after the tasks are restored, as it uses them to tell the dirs of // the live tasks from the orphaned ones runner.sweepOrphanedTaskDirs(ctx, storedTasks) + // Brings authorized_keys in line with the restored tasks, dropping the entries left + // by the tasks that are gone. Only the users of the restored tasks are reconciled; + // the users of the swept tasks are already done by sweepOrphanedTaskDirs() + if err := runner.reconcileHostSshKeys(ctx); err != nil { + log.Error(ctx, "failed to reconcile host SSH keys on startup", "err", err) + } return runner, nil } @@ -543,12 +555,13 @@ func (d *DockerRunner) Start(ctx context.Context, taskID string) (err error) { } if len(cfg.HostSshKeys) > 0 { - ak := AuthorizedKeys{user: cfg.HostSshUser, lookup: user.Lookup} - if err := ak.AppendPublicKeys(ctx, cfg.HostSshKeys); err != nil { - errMessage := fmt.Sprintf("ak.AppendPublicKeys error: %s", err.Error()) + // No user is passed: the task is already stored and not cleaned up by now, + // therefore its own keys are a part of the reconciled set + if err := d.reconcileHostSshKeys(ctx); err != nil { + errMessage := fmt.Sprintf("reconcileHostSshKeys error: %s", err.Error()) log.Error(ctx, errMessage) task.SetStatusTerminated(string(types.TerminationReasonExecutorError), errMessage) - return fmt.Errorf("append public keys: %w", err) + return fmt.Errorf("reconcile host SSH keys: %w", err) } } @@ -657,12 +670,12 @@ func (d *DockerRunner) cleanupLocked(ctx context.Context, task *Task) { return } log.Debug(ctx, "releasing task resources", "task", task.ID) - releaseTaskResources(ctx, task.config) - if len(task.gpuIDs) > 0 { - releasedGpuIDs := d.gpuLock.Release(ctx, task.gpuIDs) - log.Debug(ctx, "released GPU(s)", "task", task.ID, "gpus", releasedGpuIDs) - } task.cleanedUp = true + // The flag is committed _before_ the resources are released, so that the host SSH + // keys of this task are already out of the reconciled set by the time + // releaseTaskResources() computes it. This is safe if the shim stops running in + // between: the state file is only written at the end, therefore the task is cleaned + // up again, idempotently, after a restart. // Commit the flag without touching the rest of the local copy of the task, // which may contain uncommitted changes made by the caller if _, err := d.tasks.Modify(task.ID, func(t *Task) error { @@ -671,22 +684,70 @@ func (d *DockerRunner) cleanupLocked(ctx context.Context, task *Task) { }); err != nil && !errors.Is(err, ErrNotFound) { log.Error(ctx, "failed to commit cleaned up state", "task", task.ID, "err", err) } + d.releaseTaskResources(ctx, task.config) + if len(task.gpuIDs) > 0 { + releasedGpuIDs := d.gpuLock.Release(ctx, task.gpuIDs) + log.Debug(ctx, "released GPU(s)", "task", task.ID, "gpus", releasedGpuIDs) + } d.saveTaskState(ctx, task.ID) } // releaseTaskResources releases the host resources acquired for a task: volumes and // host SSH keys. Unlike GPU locks, which are only kept in memory, these outlive the -// shim process, therefore they are released by task config and not by task -func releaseTaskResources(ctx context.Context, cfg TaskConfig) { +// shim process, therefore they are released by task config and not by task. +// The task must already be marked as cleaned up, or gone from TaskStorage altogether, +// otherwise its host SSH keys are still considered to be in use +func (d *DockerRunner) releaseTaskResources(ctx context.Context, cfg TaskConfig) { if err := unmountVolumes(ctx, cfg); err != nil { log.Error(ctx, "failed to unmount volumes", "err", err) } if len(cfg.HostSshKeys) > 0 { - ak := AuthorizedKeys{user: cfg.HostSshUser, lookup: user.Lookup} - if err := ak.RemovePublicKeys(cfg.HostSshKeys); err != nil { - log.Error(ctx, "failed to remove public keys", "err", err) + if err := d.reconcileHostSshKeys(ctx, cfg.HostSshUser); err != nil { + log.Error(ctx, "failed to reconcile host SSH keys", "err", err) + } + } +} + +// reconcileHostSshKeys brings the shim-owned entries of the host users' authorized_keys +// files in line with the tasks that still need them, that is, with the keys of all the +// stored tasks that have not been cleaned up yet. +// +// extraUsers are reconciled on top of the users of those tasks. Only a user whose tasks +// contribute no keys needs it: once the last task of a user is cleaned up, or gone from +// TaskStorage altogether, nothing names that user anymore, yet its entries are still in +// the file and have to be dropped. +// +// A failure for one user does not keep the other users from being reconciled; all the +// failures are returned joined, for the caller to report. +func (d *DockerRunner) reconcileHostSshKeys(ctx context.Context, extraUsers ...string) error { + // The lock is held for the whole read-modify-write cycle, so that a stale set of + // keys cannot overwrite a newer one. Nothing takes a task lock while holding it, + // therefore it cannot deadlock with the task lock its callers may hold + d.authorizedKeysMu.Lock() + defer d.authorizedKeysMu.Unlock() + + // Seeding the map is what makes the loop at the end visit the extra users: a user + // that no task names has no entry in the map, and so is never reconciled + keysByUser := make(map[string][]string, len(extraUsers)) + for _, username := range extraUsers { + keysByUser[username] = nil + } + for _, task := range d.tasks.List() { + cfg := task.config + if task.cleanedUp || len(cfg.HostSshKeys) == 0 { + continue + } + keysByUser[cfg.HostSshUser] = append(keysByUser[cfg.HostSshUser], cfg.HostSshKeys...) + } + + var errs []error + for username, keys := range keysByUser { + ak := AuthorizedKeys{user: username, lookup: d.userLookup} + if err := ak.Reconcile(ctx, keys); err != nil { + errs = append(errs, fmt.Errorf("user %s: %w", username, err)) } } + return errors.Join(errs...) } // Terminate aborts running operations (pulling an image, running a container) and sets task status to terminated diff --git a/runner/internal/shim/docker_test.go b/runner/internal/shim/docker_test.go index 25d26c50b..2570f0803 100644 --- a/runner/internal/shim/docker_test.go +++ b/runner/internal/shim/docker_test.go @@ -6,7 +6,9 @@ import ( "errors" "math/rand" "os" + "os/user" "path/filepath" + "strconv" "sync" "testing" "time" @@ -597,6 +599,7 @@ func newRestoreRunner(t *testing.T, gpuIDs []string, containers ...dockertypes.C gpuVendor: gpu.GpuVendorNvidia, gpuLock: gpuLock, tasks: NewTaskStorage(), + userLookup: failingUserLookup, } } @@ -711,3 +714,117 @@ func TestRestoreState_DuplicateTaskReleasesGpus(t *testing.T) { assert.True(t, runner.gpuLock.lock["GPU-beef"], "GPU-beef") assert.False(t, runner.gpuLock.lock["GPU-f00d"], "GPU-f00d is not owned by any task") } + +const testHostSshUser = "test_user" + +// newHostSshKeysRunner returns a runner whose host user's home is a temp dir, along with +// the path of that user's authorized_keys file. The mocked lookup reports the ids of the +// current process, so that setting the ownership of the file succeeds without root +func newHostSshKeysRunner(t *testing.T) (*DockerRunner, string) { + t.Helper() + usr := &user.User{ + Username: testHostSshUser, + HomeDir: t.TempDir(), + Uid: strconv.Itoa(os.Getuid()), + Gid: strconv.Itoa(os.Getgid()), + } + runner := newTestRunner(t, nil) + runner.userLookup = func(username string) (*user.User, error) { + if username != usr.Username { + return nil, errors.New("user not found") + } + return usr, nil + } + return runner, authorizedKeysPath(usr.HomeDir) +} + +func addHostSshKeysTask(t *testing.T, runner *DockerRunner, taskID string, keys ...string) Task { + t.Helper() + task := NewTaskFromConfig(TaskConfig{ + ID: taskID, + Name: taskID, + HostSshUser: testHostSshUser, + HostSshKeys: keys, + }) + require.True(t, runner.tasks.Add(task)) + return task +} + +func cleanupHostSshKeysTask(t *testing.T, runner *DockerRunner, task *Task) { + t.Helper() + task.Lock(t.Context()) + defer task.Release(t.Context()) + runner.cleanupLocked(t.Context(), task) +} + +func readAuthorizedKeysFile(t *testing.T, path string) string { + t.Helper() + content, err := os.ReadFile(path) + require.NoError(t, err) + return string(content) +} + +// TestHostSshKeys_CoLocatedTasksShareKey covers dstackai/dstack#4174: on an instance with +// blocks, the tasks co-located on it hold the same host SSH key, and the first one to +// finish must not revoke the access of the others +func TestHostSshKeys_CoLocatedTasksShareKey(t *testing.T) { + const key = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGYzO2yHhoIzYHnGH5CT/hpTNGRHvJHkKQlXqPZ0Uxwj user@host" + entry := key + " # added by dstack-shim\n" + runner, keysPath := newHostSshKeysRunner(t) + first := addHostSshKeysTask(t, runner, "task-1", key) + second := addHostSshKeysTask(t, runner, "task-2", key) + + require.NoError(t, runner.reconcileHostSshKeys(t.Context())) + // One entry, however many tasks are using the key + assert.Equal(t, entry, readAuthorizedKeysFile(t, keysPath)) + + cleanupHostSshKeysTask(t, runner, &first) + assert.Equal(t, entry, readAuthorizedKeysFile(t, keysPath), "the key is still used by task-2") + + cleanupHostSshKeysTask(t, runner, &second) + assert.Empty(t, readAuthorizedKeysFile(t, keysPath), "the last task using the key is gone") +} + +// TestHostSshKeys_KeepsKeyAddedByHand checks that a key the user added to the file +// themselves survives a task that happened to use the same key +func TestHostSshKeys_KeepsKeyAddedByHand(t *testing.T) { + const key = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGYzO2yHhoIzYHnGH5CT/hpTNGRHvJHkKQlXqPZ0Uxwj user@host" + runner, keysPath := newHostSshKeysRunner(t) + writeTestAuthorizedKeys(t, keysPath, key+"\n") + task := addHostSshKeysTask(t, runner, "task-1", key) + + require.NoError(t, runner.reconcileHostSshKeys(t.Context())) + assert.Equal(t, key+"\n"+key+" # added by dstack-shim\n", readAuthorizedKeysFile(t, keysPath)) + + cleanupHostSshKeysTask(t, runner, &task) + assert.Equal(t, key+"\n", readAuthorizedKeysFile(t, keysPath)) +} + +// TestHostSshKeys_SweepOrphanedTaskDir checks that the keys of a task that has no +// container are removed, while the keys of a task that has one are kept +func TestHostSshKeys_SweepOrphanedTaskDir(t *testing.T) { + const ( + liveKey = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIGYzO2yHhoIzYHnGH5CT/hpTNGRHvJHkKQlXqPZ0Uxwj live" + orphanKey = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAILuLmyPGV/gcatBaZFxRPKGQVJ4vBjuqEsHIkKGrGZKS orphan" + markerLine = " # added by dstack-shim\n" + ) + runner, keysPath := newHostSshKeysRunner(t) + addHostSshKeysTask(t, runner, "task-1", liveKey) + orphan := NewTaskFromConfig(TaskConfig{ + ID: "task-2", + Name: "task-2", + HostSshUser: testHostSshUser, + HostSshKeys: []string{orphanKey}, + }) + orphanDir := filepath.Join(runner.dockerParams.TasksDir(), "task-2") + require.NoError(t, os.MkdirAll(orphanDir, 0o755)) + require.NoError(t, writeTaskState(orphanDir, newTaskState(&orphan))) + stored := scanTaskDirs(t.Context(), runner.dockerParams.TasksDir()) + // Both tasks left their keys behind before the shim restarted + writeTestAuthorizedKeys(t, keysPath, liveKey+markerLine+orphanKey+markerLine) + + runner.sweepOrphanedTaskDirs(t.Context(), stored) + + assert.Equal(t, liveKey+markerLine, readAuthorizedKeysFile(t, keysPath)) + assert.NoDirExists(t, orphanDir) +} diff --git a/runner/internal/shim/process_test.go b/runner/internal/shim/process_test.go index e88e3d20b..be822248b 100644 --- a/runner/internal/shim/process_test.go +++ b/runner/internal/shim/process_test.go @@ -3,7 +3,9 @@ package shim import ( "bytes" "context" + "fmt" "io" + "os/user" "testing" dockertypes "github.com/docker/docker/api/types" @@ -217,9 +219,17 @@ func newTestRunner(t *testing.T, client docker.APIClient) *DockerRunner { dockerParams: &dockerParametersMock{tasksDir: t.TempDir()}, gpuLock: newTestGpuLock(t), tasks: NewTaskStorage(), + userLookup: failingUserLookup, } } +// failingUserLookup is the host user lookup of a test runner that does not override it. +// No test may write to a real user's authorized_keys file, therefore reaching the lookup +// unexpectedly must be an error and not a write +func failingUserLookup(username string) (*user.User, error) { + return nil, fmt.Errorf("user %s not found", username) +} + func newTestGpuLock(t *testing.T, ids ...string) *GpuLock { t.Helper() lock := make(map[string]bool, len(ids)) diff --git a/runner/internal/shim/state.go b/runner/internal/shim/state.go index 7d7243f5a..d5c02e5e6 100644 --- a/runner/internal/shim/state.go +++ b/runner/internal/shim/state.go @@ -182,7 +182,7 @@ func (d *DockerRunner) sweepOrphanedTaskDirs(ctx context.Context, storedTasks ma if !stored.state.CleanedUp { // GPU locks are in-memory, so there is nothing to release: a task without // a container holds no GPUs after a restart - releaseTaskResources(ctx, stored.state.Config) + d.releaseTaskResources(ctx, stored.state.Config) } if err := os.RemoveAll(stored.dir); err != nil { log.Error(ctx, "failed to remove orphaned task dir", "dir", stored.dir, "err", err)