extern/zrok
1
1
Fork 0
mirror of https://github.com/openziti/zrok.git synced 2024-11-07 08:44:14 +01:00
Code Issues Packages Projects Releases Wiki Activity

centralized environment limits check (#277)

Browse Source
This commit is contained in:
Michael Quigley 2023-03-21 16:18:17 -04:00
parent 2ec228c496
commit 79e9f484dc
No known key found for this signature in database
GPG Key ID: 9B60314A9DD20A62
3 changed files with 39 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
controller/controller.go
View File

@ -20,6 +20,7 @@ import (
var cfg *config.Config
var str *store.Store
var idb influxdb2.Client
var limitsAgent *limits.Agent
func Run(inCfg *config.Config) error {
cfg = inCfg
@ -43,7 +44,7 @@ func Run(inCfg *config.Config) error {
api.AdminInviteTokenGenerateHandler = newInviteTokenGenerateHandler()
api.AdminListFrontendsHandler = newListFrontendsHandler()
api.AdminUpdateFrontendHandler = newUpdateFrontendHandler()
api.EnvironmentEnableHandler = newEnableHandler(cfg.Limits)
api.EnvironmentEnableHandler = newEnableHandler()
api.EnvironmentDisableHandler = newDisableHandler()
api.MetadataConfigurationHandler = newConfigurationHandler(cfg)
api.MetadataGetEnvironmentDetailHandler = newEnvironmentDetailHandler()
@ -83,13 +84,13 @@ func Run(inCfg *config.Config) error {
defer func() { ma.Stop() }()
if cfg.Limits != nil && cfg.Limits.Enforcing {
la, err := limits.NewAgent(cfg.Limits, cfg.Metrics.Influx, cfg.Ziti, str)
limitsAgent, err = limits.NewAgent(cfg.Limits, cfg.Metrics.Influx, cfg.Ziti, str)
if err != nil {
return errors.Wrap(err, "error creating limits agent")
}
ma.AddUsageSink(la)
la.Start()
defer func() { la.Stop() }()
ma.AddUsageSink(limitsAgent)
limitsAgent.Start()
defer func() { limitsAgent.Stop() }()
}
}

30
controller/enable.go
View File

@ -4,8 +4,6 @@ import (
"bytes"
"encoding/json"
"github.com/go-openapi/runtime/middleware"
"github.com/jmoiron/sqlx"
"github.com/openziti/zrok/controller/limits"
"github.com/openziti/zrok/controller/store"
"github.com/openziti/zrok/controller/zrokEdgeSdk"
"github.com/openziti/zrok/rest_model_zrok"
@ -14,12 +12,10 @@ import (
"github.com/sirupsen/logrus"
)
type enableHandler struct {
cfg *limits.Config
}
type enableHandler struct{}
func newEnableHandler(cfg *limits.Config) *enableHandler {
return &enableHandler{cfg: cfg}
func newEnableHandler() *enableHandler {
return &enableHandler{}
}
func (h *enableHandler) Handle(params environment.EnableParams, principal *rest_model_zrok.Principal) middleware.Responder {
@ -31,7 +27,7 @@ func (h *enableHandler) Handle(params environment.EnableParams, principal *rest_
}
defer func() { _ = tx.Rollback() }()
if err := h.checkLimits(principal, tx); err != nil {
if err := h.checkLimits(principal); err != nil {
logrus.Errorf("limits error for user '%v': %v", principal.Email, err)
return environment.NewEnableUnauthorized()
}
@ -100,14 +96,16 @@ func (h *enableHandler) Handle(params environment.EnableParams, principal *rest_
return resp
}
func (h *enableHandler) checkLimits(principal *rest_model_zrok.Principal, tx *sqlx.Tx) error {
if !principal.Limitless && h.cfg.Environments > limits.Unlimited {
envs, err := str.FindEnvironmentsForAccount(int(principal.ID), tx)
if err != nil {
return errors.Errorf("unable to find environments for account '%v': %v", principal.Email, err)
}
if len(envs)+1 > h.cfg.Environments {
return errors.Errorf("would exceed environments limit of %d for '%v'", h.cfg.Environments, principal.Email)
func (h *enableHandler) checkLimits(principal *rest_model_zrok.Principal) error {
if !principal.Limitless {
if limitsAgent != nil {
ok, err := limitsAgent.CanCreateEnvironment(int(principal.ID))
if err != nil {
return errors.Wrapf(err, "error checking limits for '%v'", principal.Email)
}
if !ok {
return errors.Wrapf(err, "environment limit check failed for '%v'", principal.Email)
}
}
}
return nil

19
controller/limits/agent.go
View File

@ -43,6 +43,25 @@ func (a *Agent) Stop() {
<-a.join
}
func (a *Agent) CanCreateEnvironment(acctId int) (bool, error) {
if a.cfg.Environments > Unlimited {
trx, err := a.str.Begin()
if err != nil {
return false, errors.Wrap(err, "error creating transaction")
}
defer func() { _ = trx.Rollback() }()
envs, err := a.str.FindEnvironmentsForAccount(acctId, trx)
if err != nil {
return false, err
}
if len(envs)+1 > a.cfg.Environments {
return false, nil
}
}
return true, nil
}
func (a *Agent) Handle(u *metrics.Usage) error {
logrus.Debugf("handling: %v", u)
a.queue <- u