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)