Webapp: Don't use DM for IAM.

bootstrap/cmd/bootstrap/app/gcpUtils.go

@@ -1,10 +1,15 @@
 package app
 
 import (
+	"fmt"
+	"time"
+
 	"golang.org/x/net/context"
 	"google.golang.org/api/deploymentmanager/v2"
 	"golang.org/x/oauth2"
 	"github.com/ghodss/yaml"
+	log "github.com/sirupsen/logrus"
+	"google.golang.org/api/cloudresourcemanager/v1"
 	"io/ioutil"
 	"path"
 )
@@ -20,6 +25,23 @@ type DmConf struct {
 	Resources []Resource `json:"resources"`
 }
 
+type IamBinding struct {
+	Members []string `type:"members`
+	Roles   []string `type:"roles"`
+}
+
+type IamConf struct {
+	IamBindings []IamBinding `json:"bindings"`
+}
+
+type ApplyIamRequest struct {
+	Project string `json:"project"`
+	Cluster string `json:"cluster"`
+	Email string `json:"email"`
+	Token string `json:"token"`
+	Action string `json:"action`
+}
+
 // TODO: handle concurrent & repetitive deployment requests.
 func (s *ksServer)InsertDeployment(ctx context.Context, req CreateRequest) error {
 	regPath := s.knownRegistries["kubeflow"].RegUri
@@ -83,4 +105,113 @@ func (s *ksServer)GetDeploymentStatus(ctx context.Context, req CreateRequest) (s
 		return "", err
 	}
 	return dm.Operation.Status, nil
-}
+}
+
+func GetUpdatedPolicy(currentPolicy *cloudresourcemanager.Policy, iamConf *IamConf, req ApplyIamRequest) cloudresourcemanager.Policy {
+	// map from role to members.
+	policyMap := map[string]map[string]bool {}
+	for _, binding := range currentPolicy.Bindings {
+		policyMap[binding.Role] = make(map[string]bool)
+		for _, member := range binding.Members {
+			policyMap[binding.Role][member] = true
+		}
+	}
+
+	// Replace placeholder with actual identity.
+	saMapping := map[string]string {
+		"set-kubeflow-admin-service-account": fmt.Sprintf("serviceAccount:%v-admin@%v.iam.gserviceaccount.com", req.Cluster, req.Project),
+		"set-kubeflow-user-service-account": fmt.Sprintf("serviceAccount:%v-user@%v.iam.gserviceaccount.com", req.Cluster, req.Project),
+		"set-kubeflow-vm-service-account": fmt.Sprintf("serviceAccount:%v-vm@%v.iam.gserviceaccount.com", req.Cluster, req.Project),
+		"set-kubeflow-iap-account": fmt.Sprintf("user:%v", req.Email),
+	}
+	for _, binding := range iamConf.IamBindings {
+		for _, member := range binding.Members {
+			actualMember := member
+			if val, ok := saMapping[member]; ok {
+				actualMember = val
+			}
+			for _, role := range binding.Roles {
+				if req.Action == "add" {
+					policyMap[role][actualMember] = true
+				} else {
+					// action == "remove"
+					policyMap[role][actualMember] = false
+				}
+			}
+		}
+	}
+	newPolicy := cloudresourcemanager.Policy{}
+	for role, memberSet := range policyMap {
+		binding := cloudresourcemanager.Binding{}
+		binding.Role = role
+		for member, exists := range memberSet {
+			if exists {
+				binding.Members = append(binding.Members, member)
+			}
+		}
+		newPolicy.Bindings = append(newPolicy.Bindings, &binding)
+	}
+	return newPolicy
+}
+
+func (s *ksServer)ApplyIamPolicy(ctx context.Context, req ApplyIamRequest) error {
+	ts := oauth2.StaticTokenSource(&oauth2.Token{
+		AccessToken: req.Token,
+	})
+	resourceManager, err := cloudresourcemanager.New(oauth2.NewClient(ctx, ts))
+	if err != nil {
+		log.Errorf("Cannot create resource manager client: %v", err)
+		return err
+	}
+	s.serverMux.Lock()
+	defer s.serverMux.Unlock()
+
+	// Get current policy
+	retry := 0
+	var saPolicy *cloudresourcemanager.Policy
+	for retry < 5 {
+		saPolicy, err = resourceManager.Projects.GetIamPolicy(
+			req.Project,
+			&cloudresourcemanager.GetIamPolicyRequest{
+			}).Do()
+		if err != nil {
+			retry += 1
+			time.Sleep(3 * time.Second)
+		}
+	}
+	if err != nil {
+		log.Errorf("Cannot get current policy: %v", err)
+		return err
+	}
+
+    // Get the iam change from config.
+	regPath := s.knownRegistries["kubeflow"].RegUri
+	templatePath := path.Join(regPath, "../components/gcp-click-to-deploy/src/configs/iam_bindings_template.yaml")
+	var iamConf IamConf
+	err = LoadConfig(templatePath, &iamConf)
+	if err != nil {
+		log.Errorf("Failed to load iam config: %v", err)
+		return err
+	}
+
+	// Get the updated policy and apply it.
+	newPolicy := GetUpdatedPolicy(saPolicy, &iamConf, req)
+	retry = 0
+	if retry < 5 {
+		_, err = resourceManager.Projects.SetIamPolicy(
+			req.Project,
+			&cloudresourcemanager.SetIamPolicyRequest{
+				Policy: &newPolicy,
+			}).Do()
+		if err != nil {
+			retry += 1
+			time.Sleep(3 * time.Second)
+		}
+	}
+	if err != nil {
+		log.Errorf("Cannot set new ploicy: %v", err)
+		return err
+	}
+
+	return nil
+}

bootstrap/cmd/bootstrap/app/ksServer.go

@@ -48,6 +48,7 @@ type KsService interface {
 	BindRole(context.Context, string, string, string) error
 	InsertDeployment(context.Context, CreateRequest) error
 	GetDeploymentStatus(context.Context, CreateRequest) (string, error)
+	ApplyIamPolicy(context.Context, ApplyIamRequest) error
 }
 
 // appInfo keeps track of information about apps.
@@ -813,6 +814,19 @@ func finishDeployment(svc KsService, req CreateRequest) {
 		return
 	}
 
+	log.Info("Patching IAM bindings...")
+	err = svc.ApplyIamPolicy(ctx, ApplyIamRequest{
+		Project: req.Project,
+		Cluster: req.Cluster,
+		Email: req.Email,
+		Token: req.Token,
+		Action: "add",
+	})
+	if err != nil {
+		log.Errorf("Failed to update IAM: %v", err)
+		return
+	}
+
 	log.Infof("Inserting sa keys...")
 	err = svc.InsertSaKeys(ctx, InsertSaKeyRequest{
 		Cluster: req.Cluster,
@@ -897,6 +911,18 @@ func makeSaKeyEndpoint(svc KsService) endpoint.Endpoint {
 	}
 }
 
+func makeIamEndpoint(svc KsService) endpoint.Endpoint {
+	return func(ctx context.Context, request interface{}) (interface{}, error) {
+		req := request.(ApplyIamRequest)
+		err := svc.ApplyIamPolicy(ctx, req)
+		r := &basicServerResponse{}
+		if err != nil {
+			r.Err = err.Error()
+		}
+		return r, nil
+	}
+}
+
 func decodeCreateAppRequest(_ context.Context, r *http.Request) (interface{}, error) {
 	var request CreateRequest
 	if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
@@ -971,6 +997,18 @@ func (s *ksServer) StartHttp(port int) {
 		encodeResponse,
 	)
 
+	applyIamHandler := httptransport.NewServer(
+		makeIamEndpoint(s),
+		func (_ context.Context, r *http.Request) (interface{}, error) {
+			var request ApplyIamRequest
+			if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
+				return nil, err
+			}
+			return request, nil
+		},
+		encodeResponse,
+	)
+
 	initProjectHandler := httptransport.NewServer(
 		makeInitProjectEndpoint(s),
 		func (_ context.Context, r *http.Request) (interface{}, error) {
@@ -995,6 +1033,7 @@ func (s *ksServer) StartHttp(port int) {
 	http.Handle("/kfctl/apps/apply", optionsHandler(applyAppHandler))
 	http.Handle("/kfctl/apps/create", optionsHandler(createAppHandler))
 	http.Handle("/kfctl/iam/insertSaKey", optionsHandler(insertSaKeyHandler))
+	http.Handle("/kfctl/iam/apply", optionsHandler(applyIamHandler))
 	http.Handle("/kfctl/initProject", optionsHandler(initProjectHandler))
 	http.Handle("/kfctl/e2eDeploy", optionsHandler(deployHandler))