environment disable now cleans up agent remote; refactored and streamlined disableHandler (#967)

controller/disable.go

@@ -1,163 +1,181 @@
 package controller
 
 import (
+	"context"
 	"fmt"
 	"github.com/go-openapi/runtime/middleware"
 	"github.com/jmoiron/sqlx"
 	"github.com/openziti/edge-api/rest_management_api_client"
+	edge_service "github.com/openziti/edge-api/rest_management_api_client/service"
+	"github.com/openziti/zrok/controller/store"
 	"github.com/openziti/zrok/controller/zrokEdgeSdk"
 	"github.com/openziti/zrok/rest_model_zrok"
 	"github.com/openziti/zrok/rest_server_zrok/operations/environment"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
+	"time"
 )
 
-type disableHandler struct {
+type disableHandler struct{}
-}
 
 func newDisableHandler() *disableHandler {
 	return &disableHandler{}
 }
 
 func (h *disableHandler) Handle(params environment.DisableParams, principal *rest_model_zrok.Principal) middleware.Responder {
-	tx, err := str.Begin()
+	trx, err := str.Begin()
 	if err != nil {
 		logrus.Errorf("error starting transaction for user '%v': %v", principal.Email, err)
 		return environment.NewDisableInternalServerError()
 	}
-	defer func() { _ = tx.Rollback() }()
+	defer func() { _ = trx.Rollback() }()
-	envId, err := h.checkZitiIdentity(params.Body.Identity, principal, tx)
+	env, err := str.FindEnvironmentForAccount(params.Body.Identity, int(principal.ID), trx)
 	if err != nil {
 		logrus.Errorf("identity check failed for user '%v': %v", principal.Email, err)
 		return environment.NewDisableUnauthorized()
 	}
-	env, err := str.GetEnvironment(envId, tx)
-	if err != nil {
-		logrus.Errorf("error getting environment for user '%v': %v", principal.Email, err)
-		return environment.NewDisableInternalServerError()
-	}
-	if env.Deleted {
-		logrus.Errorf("environment '%v' for user '%v' deleted", env.ZId, principal.Email)
-		return environment.NewDisableUnauthorized()
-	}
 	edge, err := zrokEdgeSdk.Client(cfg.Ziti)
 	if err != nil {
 		logrus.Errorf("error getting edge client for user '%v': %v", principal.Email, err)
 		return environment.NewDisableInternalServerError()
 	}
-	if err := h.removeSharesForEnvironment(envId, tx, edge); err != nil {
+	if err := h.removeSharesForEnvironment(env, trx, edge); err != nil {
 		logrus.Errorf("error removing shares for environment for user '%v': %v", principal.Email, err)
 		return environment.NewDisableInternalServerError()
 	}
-	if err := h.removeFrontendsForEnvironment(envId, tx, edge); err != nil {
+	if err := h.removeFrontendsForEnvironment(env, trx, edge); err != nil {
 		logrus.Errorf("error removing frontends for environment for user '%v': %v", principal.Email, err)
 		return environment.NewDisableInternalServerError()
 	}
-	if err := h.removeEnvironment(envId, tx); err != nil {
+	if err := h.removeAgentRemoteForEnvironment(env, trx, edge); err != nil {
-		logrus.Errorf("error removing environment for user '%v': %v", principal.Email, err)
+		logrus.Errorf("error removing agent remote for '%v' (%v): %v", env.ZId, principal.Email, err)
 		return environment.NewDisableInternalServerError()
 	}
 	if err := zrokEdgeSdk.DeleteEdgeRouterPolicy(env.ZId, edge); err != nil {
 		logrus.Errorf("error deleting edge router policy for user '%v': %v", principal.Email, err)
 		return environment.NewDisableInternalServerError()
 	}
-	if err := zrokEdgeSdk.DeleteIdentity(params.Body.Identity, edge); err != nil {
+	if err := zrokEdgeSdk.DeleteIdentity(env.ZId, edge); err != nil {
 		logrus.Errorf("error deleting identity for user '%v': %v", principal.Email, err)
 		return environment.NewDisableInternalServerError()
 	}
-	if err := tx.Commit(); err != nil {
+	if err := h.removeEnvironmentFromStore(env, trx); err != nil {
+		logrus.Errorf("error removing environment for user '%v': %v", principal.Email, err)
+		return environment.NewDisableInternalServerError()
+	}
+	if err := trx.Commit(); err != nil {
 		logrus.Errorf("error committing for user '%v': %v", principal.Email, err)
+		return environment.NewDisableInternalServerError()
 	}
 	return environment.NewDisableOK()
 }
 
-func (h *disableHandler) checkZitiIdentity(id string, principal *rest_model_zrok.Principal, tx *sqlx.Tx) (int, error) {
+func (h *disableHandler) removeSharesForEnvironment(env *store.Environment, trx *sqlx.Tx, edge *rest_management_api_client.ZitiEdgeManagement) error {
-	envs, err := str.FindEnvironmentsForAccount(int(principal.ID), tx)
+	shrs, err := str.FindSharesForEnvironment(env.Id, trx)
-	if err != nil {
-		return -1, err
-	}
-	for _, env := range envs {
-		if env.ZId == id {
-			return env.Id, nil
-		}
-	}
-	return -1, errors.Errorf("no such environment '%v'", id)
-}
-
-func (h *disableHandler) removeSharesForEnvironment(envId int, tx *sqlx.Tx, edge *rest_management_api_client.ZitiEdgeManagement) error {
-	env, err := str.GetEnvironment(envId, tx)
 	if err != nil {
 		return err
 	}
-	if !env.Deleted {
-		shrs, err := str.FindSharesForEnvironment(envId, tx)
-		if err != nil {
-			return err
-		}
-		for _, shr := range shrs {
-			shrToken := shr.Token
-			logrus.Infof("garbage collecting share '%v' for environment '%v'", shrToken, env.ZId)
-			if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicyForShare(env.ZId, shrToken, edge); err != nil {
-				logrus.Error(err)
-			}
-			if err := zrokEdgeSdk.DeleteServicePoliciesDialForShare(env.ZId, shrToken, edge); err != nil {
-				logrus.Error(err)
-			}
-			if err := zrokEdgeSdk.DeleteServicePoliciesBindForShare(env.ZId, shrToken, edge); err != nil {
-				logrus.Error(err)
-			}
-			if err := zrokEdgeSdk.DeleteConfig(env.ZId, shrToken, edge); err != nil {
-				logrus.Error(err)
-			}
-			if err := zrokEdgeSdk.DeleteService(env.ZId, shr.ZId, edge); err != nil {
-				logrus.Error(err)
-			}
-			logrus.Infof("removed share '%v' for environment '%v'", shr.Token, env.ZId)
-		}
-	}
-	return nil
-}
-
-func (h *disableHandler) removeFrontendsForEnvironment(envId int, tx *sqlx.Tx, edge *rest_management_api_client.ZitiEdgeManagement) error {
-	env, err := str.GetEnvironment(envId, tx)
-	if err != nil {
-		return err
-	}
-	if !env.Deleted {
-		fes, err := str.FindFrontendsForEnvironment(envId, tx)
-		if err != nil {
-			return err
-		}
-		for _, fe := range fes {
-			if err := zrokEdgeSdk.DeleteServicePolicies(env.ZId, fmt.Sprintf("tags.zrokFrontendToken=\"%v\" and type=1", fe.Token), edge); err != nil {
-				logrus.Errorf("error removing frontend access for '%v': %v", fe.Token, err)
-			}
-		}
-	}
-	return nil
-}
-
-func (h *disableHandler) removeEnvironment(envId int, tx *sqlx.Tx) error {
-	shrs, err := str.FindSharesForEnvironment(envId, tx)
-	if err != nil {
-		return errors.Wrapf(err, "error finding shares for environment '%d'", envId)
-	}
 	for _, shr := range shrs {
-		if err := str.DeleteShare(shr.Id, tx); err != nil {
+		shrToken := shr.Token
-			return errors.Wrapf(err, "error deleting share '%d' for environment '%d'", shr.Id, envId)
+		logrus.Infof("garbage collecting share '%v' for environment '%v'", shrToken, env.ZId)
+		if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicyForShare(env.ZId, shrToken, edge); err != nil {
+			logrus.Error(err)
 		}
-	}
+		if err := zrokEdgeSdk.DeleteServicePoliciesDialForShare(env.ZId, shrToken, edge); err != nil {
-	fes, err := str.FindFrontendsForEnvironment(envId, tx)
+			logrus.Error(err)
-	if err != nil {
-		return errors.Wrapf(err, "error finding frontends for environment '%d'", envId)
-	}
-	for _, fe := range fes {
-		if err := str.DeleteFrontend(fe.Id, tx); err != nil {
-			return errors.Wrapf(err, "error deleting frontend '%d' for environment '%d'", fe.Id, envId)
 		}
-	}
+		if err := zrokEdgeSdk.DeleteServicePoliciesBindForShare(env.ZId, shrToken, edge); err != nil {
-	if err := str.DeleteEnvironment(envId, tx); err != nil {
+			logrus.Error(err)
-		return errors.Wrapf(err, "error deleting environment '%d'", envId)
+		}
+		if err := zrokEdgeSdk.DeleteConfig(env.ZId, shrToken, edge); err != nil {
+			logrus.Error(err)
+		}
+		if err := zrokEdgeSdk.DeleteService(env.ZId, shr.ZId, edge); err != nil {
+			logrus.Error(err)
+		}
+		logrus.Infof("removed share '%v' for environment '%v'", shr.Token, env.ZId)
+	}
+	return nil
+}
+
+func (h *disableHandler) removeFrontendsForEnvironment(env *store.Environment, trx *sqlx.Tx, edge *rest_management_api_client.ZitiEdgeManagement) error {
+	fes, err := str.FindFrontendsForEnvironment(env.Id, trx)
+	if err != nil {
+		return err
+	}
+	for _, fe := range fes {
+		if err := zrokEdgeSdk.DeleteServicePolicies(env.ZId, fmt.Sprintf("tags.zrokFrontendToken=\"%v\" and type=1", fe.Token), edge); err != nil {
+			logrus.Errorf("error removing frontend access for '%v': %v", fe.Token, err)
+		}
+	}
+	return nil
+}
+
+func (h *disableHandler) removeAgentRemoteForEnvironment(env *store.Environment, trx *sqlx.Tx, edge *rest_management_api_client.ZitiEdgeManagement) error {
+	enrolled, err := str.IsAgentEnrolledForEnvironment(env.Id, trx)
+	if err != nil {
+		return err
+	}
+	if enrolled {
+		ae, err := str.FindAgentEnrollmentForEnvironment(env.Id, trx)
+		if err != nil {
+			return err
+		}
+		if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicyForAgentRemote(env.ZId, ae.Token, edge); err != nil {
+			return err
+		}
+		if err := zrokEdgeSdk.DeleteServicePoliciesDialForAgentRemote(env.ZId, ae.Token, edge); err != nil {
+			return err
+		}
+		if err := zrokEdgeSdk.DeleteServicePoliciesBindForAgentRemote(env.ZId, ae.Token, edge); err != nil {
+			return err
+		}
+		filter := fmt.Sprintf("name=\"%v\"", ae.Token)
+		limit := int64(1)
+		offset := int64(0)
+		listReq := &edge_service.ListServicesParams{
+			Filter:  &filter,
+			Limit:   &limit,
+			Offset:  &offset,
+			Context: context.Background(),
+		}
+		listReq.SetTimeout(30 * time.Second)
+		listResp, err := edge.Service.ListServices(listReq, nil)
+		if err != nil {
+			return err
+		}
+		aeZId := *(listResp.Payload.Data[0].ID)
+		if err := zrokEdgeSdk.DeleteService(env.ZId, aeZId, edge); err != nil {
+			return err
+		}
+		if err := str.DeleteAgentEnrollment(ae.Id, trx); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (h *disableHandler) removeEnvironmentFromStore(env *store.Environment, trx *sqlx.Tx) error {
+	shrs, err := str.FindSharesForEnvironment(env.Id, trx)
+	if err != nil {
+		return errors.Wrapf(err, "error finding shares for environment '%d'", env.Id)
+	}
+	for _, shr := range shrs {
+		if err := str.DeleteShare(shr.Id, trx); err != nil {
+			return errors.Wrapf(err, "error deleting share '%d' for environment '%d'", shr.Id, env.Id)
+		}
+	}
+	fes, err := str.FindFrontendsForEnvironment(env.Id, trx)
+	if err != nil {
+		return errors.Wrapf(err, "error finding frontends for environment '%d'", env.Id)
+	}
+	for _, fe := range fes {
+		if err := str.DeleteFrontend(fe.Id, trx); err != nil {
+			return errors.Wrapf(err, "error deleting frontend '%d' for environment '%d'", fe.Id, env.Id)
+		}
+	}
+	if err := str.DeleteEnvironment(env.Id, trx); err != nil {
+		return errors.Wrapf(err, "error deleting environment '%d'", env.Id)
 	}
 	return nil
 }

controller/enable.go

@@ -20,7 +20,6 @@ func newEnableHandler() *enableHandler {
 }
 
 func (h *enableHandler) Handle(params environment.EnableParams, principal *rest_model_zrok.Principal) middleware.Responder {
-	// start transaction early; if it fails, don't bother creating ziti resources
 	trx, err := str.Begin()
 	if err != nil {
 		logrus.Errorf("error starting transaction for user '%v': %v", principal.Email, err)

controller/store/agentEnrollment.go

@@ -23,6 +23,14 @@ func (str *Store) CreateAgentEnrollment(envId int, token string, trx *sqlx.Tx) (
 	return id, nil
 }
 
+func (str *Store) IsAgentEnrolledForEnvironment(envId int, trx *sqlx.Tx) (bool, error) {
+	count := 0
+	if err := trx.QueryRowx("select count(0) from agent_enrollments where environment_id = $1 and not deleted", envId).Scan(&count); err != nil {
+		return false, err
+	}
+	return count == 0, nil
+}
+
 func (str *Store) FindAgentEnrollmentForEnvironment(envId int, trx *sqlx.Tx) (*AgentEnrollment, error) {
 	ae := &AgentEnrollment{}
 	if err := trx.QueryRowx("select * from agent_enrollments where environment_id = $1 and not deleted", envId).StructScan(ae); err != nil {