From 79e9f484dca76ef79f4baba126ffce338a67f3c9 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Tue, 21 Mar 2023 16:18:17 -0400
Subject: [PATCH] centralized environment limits check (#277)

---
 controller/controller.go   | 11 ++++++-----
 controller/enable.go       | 30 ++++++++++++++----------------
 controller/limits/agent.go | 19 +++++++++++++++++++
 3 files changed, 39 insertions(+), 21 deletions(-)

diff --git a/controller/controller.go b/controller/controller.go
index f02ecd33..01c61f9e 100644
--- a/controller/controller.go
+++ b/controller/controller.go
@@ -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() }()
 		}
 	}
 
diff --git a/controller/enable.go b/controller/enable.go
index 1025c562..cf275d80 100644
--- a/controller/enable.go
+++ b/controller/enable.go
@@ -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
diff --git a/controller/limits/agent.go b/controller/limits/agent.go
index 5c144473..01d07cd8 100644
--- a/controller/limits/agent.go
+++ b/controller/limits/agent.go
@@ -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