[ws-manager-mk2] Support public SSH keys (#16413)

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:"sshPublicKeys,omitempty"`
 }
 
 type Ownership struct {

components/ws-manager-api/go/crd/v1/zz_generated.deepcopy.go

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

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

@@ -246,7 +246,8 @@ func (wsm *WorkspaceManagerServer) StartWorkspace(ctx context.Context, req *wsma
 			Admission: workspacev1.AdmissionSpec{
 				Level: admissionLevel,
 			},
-			Ports: ports,
+			Ports:         ports,
+			SshPublicKeys: req.Spec.SshPublicKeys,
 		},
 	}
 	controllerutil.AddFinalizer(&ws, workspacev1.GitpodFinalizerName)
@@ -684,6 +685,23 @@ func (wsm *WorkspaceManagerServer) ControlAdmission(ctx context.Context, req *ws
 	return &wsmanapi.ControlAdmissionResponse{}, nil
 }
 
+func (wsm *WorkspaceManagerServer) UpdateSSHKey(ctx context.Context, req *wsmanapi.UpdateSSHKeyRequest) (res *wsmanapi.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 &wsmanapi.UpdateSSHKeyResponse{}, err
+	}
+
+	err = wsm.modifyWorkspace(ctx, req.Id, false, func(ws *workspacev1.Workspace) error {
+		ws.Spec.SshPublicKeys = req.Keys
+		return nil
+	})
+
+	return &wsmanapi.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 {
@@ -747,6 +765,19 @@ func validateStartWorkspaceRequest(req *wsmanapi.StartWorkspaceRequest) error {
 	return nil
 }
 
+func validateUpdateSSHKeyRequest(req *wsmanapi.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.(wsmanapi.WorkspaceType)
 	if !ok {

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

@@ -80,7 +80,7 @@ var runCmd = &cobra.Command{
 		infoprov = append(infoprov, podInfoProv)
 
 		if cfg.EnableWorkspaceCRD {
-			crdInfoProv, err := proxy.NewCRDWorkspaceInfoProvider(context.TODO(), mgr.GetClient(), mgr.GetScheme())
+			crdInfoProv, err := proxy.NewCRDWorkspaceInfoProvider(mgr.GetClient(), mgr.GetScheme())
 			if err == nil {
 				if err = crdInfoProv.SetupWithManager(mgr); err != nil {
 					log.WithError(err).Warn(err, "unable to create CRD-based info provider", "controller", "Workspace")

components/ws-proxy/cmd/run.go

@@ -214,7 +214,7 @@ type CRDWorkspaceInfoProvider struct {
 }
 
 // NewRemoteWorkspaceInfoProvider creates a fresh WorkspaceInfoProvider.
-func NewCRDWorkspaceInfoProvider(ctx context.Context, client client.Client, scheme *runtime.Scheme) (*CRDWorkspaceInfoProvider, error) {
+func NewCRDWorkspaceInfoProvider(client client.Client, scheme *runtime.Scheme) (*CRDWorkspaceInfoProvider, error) {
 	// create custom indexer for searches
 	indexers := cache.Indexers{
 		workspaceIndex: func(obj interface{}) ([]string, error) {
@@ -291,6 +291,7 @@ func (r *CRDWorkspaceInfoProvider) Reconcile(ctx context.Context, req ctrl.Reque
 		Ports:           ports,
 		Auth:            &wsapi.WorkspaceAuthentication{Admission: admission, OwnerToken: ws.Status.OwnerToken},
 		StartedAt:       ws.CreationTimestamp.Time,
+		SSHPublicKeys:   ws.Spec.SshPublicKeys,
 	}
 
 	r.store.Update(req.Name, wsinfo)
@@ -350,12 +351,16 @@ func extractUserSSHPublicKeys(pod *corev1.Pod) []string {
 		if err != nil {
 			return nil
 		}
-		var spec api.SSHPublicKeys
-		err = proto.Unmarshal(specPB, &spec)
-		if err != nil {
-			return nil
-		}
-		return spec.Keys
+		return unmarshalUserSSHPublicKey(specPB)
 	}
 	return nil
 }
+
+func unmarshalUserSSHPublicKey(keys []byte) []string {
+	var spec api.SSHPublicKeys
+	err := proto.Unmarshal(keys, &spec)
+	if err != nil {
+		return nil
+	}
+	return spec.Keys
+}