[shim] Mark authorized_keys entries added by the shim

runner/internal/shim/authorized_keys.go

@@ -2,16 +2,37 @@ package shim
 
 import (
 	"bufio"
+	"bytes"
+	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
 	"os/user"
 	"path/filepath"
 	"slices"
+	"strings"
 
 	"golang.org/x/crypto/ssh"
+
+	"github.com/dstackai/dstack/runner/internal/common/log"
 )
 
+// publicKeyMarker is appended to the comment field of every authorized_keys entry
+// added by the shim, mirroring the `# added by dstack` marker the server writes when
+// 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 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.
+const publicKeyMarker = "# added by dstack-shim"
+
 func PublicKeyFingerprint(key string) (string, error) {
 	pk, _, _, _, err := ssh.ParseAuthorizedKey([]byte(key))
 	if err != nil {
@@ -52,15 +73,66 @@ func AppendPublicKeys(fileKeys []string, keysToAppend []string) []string {
 	return newKeys
 }
 
+// canonicalizePublicKey validates a public key received from the server and returns the
+// authorized_keys line to write for it, marked with publicKeyMarker.
+//
+// The line is rebuilt from the parsed key instead of reusing the original entry, so that
+// nothing unvalidated reaches the file.
+func canonicalizePublicKey(publicKey string) (string, error) {
+	key, comment, options, rest, err := ssh.ParseAuthorizedKey([]byte(publicKey))
+	if err != nil {
+		return "", fmt.Errorf("parse public key: %w", err)
+	}
+	// ParseAuthorizedKey stops at the end of the first key it finds
+	if len(bytes.TrimSpace(rest)) > 0 {
+		return "", errors.New("more than one key in a single entry")
+	}
+	// Options are a feature of the authorized_keys format, not of the on-disk public key
+	// format, therefore an entry carrying them is not a public key
+	if len(options) > 0 {
+		return "", errors.New("unexpected authorized_keys options")
+	}
+	// MarshalAuthorizedKey returns a "<type> <base64>\n" line
+	keyLine := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(key)))
+	// The comment cannot span lines, but may contain other whitespace, e.g., \t or \r
+	commentFields := strings.Fields(comment)
+	commentFields = append(commentFields, publicKeyMarker)
+	return keyLine + " " + strings.Join(commentFields, " "), nil
+}
+
 type AuthorizedKeys struct {
 	user   string
 	lookup func(username string) (*user.User, error)
 }
 
-func (ak AuthorizedKeys) AppendPublicKeys(publicKeys []string) error {
-	return ak.transformAuthorizedKeys(AppendPublicKeys, publicKeys)
+// 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 {
+	lines := make([]string, 0, len(publicKeys))
+	for _, publicKey := range publicKeys {
+		line, err := canonicalizePublicKey(publicKey)
+		if err != nil {
+			// Whitespace is collapsed to keep the entry on a single log line
+			log.Error(
+				ctx, "skipping invalid public key",
+				"key", strings.Join(strings.Fields(publicKey), " "), "err", err,
+			)
+			continue
+		}
+		lines = append(lines, line)
+	}
+	if len(lines) == 0 {
+		return nil
+	}
+	return ak.transformAuthorizedKeys(AppendPublicKeys, 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)
 }

runner/internal/shim/authorized_keys_test.go

@@ -1,6 +1,7 @@
 package shim
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/user"
@@ -161,13 +162,103 @@ func TestAppendKey(t *testing.T) {
 	err = os.WriteFile(filePath, []byte(key), os.ModePerm)
 	require.NoError(t, err)
 
-	commentLine := "# comment line"
-	err = ak.AppendPublicKeys([]string{commentLine})
+	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), commentLine)
+	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"
+
+	testCases := []struct {
+		name     string
+		key      string
+		expected string
+		isError  bool
+	}{
+		{
+			name:     "with comment",
+			key:      blob + " user@host",
+			expected: blob + " user@host # added by dstack-shim",
+		},
+		{
+			name:     "without comment",
+			key:      blob,
+			expected: blob + " # added by dstack-shim",
+		},
+		{
+			name:     "surrounding whitespace",
+			key:      "  " + blob + "\tuser@host \r\n",
+			expected: blob + " user@host # added by dstack-shim",
+		},
+		{
+			name:    "two keys in one entry",
+			key:     blob + " user@host\n" + blob + " user@host",
+			isError: true,
+		},
+		{
+			name: "authorized_keys options",
+			// options are not part of the on-disk public key format
+			key:     `restrict,command="/bin/false" ` + blob + " user@host",
+			isError: true,
+		},
+		{
+			name:    "comment line",
+			key:     "# comment line",
+			isError: true,
+		},
+		{
+			name:    "malformed",
+			key:     "ssh-ed25519 AAAAP66um5MadfhB5dSnEM=",
+			isError: true,
+		},
+		{
+			name:    "empty",
+			key:     "",
+			isError: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			line, err := canonicalizePublicKey(tc.key)
+			if tc.isError {
+				require.Error(t, err)
+				require.Empty(t, line)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tc.expected, line)
+			}
+		})
+	}
 }
 
 func TestRemoveKey(t *testing.T) {
@@ -234,7 +325,7 @@ func TestAppendTwoKey(t *testing.T) {
 	err = os.WriteFile(filePath, []byte(first), os.ModePerm)
 	require.NoError(t, err)
 
-	err = ak.AppendPublicKeys([]string{second, third})
+	err = ak.AppendPublicKeys(context.Background(), []string{second, third})
 	require.NoError(t, err)
 
 	b, err := os.ReadFile(filePath)

runner/internal/shim/docker.go

@@ -488,7 +488,7 @@ 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(cfg.HostSshKeys); err != nil {
+		if err := ak.AppendPublicKeys(ctx, cfg.HostSshKeys); err != nil {
 			errMessage := fmt.Sprintf("ak.AppendPublicKeys error: %s", err.Error())
 			log.Error(ctx, errMessage)
 			task.SetStatusTerminated(string(types.TerminationReasonExecutorError), errMessage)