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

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

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

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

@@ -188,6 +188,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 +235,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 +529,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 +614,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 +694,21 @@ func setEnvironment(envs []*wsmanapi.EnvironmentVariable) []corev1.EnvVar {
 	return envVars
 }
 
+func setSshPublicKeys(keys []string) ([]byte, error) {
+	if len(keys) != 0 {
+		spec := &api.SSHPublicKeys{
+			Keys: keys,
+		}
+		sshSpec, err := proto.Marshal(spec)
+		if err != nil {
+			return nil, xerrors.Errorf("cannot create remarshal of ssh key spec: %w", err)
+		}
+		return sshSpec, nil
+	}
+
+	return nil, nil
+}
+
 func extractWorkspaceStatus(ws *workspacev1.Workspace) *wsmanapi.WorkspaceStatus {
 	version, _ := strconv.ParseUint(ws.ResourceVersion, 10, 64)
 

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:   unmarshalUserSSHPublicKey(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
+}