From a6c2841cf2e3885ca3e8535d525518bf25ea98d0 Mon Sep 17 00:00:00 2001 From: Michael Quigley Date: Mon, 27 Mar 2023 14:01:31 -0400 Subject: [PATCH] environment limit and relax actions (#276) --- controller/limits/environmentLimitAction.go | 19 ++++++++++-- controller/limits/environmentRelaxAction.go | 34 +++++++++++++++++++-- controller/limits/shareLimitAction.go | 10 +++--- controller/limits/shareRelaxAction.go | 20 ++++++------ 4 files changed, 64 insertions(+), 19 deletions(-) diff --git a/controller/limits/environmentLimitAction.go b/controller/limits/environmentLimitAction.go index 79c78517..ce26cafc 100644 --- a/controller/limits/environmentLimitAction.go +++ b/controller/limits/environmentLimitAction.go @@ -4,6 +4,8 @@ import ( "github.com/jmoiron/sqlx" "github.com/openziti/edge/rest_management_api_client" "github.com/openziti/zrok/controller/store" + "github.com/openziti/zrok/controller/zrokEdgeSdk" + "github.com/pkg/errors" "github.com/sirupsen/logrus" ) @@ -16,7 +18,20 @@ func newEnvironmentLimitAction(str *store.Store, edge *rest_management_api_clien return &environmentLimitAction{str, edge} } -func (a *environmentLimitAction) HandleEnvironment(e *store.Environment, rxBytes, txBytes int64, limit *BandwidthPerPeriod, trx *sqlx.Tx) error { - logrus.Infof("limiting '%v'", e.ZId) +func (a *environmentLimitAction) HandleEnvironment(env *store.Environment, _, _ int64, _ *BandwidthPerPeriod, trx *sqlx.Tx) error { + logrus.Infof("limiting '%v'", env.ZId) + + shrs, err := a.str.FindSharesForEnvironment(env.Id, trx) + if err != nil { + return errors.Wrapf(err, "error finding shares for environment '%v'", env.ZId) + } + + for _, shr := range shrs { + if err := zrokEdgeSdk.DeleteServicePolicyDial(env.ZId, shr.Token, a.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) + } + return nil } diff --git a/controller/limits/environmentRelaxAction.go b/controller/limits/environmentRelaxAction.go index 0721b508..aa79ade3 100644 --- a/controller/limits/environmentRelaxAction.go +++ b/controller/limits/environmentRelaxAction.go @@ -4,6 +4,8 @@ import ( "github.com/jmoiron/sqlx" "github.com/openziti/edge/rest_management_api_client" "github.com/openziti/zrok/controller/store" + "github.com/openziti/zrok/controller/zrokEdgeSdk" + "github.com/pkg/errors" "github.com/sirupsen/logrus" ) @@ -16,7 +18,35 @@ func newEnvironmentRelaxAction(str *store.Store, edge *rest_management_api_clien return &environmentRelaxAction{str, edge} } -func (a *environmentRelaxAction) HandleEnvironment(e *store.Environment, rxBytes, txBytes int64, limit *BandwidthPerPeriod, trx *sqlx.Tx) error { - logrus.Infof("relaxing '%v'", e.ZId) +func (a *environmentRelaxAction) HandleEnvironment(env *store.Environment, rxBytes, txBytes int64, limit *BandwidthPerPeriod, trx *sqlx.Tx) error { + logrus.Infof("relaxing '%v'", env.ZId) + + shrs, err := a.str.FindSharesForEnvironment(env.Id, trx) + if err != nil { + return errors.Wrapf(err, "error finding shares for environment '%v'", env.ZId) + } + + for _, shr := range shrs { + if shr.ShareMode == "public" { + env, err := a.str.GetEnvironment(shr.EnvironmentId, trx) + if err != nil { + return errors.Wrap(err, "error finding environment") + } + + fe, err := a.str.FindFrontendPubliclyNamed(*shr.FrontendSelection, trx) + if err != nil { + return errors.Wrapf(err, "error finding frontend name '%v' for '%v'", *shr.FrontendSelection, shr.Token) + } + + if err := zrokEdgeSdk.CreateServicePolicyDial(env.ZId+"-"+shr.ZId+"-dial", shr.ZId, []string{fe.ZId}, zrokEdgeSdk.ZrokShareTags(shr.Token).SubTags, a.edge); err != nil { + return errors.Wrapf(err, "error creating dial service policy for '%v'", shr.Token) + } + logrus.Infof("added dial service policy for '%v'", shr.Token) + + } else if shr.ShareMode == "private" { + return errors.New("share relax for private share not implemented") + } + } + return nil } diff --git a/controller/limits/shareLimitAction.go b/controller/limits/shareLimitAction.go index 4daf2551..7eefffb2 100644 --- a/controller/limits/shareLimitAction.go +++ b/controller/limits/shareLimitAction.go @@ -17,18 +17,18 @@ func newShareLimitAction(str *store.Store, edge *rest_management_api_client.Ziti return &shareLimitAction{str, edge} } -func (a *shareLimitAction) HandleShare(s *store.Share, _, _ int64, _ *BandwidthPerPeriod, trx *sqlx.Tx) error { - logrus.Infof("limiting '%v'", s.Token) +func (a *shareLimitAction) HandleShare(shr *store.Share, _, _ int64, _ *BandwidthPerPeriod, trx *sqlx.Tx) error { + logrus.Infof("limiting '%v'", shr.Token) - env, err := a.str.GetEnvironment(s.EnvironmentId, trx) + env, err := a.str.GetEnvironment(shr.EnvironmentId, trx) if err != nil { return err } - if err := zrokEdgeSdk.DeleteServicePolicyDial(env.ZId, s.Token, a.edge); err != nil { + if err := zrokEdgeSdk.DeleteServicePolicyDial(env.ZId, shr.Token, a.edge); err != nil { return err } - logrus.Infof("removed service dial policy for '%v'", s.Token) + logrus.Infof("removed dial service policy for '%v'", shr.Token) return nil } diff --git a/controller/limits/shareRelaxAction.go b/controller/limits/shareRelaxAction.go index a13c1a7a..e702cf01 100644 --- a/controller/limits/shareRelaxAction.go +++ b/controller/limits/shareRelaxAction.go @@ -18,26 +18,26 @@ func newShareRelaxAction(str *store.Store, edge *rest_management_api_client.Ziti return &shareRelaxAction{str, edge} } -func (a *shareRelaxAction) HandleShare(s *store.Share, rxBytes, txBytes int64, limit *BandwidthPerPeriod, trx *sqlx.Tx) error { - logrus.Infof("relaxing '%v'", s.Token) +func (a *shareRelaxAction) HandleShare(shr *store.Share, rxBytes, txBytes int64, limit *BandwidthPerPeriod, trx *sqlx.Tx) error { + logrus.Infof("relaxing '%v'", shr.Token) - if s.ShareMode == "public" { - env, err := a.str.GetEnvironment(s.EnvironmentId, trx) + if shr.ShareMode == "public" { + env, err := a.str.GetEnvironment(shr.EnvironmentId, trx) if err != nil { return errors.Wrap(err, "error finding environment") } - fe, err := a.str.FindFrontendPubliclyNamed(*s.FrontendSelection, trx) + fe, err := a.str.FindFrontendPubliclyNamed(*shr.FrontendSelection, trx) if err != nil { - return errors.Wrapf(err, "error finding frontend name '%v' for '%v'", *s.FrontendSelection, s.Token) + return errors.Wrapf(err, "error finding frontend name '%v' for '%v'", *shr.FrontendSelection, shr.Token) } - if err := zrokEdgeSdk.CreateServicePolicyDial(env.ZId+"-"+s.ZId+"-dial", s.ZId, []string{fe.ZId}, zrokEdgeSdk.ZrokShareTags(s.Token).SubTags, a.edge); err != nil { - return errors.Wrapf(err, "error creating dial service policy for '%v'", s.Token) + if err := zrokEdgeSdk.CreateServicePolicyDial(env.ZId+"-"+shr.ZId+"-dial", shr.ZId, []string{fe.ZId}, zrokEdgeSdk.ZrokShareTags(shr.Token).SubTags, a.edge); err != nil { + return errors.Wrapf(err, "error creating dial service policy for '%v'", shr.Token) } - logrus.Infof("added dial service policy for '%v'", s.Token) + logrus.Infof("added dial service policy for '%v'", shr.Token) - } else if s.ShareMode == "private" { + } else if shr.ShareMode == "private" { return errors.New("share relax for private share not implemented") }