From f9dc0f6ba17e3ae555e4c703e70ed85f52baec20 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Wed, 29 Mar 2023 13:29:12 -0400
Subject: [PATCH] check limit journals for creating environments or shares
 (#281)

---
 controller/limits/agent.go | 68 +++++++++++++++++++++++++++++---------
 controller/share.go        |  6 ++--
 2 files changed, 55 insertions(+), 19 deletions(-)

diff --git a/controller/limits/agent.go b/controller/limits/agent.go
index fd58a8f4..9ec418a3 100644
--- a/controller/limits/agent.go
+++ b/controller/limits/agent.go
@@ -69,35 +69,71 @@ func (a *Agent) Stop() {
 }
 
 func (a *Agent) CanCreateEnvironment(acctId int, trx *sqlx.Tx) (bool, error) {
-	if a.cfg.Enforcing && a.cfg.Environments > Unlimited {
-		envs, err := a.str.FindEnvironmentsForAccount(acctId, trx)
+	if a.cfg.Enforcing {
+		alj, err := a.str.FindLatestAccountLimitJournal(acctId, trx)
 		if err != nil {
 			return false, err
 		}
-		if len(envs)+1 > a.cfg.Environments {
+		if alj.Action == store.LimitAction {
 			return false, nil
 		}
+
+		if a.cfg.Environments > Unlimited {
+			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) CanCreateShare(acctId int, trx *sqlx.Tx) (bool, error) {
-	if a.cfg.Enforcing && a.cfg.Shares > Unlimited {
-		envs, err := a.str.FindEnvironmentsForAccount(acctId, trx)
-		if err != nil {
-			return false, err
-		}
-		total := 0
-		for i := range envs {
-			shrs, err := a.str.FindSharesForEnvironment(envs[i].Id, trx)
+func (a *Agent) CanCreateShare(acctId, envId int, trx *sqlx.Tx) (bool, error) {
+	if a.cfg.Enforcing {
+		if empty, err := a.str.IsAccountLimitJournalEmpty(acctId, trx); err == nil && !empty {
+			alj, err := a.str.FindLatestAccountLimitJournal(acctId, trx)
 			if err != nil {
-				return false, errors.Wrapf(err, "unable to find shares for environment '%v'", envs[i].ZId)
+				return false, err
 			}
-			total += len(shrs)
-			if total+1 > a.cfg.Shares {
+			if alj.Action == store.LimitAction {
 				return false, nil
 			}
-			logrus.Infof("total = %d", total)
+		} else if err != nil {
+			return false, err
+		}
+
+		if empty, err := a.str.IsEnvironmentLimitJournalEmpty(envId, trx); err == nil && !empty {
+			elj, err := a.str.FindLatestEnvironmentLimitJournal(envId, trx)
+			if err != nil {
+				return false, err
+			}
+			if elj.Action == store.LimitAction {
+				return false, nil
+			}
+		} else if err != nil {
+			return false, err
+		}
+
+		if a.cfg.Shares > Unlimited {
+			envs, err := a.str.FindEnvironmentsForAccount(acctId, trx)
+			if err != nil {
+				return false, err
+			}
+			total := 0
+			for i := range envs {
+				shrs, err := a.str.FindSharesForEnvironment(envs[i].Id, trx)
+				if err != nil {
+					return false, errors.Wrapf(err, "unable to find shares for environment '%v'", envs[i].ZId)
+				}
+				total += len(shrs)
+				if total+1 > a.cfg.Shares {
+					return false, nil
+				}
+				logrus.Infof("total = %d", total)
+			}
 		}
 	}
 	return true, nil
diff --git a/controller/share.go b/controller/share.go
index 4b5b5b1b..bdfc6964 100644
--- a/controller/share.go
+++ b/controller/share.go
@@ -47,7 +47,7 @@ func (h *shareHandler) Handle(params share.ShareParams, principal *rest_model_zr
 		return share.NewShareInternalServerError()
 	}
 
-	if err := h.checkLimits(principal, trx); err != nil {
+	if err := h.checkLimits(envId, principal, trx); err != nil {
 		logrus.Errorf("limits error: %v", err)
 		return share.NewShareUnauthorized()
 	}
@@ -142,10 +142,10 @@ func (h *shareHandler) Handle(params share.ShareParams, principal *rest_model_zr
 	})
 }
 
-func (h *shareHandler) checkLimits(principal *rest_model_zrok.Principal, trx *sqlx.Tx) error {
+func (h *shareHandler) checkLimits(envId int, principal *rest_model_zrok.Principal, trx *sqlx.Tx) error {
 	if !principal.Limitless {
 		if limitsAgent != nil {
-			ok, err := limitsAgent.CanCreateShare(int(principal.ID), trx)
+			ok, err := limitsAgent.CanCreateShare(int(principal.ID), envId, trx)
 			if err != nil {
 				return errors.Wrapf(err, "error checking share limits for '%v'", principal.Email)
 			}