[ws-manager-mk2] Support public SSH keys

Author: Furisto

components/ws-manager-api/go/crd/v1/workspace_types.go

@@ -45,6 +45,8 @@ type WorkspaceSpec struct {
 
 	// +kubebuilder:validation:MinItems=0
 	Ports []PortSpec `json:"ports"`
+
+	SshPublicKeys string `json:"sshPublicKey,omitempty"`
 }
 
 type Ownership struct {

components/ws-manager-mk2/config/crd/bases/workspace.gitpod.io_workspaces.yaml

@@ -151,6 +151,8 @@ spec:
                   type: object
                 minItems: 0
                 type: array
+              sshPublicKey:
+                type: string
               sysEnvVars:
                 items:
                   description: EnvVar represents an environment variable present in

components/ws-manager-mk2/controllers/create.go

@@ -24,6 +24,7 @@ import (
 	"k8s.io/utils/pointer"
 	ctrl "sigs.k8s.io/controller-runtime"
 
+	"github.com/gitpod-io/gitpod/common-go/kubernetes"
 	wsk8s "github.com/gitpod-io/gitpod/common-go/kubernetes"
 	"github.com/gitpod-io/gitpod/common-go/tracing"
 	csapi "github.com/gitpod-io/gitpod/content-service/api"
@@ -295,6 +296,10 @@ func createDefiniteWorkspacePod(sctx *startWorkspaceContext) (*corev1.Pod, error
 		annotations[k] = v
 	}
 
+	if sctx.Workspace.Spec.SshPublicKeys != "" {
+		annotations[kubernetes.WorkspaceSSHPublicKeys] = sctx.Workspace.Spec.SshPublicKeys
+	}
+
 	// By default we embue our workspace pods with some tolerance towards pressure taints,
 	// see https://kubernetes.io/docs/concepts/configuration/taint-and-toleration/#taint-based-evictions
 	// for more details. As hope/assume that the pressure might go away in this time.

components/ws-manager-mk2/controllers/workspace_controller.go

@@ -18,6 +18,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
+	"github.com/gitpod-io/gitpod/common-go/kubernetes"
 	wsk8s "github.com/gitpod-io/gitpod/common-go/kubernetes"
 	config "github.com/gitpod-io/gitpod/ws-manager/api/config"
 	workspacev1 "github.com/gitpod-io/gitpod/ws-manager/api/crd/v1"
@@ -241,6 +242,13 @@ func (r *WorkspaceReconciler) actOnStatus(ctx context.Context, workspace *worksp
 			return ctrl.Result{Requeue: true}, err
 		}
 
+	case workspace.Spec.SshPublicKeys != pod.Annotations[kubernetes.WorkspaceSSHPublicKeys]:
+		pod.Annotations[kubernetes.WorkspaceSSHPublicKeys] = workspace.Spec.SshPublicKeys
+		err := r.Client.Update(ctx, pod)
+		if err != nil {
+			return ctrl.Result{Requeue: true}, err
+		}
+
 	// we've disposed already - try to remove the finalizer and call it a day
 	case workspace.Status.Phase == workspacev1.WorkspacePhaseStopped:
 		hadFinalizer := controllerutil.ContainsFinalizer(pod, workspacev1.GitpodFinalizerName)

components/ws-manager-mk2/service/manager.go

@@ -6,6 +6,7 @@ package service
 
 import (
 	"context"
+	"encoding/base64"
 	"fmt"
 	"strconv"
 	"sync"
@@ -188,6 +189,11 @@ func (wsm *WorkspaceManagerServer) StartWorkspace(ctx context.Context, req *wsma
 		}
 	}
 
+	sshPublicKeys, err := setSshPublicKeys(req.Spec.SshPublicKeys)
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid ssh public keys")
+	}
+
 	ws := workspacev1.Workspace{
 		TypeMeta: metav1.TypeMeta{
 			APIVersion: workspacev1.GroupVersion.String(),
@@ -230,7 +236,8 @@ func (wsm *WorkspaceManagerServer) StartWorkspace(ctx context.Context, req *wsma
 			Admission: workspacev1.AdmissionSpec{
 				Level: admissionLevel,
 			},
-			Ports: ports,
+			Ports:         ports,
+			SshPublicKeys: sshPublicKeys,
 		},
 	}
 	controllerutil.AddFinalizer(&ws, workspacev1.GitpodFinalizerName)
@@ -523,6 +530,28 @@ func (wsm *WorkspaceManagerServer) ControlAdmission(ctx context.Context, req *ws
 	return &wsmanapi.ControlAdmissionResponse{}, nil
 }
 
+func (wsm *WorkspaceManagerServer) UpdateSSHKey(ctx context.Context, req *api.UpdateSSHKeyRequest) (res *api.UpdateSSHKeyResponse, err error) {
+	span, ctx := tracing.FromContext(ctx, "UpdateSSHKey")
+	tracing.ApplyOWI(span, log.OWI("", "", req.Id))
+	defer tracing.FinishSpan(span, &err)
+
+	if err = validateUpdateSSHKeyRequest(req); err != nil {
+		return &api.UpdateSSHKeyResponse{}, err
+	}
+
+	err = wsm.modifyWorkspace(ctx, req.Id, false, func(ws *workspacev1.Workspace) error {
+		sshKeys, err := setSshPublicKeys(req.Keys)
+		if err != nil {
+			return err
+		}
+
+		ws.Spec.SshPublicKeys = sshKeys
+		return nil
+	})
+
+	return &api.UpdateSSHKeyResponse{}, err
+}
+
 // modifyWorkspace modifies a workspace object using the mod function. If the mod function returns a gRPC status error, that error
 // is returned directly. If mod returns a non-gRPC error it is turned into one.
 func (wsm *WorkspaceManagerServer) modifyWorkspace(ctx context.Context, id string, updateStatus bool, mod func(ws *workspacev1.Workspace) error) error {
@@ -586,6 +615,19 @@ func validateStartWorkspaceRequest(req *api.StartWorkspaceRequest) error {
 	return nil
 }
 
+func validateUpdateSSHKeyRequest(req *api.UpdateSSHKeyRequest) error {
+	err := validation.ValidateStruct(req,
+		validation.Field(&req.Id, validation.Required),
+		validation.Field(&req.Keys, validation.Required),
+	)
+
+	if err != nil {
+		return status.Errorf(codes.InvalidArgument, "invalid request: %v", err)
+	}
+
+	return nil
+}
+
 func isValidWorkspaceType(value interface{}) error {
 	s, ok := value.(api.WorkspaceType)
 	if !ok {
@@ -653,6 +695,21 @@ func setEnvironment(envs []*wsmanapi.EnvironmentVariable) []corev1.EnvVar {
 	return envVars
 }
 
+func setSshPublicKeys(keys []string) (string, error) {
+	if len(keys) != 0 {
+		spec := &api.SSHPublicKeys{
+			Keys: keys,
+		}
+		sshSpec, err := proto.Marshal(spec)
+		if err != nil {
+			return "", xerrors.Errorf("cannot create remarshal of ssh key spec: %w", err)
+		}
+		return base64.StdEncoding.EncodeToString(sshSpec), nil
+	}
+
+	return "", nil
+}
+
 func extractWorkspaceStatus(ws *workspacev1.Workspace) *wsmanapi.WorkspaceStatus {
 	version, _ := strconv.ParseUint(ws.ResourceVersion, 10, 64)