HP-682 Feat/multiple paymodels

hatchery/config.go

@@ -56,19 +56,29 @@ type AppConfigInfo struct {
 
 // TODO remove PayModel from config once DynamoDB contains all necessary data
 type PayModel struct {
-	Name         string `json:"name"`
-	User         string `json:"user_id"`
-	AWSAccountId string `json:"aws_account_id"`
-	Region       string `json:"region"`
-	Ecs          string `json:"ecs"`
-	VpcId        string `json:"vpcid"`
-	Subnet       int    `json:"subnet"`
+	Id              string  `json:"bmh_workspace_id"`
+	Name            string  `json:"workspace_type"`
+	User            string  `json:"user_id"`
+	AWSAccountId    string  `json:"account_id"`
+	Region          string  `json:"region"`
+	Ecs             string  `json:"ecs"`
+	Subnet          int     `json:"subnet"`
+	HardLimit       float32 `json:"hard-limit"`
+	SoftLimit       float32 `json:"soft-limit"`
+	TotalUsage      float32 `json:"total-usage"`
+	CurrentPayModel bool    `json:"current_pay_model"`
+}
+
+type AllPayModels struct {
+	CurrentPayModel *PayModel  `json:"current_pay_model"`
+	PayModels       []PayModel `json:"all_pay_models"`
 }
 
 // HatcheryConfig is the root of all the configuration
 type HatcheryConfig struct {
 	UserNamespace          string           `json:"user-namespace"`
 	DefaultPayModel        PayModel         `json:"default-pay-model"`
+	DisableLocalWS         bool             `json:"disable-local-ws"`
 	PayModels              []PayModel       `json:"pay-models"`
 	PayModelsDynamodbTable string           `json:"pay-models-dynamodb-table"`
 	SubDir                 string           `json:"sub-dir"`

hatchery/ecs.go

@@ -339,8 +339,6 @@ func terminateEcsWorkspace(ctx context.Context, userName string, accessToken str
 }
 
 func launchEcsWorkspace(ctx context.Context, userName string, hash string, accessToken string, payModel PayModel) error {
-	// TODO: Setup EBS volume as pd
-	// Must create volume using SDK too.. :(
 	roleARN := "arn:aws:iam::" + payModel.AWSAccountId + ":role/csoc_adminvm"
 	sess := session.Must(session.NewSession(&aws.Config{
 		// TODO: Make this configurable
@@ -486,10 +484,6 @@ func launchEcsWorkspace(ctx context.Context, userName string, hash string, acces
 		}
 		return err
 	}
-	err = setupTransitGateway(userName)
-	if err != nil {
-		return err
-	}
 
 	launchTask, err := svc.launchService(ctx, taskDefResult, userName, hash, payModel)
 	if err != nil {
@@ -500,6 +494,11 @@ func launchEcsWorkspace(ctx context.Context, userName string, hash string, acces
 		return err
 	}
 
+	err = setupTransitGateway(userName)
+	if err != nil {
+		return err
+	}
+
 	fmt.Printf("Launched ECS workspace service at %s for user %s\n", launchTask, userName)
 	return nil
 }

hatchery/efs.go

@@ -63,10 +63,10 @@ func (creds *CREDS) createAccessPoint(FileSystemId string, userName string, svc
 	if err != nil {
 		return nil, err
 	}
-
+	ap := userToResourceName(userName, "service") + "-" + strings.ReplaceAll(os.Getenv("GEN3_ENDPOINT"), ".", "-") + "-accesspoint"
 	if len(exResult.AccessPoints) == 0 {
 		input := &efs.CreateAccessPointInput{
-			ClientToken:  aws.String(fmt.Sprintf("ap-%s", userToResourceName(userName, "pod"))),
+			ClientToken:  aws.String(ap),
 			FileSystemId: aws.String(FileSystemId),
 			PosixUser: &efs.PosixUser{
 				Gid: aws.Int64(100),

hatchery/hatchery.go

@@ -24,6 +24,8 @@ func RegisterHatchery(mux *httptrace.ServeMux) {
 	mux.HandleFunc("/status", status)
 	mux.HandleFunc("/options", options)
 	mux.HandleFunc("/paymodels", paymodels)
+	mux.HandleFunc("/setpaymodel", setpaymodel)
+	mux.HandleFunc("/allpaymodels", allpaymodels)
 
 	// ECS functions
 	mux.HandleFunc("/create-ecs-cluster", createECSCluster)
@@ -55,16 +57,65 @@ func paymodels(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	userName := getCurrentUserName(r)
-	payModel, err := getPayModelForUser(userName)
+
+	payModel, err := getCurrentPayModel(userName)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
 	if payModel == nil {
-		http.Error(w, err.Error(), http.StatusNotFound)
+		http.Error(w, "Current paymodel not set", http.StatusNotFound)
 		return
 	}
+	out, err := json.Marshal(payModel)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
-	out, err := json.Marshal(payModel)
+	fmt.Fprint(w, string(out))
+}
+
+func allpaymodels(w http.ResponseWriter, r *http.Request) {
+	if r.Method != "GET" {
+		http.Error(w, "Not Found", http.StatusNotFound)
+		return
+	}
+	userName := getCurrentUserName(r)
+
+	payModels, err := getPayModelsForUser(userName)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	if payModels == nil {
+		http.Error(w, "No paymodel set", http.StatusNotFound)
+		return
+	}
+	out, err := json.Marshal(payModels)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	fmt.Fprint(w, string(out))
+}
+
+func setpaymodel(w http.ResponseWriter, r *http.Request) {
+	if r.Method != "GET" {
+		http.Error(w, "Not Found", http.StatusNotFound)
+		return
+	}
+	userName := getCurrentUserName(r)
+	id := r.URL.Query().Get("id")
+	if id == "" {
+		http.Error(w, "Missing ID argument", http.StatusBadRequest)
+		return
+	}
+	pm, err := setCurrentPaymodel(userName, id)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	out, err := json.Marshal(pm)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
@@ -76,19 +127,27 @@ func status(w http.ResponseWriter, r *http.Request) {
 	userName := getCurrentUserName(r)
 	accessToken := getBearerToken(r)
 
-	payModel, err := getPayModelForUser(userName)
+	payModel, err := getCurrentPayModel(userName)
 	if err != nil {
-		Config.Logger.Printf(err.Error())
+		if err != NopaymodelsError {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
 	}
 	var result *WorkspaceStatus
-	if payModel != nil && payModel.Ecs == "true" {
+
+	if payModel.Ecs == "true" {
 		result, err = statusEcs(r.Context(), userName, accessToken, payModel.AWSAccountId)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
 	} else {
 		result, err = statusK8sPod(r.Context(), userName, accessToken, payModel)
-	}
-	if err != nil {
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
 	}
 
 	out, err := json.Marshal(result)
@@ -154,7 +213,7 @@ func launch(w http.ResponseWriter, r *http.Request) {
 	}
 
 	userName := getCurrentUserName(r)
-	payModel, err := getPayModelForUser(userName)
+	payModel, err := getCurrentPayModel(userName)
 	if err != nil {
 		Config.Logger.Printf(err.Error())
 	}
@@ -166,6 +225,7 @@ func launch(w http.ResponseWriter, r *http.Request) {
 		err = createExternalK8sPod(r.Context(), hash, userName, accessToken, *payModel)
 	}
 	if err != nil {
+		Config.Logger.Printf("error during launch: %-v", err)
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
@@ -179,7 +239,7 @@ func terminate(w http.ResponseWriter, r *http.Request) {
 	}
 	accessToken := getBearerToken(r)
 	userName := getCurrentUserName(r)
-	payModel, err := getPayModelForUser(userName)
+	payModel, err := getCurrentPayModel(userName)
 	if err != nil {
 		Config.Logger.Printf(err.Error())
 	}
@@ -219,7 +279,7 @@ func getBearerToken(r *http.Request) string {
 // TODO: NEED TO CALL THIS FUNCTION IF IT DOESN'T EXIST!!!
 func createECSCluster(w http.ResponseWriter, r *http.Request) {
 	userName := getCurrentUserName(r)
-	payModel, err := getPayModelForUser(userName)
+	payModel, err := getCurrentPayModel(userName)
 	if payModel == nil {
 		http.Error(w, "Paymodel has not been setup for user", http.StatusNotFound)
 		return

hatchery/iam.go

@@ -14,7 +14,10 @@ func (creds *CREDS) taskRole(userName string) (*string, error) {
 		Credentials: creds.creds,
 		Region:      aws.String("us-east-1"),
 	})))
-	pm := Config.PayModelMap[userName]
+	pm, err := getCurrentPayModel(userName)
+	if err != nil {
+		return nil, err
+	}
 	policyArn := fmt.Sprintf("arn:aws:iam::%s:policy/%s", pm.AWSAccountId, fmt.Sprintf("ws-task-policy-%s", userName))
 	taskRoleInput := &iam.GetRoleInput{
 		RoleName: aws.String(userToResourceName(userName, "pod")),
@@ -96,8 +99,9 @@ func (creds *CREDS) taskRole(userName string) (*string, error) {
 	}
 
 }
-// https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_execution_IAM_role.html 
-// The task execution role grants the Amazon ECS container and Fargate agents permission to make AWS API calls on your behalf. 
+
+// https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_execution_IAM_role.html
+// The task execution role grants the Amazon ECS container and Fargate agents permission to make AWS API calls on your behalf.
 const ecsTaskExecutionRoleName = "ecsTaskExecutionRole"
 const ecsTaskExecutionPolicyArn = "arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy"
 const ecsTaskExecutionRoleAssumeRolePolicyDocument = `{