Allow to run.attach()

runner/internal/runner/api/http_test.go

@@ -21,7 +21,7 @@ func (ds DummyRunner) GetState() (shim.RunnerStatus, shim.ContainerStatus, strin
 	return ds.State, ds.ContainerStatus, "", ds.JobResult
 }
 
-func (ds DummyRunner) Run(context.Context, shim.DockerImageConfig) error {
+func (ds DummyRunner) Run(context.Context, shim.TaskConfig) error {
 	return nil
 }
 

runner/internal/shim/api/http.go

@@ -28,18 +28,18 @@ func (s *ShimServer) SubmitPostHandler(w http.ResponseWriter, r *http.Request) (
 		return nil, &api.Error{Status: http.StatusConflict}
 	}
 
-	var body DockerTaskBody
+	var body TaskConfigBody
 	if err := api.DecodeJSONBody(w, r, &body, true); err != nil {
 		log.Println("Failed to decode submit body", "err", err)
 		return nil, err
 	}
 
-	go func(taskParams shim.DockerImageConfig) {
-		err := s.runner.Run(context.Background(), taskParams)
+	go func(taskConfig shim.TaskConfig) {
+		err := s.runner.Run(context.Background(), taskConfig)
 		if err != nil {
 			fmt.Printf("failed Run %v\n", err)
 		}
-	}(body.GetDockerImageConfig())
+	}(body.GetTaskConfig())
 
 	return nil, nil
 }

runner/internal/shim/api/schemas.go

@@ -2,13 +2,15 @@ package api
 
 import "github.com/dstackai/dstack/runner/internal/shim"
 
-type DockerTaskBody struct {
+type TaskConfigBody struct {
 	Username      string   `json:"username"`
 	Password      string   `json:"password"`
 	ImageName     string   `json:"image_name"`
 	ContainerName string   `json:"container_name"`
 	ShmSize       int64    `json:"shm_size"`
 	PublicKeys    []string `json:"public_keys"`
+	SshUser       string   `json:"ssh_user"`
+	SshKey        string   `json:"ssh_key"`
 }
 
 type StopBody struct {
@@ -37,14 +39,16 @@ type StopResponse struct {
 	State string `json:"state"`
 }
 
-func (ra DockerTaskBody) GetDockerImageConfig() shim.DockerImageConfig {
-	res := shim.DockerImageConfig{
+func (ra TaskConfigBody) GetTaskConfig() shim.TaskConfig {
+	res := shim.TaskConfig{
 		ImageName:     ra.ImageName,
 		Username:      ra.Username,
 		Password:      ra.Password,
 		ContainerName: ra.ContainerName,
 		ShmSize:       ra.ShmSize,
 		PublicKeys:    ra.PublicKeys,
+		SshUser:       ra.SshUser,
+		SshKey:        ra.SshKey,
 	}
 	return res
 }

runner/internal/shim/api/server.go

@@ -10,7 +10,7 @@ import (
 )
 
 type TaskRunner interface {
-	Run(context.Context, shim.DockerImageConfig) error
+	Run(context.Context, shim.TaskConfig) error
 	GetState() (shim.RunnerStatus, shim.ContainerStatus, string, shim.JobResult)
 	Stop(bool)
 }

runner/internal/shim/authorized_keys.go

@@ -0,0 +1,141 @@
+package shim
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"os"
+	"slices"
+
+	"github.com/ztrue/tracerr"
+	"golang.org/x/crypto/ssh"
+)
+
+func PublicKeyFingerprint(key string) (string, error) {
+	pk, _, _, _, err := ssh.ParseAuthorizedKey([]byte(key))
+	if err != nil {
+		return "", tracerr.Wrap(err)
+	}
+	keyFingerprint := ssh.FingerprintSHA256(pk)
+	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
+}
+
+type AuthorizedKeys struct {
+	user     string
+	rootPath string
+}
+
+func (ak AuthorizedKeys) AppendPublicKeys(publicKeys []string) error {
+	return ak.transformAuthorizedKeys(AppendPublicKeys, publicKeys)
+}
+
+func (ak AuthorizedKeys) RemovePublicKeys(publicKeys []string) error {
+	return ak.transformAuthorizedKeys(RemovePublicKeys, publicKeys)
+}
+
+func (ak AuthorizedKeys) read(r io.Reader) ([]string, error) {
+	lines := []string{}
+	scanner := bufio.NewScanner(r)
+	for scanner.Scan() {
+		text := scanner.Text()
+		lines = append(lines, text)
+	}
+	if err := scanner.Err(); err != nil {
+		return []string{}, tracerr.Wrap(err)
+	}
+	return lines, nil
+}
+
+func (ak AuthorizedKeys) write(w io.Writer, lines []string) error {
+	wr := bufio.NewWriter(w)
+	for _, line := range lines {
+		_, err := fmt.Fprintln(wr, line)
+		if err != nil {
+			return tracerr.Wrap(err)
+
+		}
+	}
+	return wr.Flush()
+}
+
+func (ak AuthorizedKeys) GetAuthorizedKeysPath() string {
+	return fmt.Sprintf("%s/home/%s/.ssh/authorized_keys", ak.rootPath, ak.user)
+}
+
+func (ak AuthorizedKeys) transformAuthorizedKeys(transform func([]string, []string) []string, publicKeys []string) error {
+	authorizedKeysPath := ak.GetAuthorizedKeysPath()
+	info, err := os.Stat(authorizedKeysPath)
+	if err != nil {
+		return tracerr.Wrap(err)
+	}
+	fileMode := info.Mode().Perm()
+
+	authorizedKeysFile, err := os.OpenFile(authorizedKeysPath, os.O_RDWR, fileMode)
+	if err != nil {
+		return tracerr.Wrap(err)
+	}
+	defer authorizedKeysFile.Close()
+
+	lines, err := ak.read(authorizedKeysFile)
+	if err != nil {
+		return tracerr.Wrap(err)
+	}
+
+	// write backup
+	authorizedKeysPathBackup := ak.GetAuthorizedKeysPath() + ".bak"
+	authorizedKeysBackup, err := os.OpenFile(authorizedKeysPathBackup, os.O_RDWR|os.O_CREATE|os.O_TRUNC, fileMode)
+	if err != nil {
+		return tracerr.Wrap(err)
+	}
+	defer authorizedKeysBackup.Close()
+	if err := ak.write(authorizedKeysBackup, lines); err != nil {
+		return tracerr.Wrap(err)
+	}
+
+	// transform lines
+	newLines := transform(lines, publicKeys)
+
+	// write authorized_keys
+	if err := authorizedKeysFile.Truncate(0); err != nil {
+		return tracerr.Wrap(err)
+	}
+	if _, err := authorizedKeysFile.Seek(0, 0); err != nil {
+		return tracerr.Wrap(err)
+	}
+	if err := ak.write(authorizedKeysFile, newLines); err != nil {
+		return tracerr.Wrap(err)
+	}
+
+	return nil
+}