From b2f051834e157afa70a4deb5956c576e2ae3dcbb Mon Sep 17 00:00:00 2001 From: Michael Quigley Date: Fri, 30 May 2025 15:52:44 -0400 Subject: [PATCH] fledged enroll/unenroll handlers for agent remoting (#967) --- controller/agentEnroll.go | 5 ++ controller/agentUnenroll.go | 108 ++++++++++++++++++++++++++++ controller/controller.go | 1 + controller/disable.go | 6 +- controller/gc.go | 8 +-- controller/limits/limitAction.go | 2 +- controller/store/agentEnrollment.go | 14 +++- controller/unshare.go | 6 +- controller/zrokEdgeSdk/serp.go | 11 ++- controller/zrokEdgeSdk/sp.go | 12 +++- 10 files changed, 158 insertions(+), 15 deletions(-) create mode 100644 controller/agentUnenroll.go diff --git a/controller/agentEnroll.go b/controller/agentEnroll.go index 525dc248..54e0db82 100644 --- a/controller/agentEnroll.go +++ b/controller/agentEnroll.go @@ -68,6 +68,11 @@ func (h *agentEnrollHandler) Handle(params agent.EnrollParams, principal *rest_m return agent.NewEnrollInternalServerError() } + if _, err := str.CreateAgentEnrollment(env.Id, token, trx); err != nil { + logrus.Errorf("error storing agent enrollment for '%v' (%v): %v", env.ZId, principal.Email, err) + return agent.NewEnrollInternalServerError() + } + if err := trx.Commit(); err != nil { logrus.Errorf("error committing agent enrollment record for '%v' (%v): %v", env.ZId, principal.Email, err) return agent.NewEnrollInternalServerError() diff --git a/controller/agentUnenroll.go b/controller/agentUnenroll.go new file mode 100644 index 00000000..e77448a1 --- /dev/null +++ b/controller/agentUnenroll.go @@ -0,0 +1,108 @@ +package controller + +import ( + "context" + "fmt" + "github.com/go-openapi/runtime/middleware" + "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/zrokEdgeSdk" + "github.com/openziti/zrok/rest_model_zrok" + "github.com/openziti/zrok/rest_server_zrok/operations/agent" + "github.com/pkg/errors" + "github.com/sirupsen/logrus" + "time" +) + +type agentUnenrollHandler struct{} + +func newAgentUnenrollHandler() *agentUnenrollHandler { + return &agentUnenrollHandler{} +} + +func (h *agentUnenrollHandler) Handle(params agent.UnenrollParams, 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 '%v': %v", principal.Email, err) + return agent.NewUnenrollInternalServerError() + } + defer trx.Rollback() + + env, err := str.FindEnvironmentForAccount(params.Body.EnvZID, int(principal.ID), trx) + if err != nil { + logrus.Errorf("error finding environment '%v' for '%v': %v", params.Body.EnvZID, principal.Email, err) + return agent.NewUnenrollUnauthorized() + } + + ae, err := str.FindAgentEnrollmentForEnvironment(env.Id, trx) + if err != nil { + logrus.Errorf("error finding agent enrollment for '%v' (%v): %v", params.Body.EnvZID, principal.Email, err) + return agent.NewUnenrollBadRequest() + } + + client, err := zrokEdgeSdk.Client(cfg.Ziti) + if err != nil { + logrus.Errorf("error getting ziti client for '%v': %v", principal.Email, err) + return agent.NewUnenrollInternalServerError() + } + + if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicyForAgentRemote(env.ZId, ae.Token, client); err != nil { + logrus.Errorf("error removing agent remote serp for '%v' (%v): %v", env.ZId, principal.Email, err) + return agent.NewUnenrollInternalServerError() + } + + if err := zrokEdgeSdk.DeleteServicePoliciesDialForAgentRemote(env.ZId, ae.Token, client); err != nil { + logrus.Errorf("error removing agent remote dial service policy for '%v' (%v): %v", env.ZId, principal.Email, err) + return agent.NewUnenrollInternalServerError() + } + + if err := zrokEdgeSdk.DeleteServicePoliciesBindForAgentRemote(env.ZId, ae.Token, client); err != nil { + logrus.Errorf("error removing agent remote bind service policy for '%v' (%v): %v", env.ZId, principal.Email, err) + return agent.NewUnenrollInternalServerError() + } + + aeZId, err := h.findAgentRemoteZId(ae.Token, client) + if err != nil { + logrus.Errorf("error finding zId for agent remote service for '%v' (%v): %v", env.ZId, principal.Email, err) + return agent.NewUnenrollInternalServerError() + } + + if err := zrokEdgeSdk.DeleteService(env.ZId, aeZId, client); err != nil { + logrus.Errorf("error removing agent remote service for '%v' (%v): %v", env.ZId, principal.Email, err) + return agent.NewUnenrollInternalServerError() + } + + if err := str.DeleteAgentEnrollment(ae.Id, trx); err != nil { + logrus.Errorf("error deleting agent enrollment for '%v' (%v): %v", env.ZId, principal.Email, err) + return agent.NewUnenrollInternalServerError() + } + + if err := trx.Commit(); err != nil { + logrus.Errorf("error committing agent unenrollment for '%v' (%v): %v", env.ZId, principal.Email, err) + return agent.NewUnenrollInternalServerError() + } + + return agent.NewUnenrollOK() +} + +func (h *agentUnenrollHandler) findAgentRemoteZId(enrollmentToken string, edge *rest_management_api_client.ZitiEdgeManagement) (string, error) { + filter := fmt.Sprintf("name=\"%v\"", enrollmentToken) + 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 + } + if len(listResp.Payload.Data) == 1 { + return *(listResp.Payload.Data[0].ID), nil + } + return "", errors.Errorf("agent remote service '%v' not found", enrollmentToken) +} diff --git a/controller/controller.go b/controller/controller.go index a2d8a25a..a9ea45e0 100644 --- a/controller/controller.go +++ b/controller/controller.go @@ -68,6 +68,7 @@ func Run(inCfg *config.Config) error { if cfg.AgentController != nil { api.AgentEnrollHandler = newAgentEnrollHandler() api.AgentPingHandler = newAgentPingHandler() + api.AgentUnenrollHandler = newAgentUnenrollHandler() } api.EnvironmentEnableHandler = newEnableHandler() api.EnvironmentDisableHandler = newDisableHandler() diff --git a/controller/disable.go b/controller/disable.go index 6b16f5f1..d2d48835 100644 --- a/controller/disable.go +++ b/controller/disable.go @@ -97,13 +97,13 @@ func (h *disableHandler) removeSharesForEnvironment(envId int, tx *sqlx.Tx, edge for _, shr := range shrs { shrToken := shr.Token logrus.Infof("garbage collecting share '%v' for environment '%v'", shrToken, env.ZId) - if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicy(env.ZId, shrToken, edge); err != nil { + if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicyForShare(env.ZId, shrToken, edge); err != nil { logrus.Error(err) } - if err := zrokEdgeSdk.DeleteServicePoliciesDial(env.ZId, shrToken, edge); err != nil { + if err := zrokEdgeSdk.DeleteServicePoliciesDialForShare(env.ZId, shrToken, edge); err != nil { logrus.Error(err) } - if err := zrokEdgeSdk.DeleteServicePoliciesBind(env.ZId, shrToken, edge); err != nil { + if err := zrokEdgeSdk.DeleteServicePoliciesBindForShare(env.ZId, shrToken, edge); err != nil { logrus.Error(err) } if err := zrokEdgeSdk.DeleteConfig(env.ZId, shrToken, edge); err != nil { diff --git a/controller/gc.go b/controller/gc.go index c4682f4a..820e533a 100644 --- a/controller/gc.go +++ b/controller/gc.go @@ -73,13 +73,13 @@ func gcServices(edge *rest_management_api_client.ZitiEdgeManagement, liveMap map for _, svc := range listResp.Payload.Data { if _, found := liveMap[*svc.Name]; !found { logrus.Infof("garbage collecting, zitiSvcId='%v', zrokSvcId='%v'", *svc.ID, *svc.Name) - if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicy("gc", *svc.Name, edge); err != nil { + if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicyForShare("gc", *svc.Name, edge); err != nil { logrus.Errorf("error garbage collecting service edge router policy: %v", err) } - if err := zrokEdgeSdk.DeleteServicePoliciesDial("gc", *svc.Name, edge); err != nil { + if err := zrokEdgeSdk.DeleteServicePoliciesDialForShare("gc", *svc.Name, edge); err != nil { logrus.Errorf("error garbage collecting service dial policy: %v", err) } - if err := zrokEdgeSdk.DeleteServicePoliciesBind("gc", *svc.Name, edge); err != nil { + if err := zrokEdgeSdk.DeleteServicePoliciesBindForShare("gc", *svc.Name, edge); err != nil { logrus.Errorf("error garbage collecting service bind policy: %v", err) } if err := zrokEdgeSdk.DeleteConfig("gc", *svc.Name, edge); err != nil { @@ -110,7 +110,7 @@ func gcServiceEdgeRouterPolicies(edge *rest_management_api_client.ZitiEdgeManage for _, serp := range listResp.Payload.Data { if _, found := liveMap[*serp.Name]; !found { logrus.Infof("garbage collecting, svcId='%v'", *serp.Name) - if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicy("gc", *serp.Name, edge); err != nil { + if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicyForShare("gc", *serp.Name, edge); err != nil { logrus.Errorf("error garbage collecting service edge router policy: %v", err) } } else { diff --git a/controller/limits/limitAction.go b/controller/limits/limitAction.go index f0392239..bce18573 100644 --- a/controller/limits/limitAction.go +++ b/controller/limits/limitAction.go @@ -38,7 +38,7 @@ func (a *limitAction) HandleAccount(acct *store.Account, _, _ int64, bwc store.B for _, shr := range shrs { if _, ignore := ignoreBackends[sdk.BackendMode(shr.BackendMode)]; !ignore { - if err := zrokEdgeSdk.DeleteServicePoliciesDial(env.ZId, shr.Token, edge); err != nil { + if err := zrokEdgeSdk.DeleteServicePoliciesDialForShare(env.ZId, shr.Token, edge); err != nil { return errors.Wrapf(err, "error deleting dial service policy for '%v'", shr.Token) } logrus.Infof("removed dial service policy for share '%v' of environment '%v'", shr.Token, env.ZId) diff --git a/controller/store/agentEnrollment.go b/controller/store/agentEnrollment.go index 99d5ce22..a49d0a24 100644 --- a/controller/store/agentEnrollment.go +++ b/controller/store/agentEnrollment.go @@ -12,7 +12,7 @@ type AgentEnrollment struct { } func (str *Store) CreateAgentEnrollment(envId int, token string, trx *sqlx.Tx) (int, error) { - stmt, err := trx.Prepare("insert into agent_enrollments (environment_id, token) values ($1, $2) returing id") + stmt, err := trx.Prepare("insert into agent_enrollments (environment_id, token) values ($1, $2) returning id") if err != nil { return 0, errors.Wrap(err, "error preparing agent enrollments insert statement") } @@ -30,3 +30,15 @@ func (str *Store) FindAgentEnrollmentForEnvironment(envId int, trx *sqlx.Tx) (*A } return ae, nil } + +func (str *Store) DeleteAgentEnrollment(id int, trx *sqlx.Tx) error { + stmt, err := trx.Prepare("update agent_enrollments set updated_at = current_timestamp, deleted = true where id = $1") + if err != nil { + return errors.Wrap(err, "error preparing agent enrollments delete statement") + } + _, err = stmt.Exec(id) + if err != nil { + return errors.Wrap(err, "error executing agent enrollments delete statement") + } + return nil +} diff --git a/controller/unshare.go b/controller/unshare.go index a2633437..f120243b 100644 --- a/controller/unshare.go +++ b/controller/unshare.go @@ -121,13 +121,13 @@ func (h *unshareHandler) findShareZId(shrToken string, edge *rest_management_api } func (h *unshareHandler) deallocateResources(senv *store.Environment, shrToken, shrZId string, edge *rest_management_api_client.ZitiEdgeManagement) { - if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicy(senv.ZId, shrToken, edge); err != nil { + if err := zrokEdgeSdk.DeleteServiceEdgeRouterPolicyForShare(senv.ZId, shrToken, edge); err != nil { logrus.Warnf("error deleting service edge router policies for share '%v' in environment '%v': %v", shrToken, senv.ZId, err) } - if err := zrokEdgeSdk.DeleteServicePoliciesDial(senv.ZId, shrToken, edge); err != nil { + if err := zrokEdgeSdk.DeleteServicePoliciesDialForShare(senv.ZId, shrToken, edge); err != nil { logrus.Warnf("error deleting dial service policies for share '%v' in environment '%v': %v", shrToken, senv.ZId, err) } - if err := zrokEdgeSdk.DeleteServicePoliciesBind(senv.ZId, shrToken, edge); err != nil { + if err := zrokEdgeSdk.DeleteServicePoliciesBindForShare(senv.ZId, shrToken, edge); err != nil { logrus.Warnf("error deleting bind service policies for share '%v' in environment '%v': %v", shrToken, senv.ZId, err) } if err := zrokEdgeSdk.DeleteConfig(senv.ZId, shrToken, edge); err != nil { diff --git a/controller/zrokEdgeSdk/serp.go b/controller/zrokEdgeSdk/serp.go index 85712e4c..d3f12584 100644 --- a/controller/zrokEdgeSdk/serp.go +++ b/controller/zrokEdgeSdk/serp.go @@ -47,8 +47,17 @@ func CreateServiceEdgeRouterPolicy(name, shrZId string, moreTags map[string]inte return resp.Payload.Data.ID, nil } -func DeleteServiceEdgeRouterPolicy(envZId, shrToken string, edge *rest_management_api_client.ZitiEdgeManagement) error { +func DeleteServiceEdgeRouterPolicyForAgentRemote(envZId, enrollmentToken string, edge *rest_management_api_client.ZitiEdgeManagement) error { + filter := fmt.Sprintf("tags.zrokAgentRemote=\"%v\"", enrollmentToken) + return DeleteServiceEdgeRouterPolicy(envZId, filter, edge) +} + +func DeleteServiceEdgeRouterPolicyForShare(envZId, shrToken string, edge *rest_management_api_client.ZitiEdgeManagement) error { filter := fmt.Sprintf("tags.zrokShareToken=\"%v\"", shrToken) + return DeleteServiceEdgeRouterPolicy(envZId, filter, edge) +} + +func DeleteServiceEdgeRouterPolicy(envZId, filter string, edge *rest_management_api_client.ZitiEdgeManagement) error { limit := int64(1) offset := int64(0) listReq := &service_edge_router_policy.ListServiceEdgeRouterPoliciesParams{ diff --git a/controller/zrokEdgeSdk/sp.go b/controller/zrokEdgeSdk/sp.go index e6a49bac..4df7e3d4 100644 --- a/controller/zrokEdgeSdk/sp.go +++ b/controller/zrokEdgeSdk/sp.go @@ -78,11 +78,19 @@ func createServicePolicy(name string, semantic rest_model.Semantic, identityRole return resp.Payload.Data.ID, nil } -func DeleteServicePoliciesBind(envZId, shrToken string, edge *rest_management_api_client.ZitiEdgeManagement) error { +func DeleteServicePoliciesBindForAgentRemote(envZId, enrollmentToken string, edge *rest_management_api_client.ZitiEdgeManagement) error { + return DeleteServicePolicies(envZId, fmt.Sprintf("tags.zrokAgentRemote=\"%v\" and type=%d", enrollmentToken, servicePolicyBind), edge) +} + +func DeleteServicePoliciesDialForAgentRemote(envZId, enrollmentToken string, edge *rest_management_api_client.ZitiEdgeManagement) error { + return DeleteServicePolicies(envZId, fmt.Sprintf("tags.zrokAgentRemote=\"%v\" and type=%d", enrollmentToken, servicePolicyDial), edge) +} + +func DeleteServicePoliciesBindForShare(envZId, shrToken string, edge *rest_management_api_client.ZitiEdgeManagement) error { return DeleteServicePolicies(envZId, fmt.Sprintf("tags.zrokShareToken=\"%v\" and type=%d", shrToken, servicePolicyBind), edge) } -func DeleteServicePoliciesDial(envZId, shrToken string, edge *rest_management_api_client.ZitiEdgeManagement) error { +func DeleteServicePoliciesDialForShare(envZId, shrToken string, edge *rest_management_api_client.ZitiEdgeManagement) error { return DeleteServicePolicies(envZId, fmt.Sprintf("tags.zrokShareToken=\"%v\" and type=%d", shrToken, servicePolicyDial), edge) }