wire in checks for reserved share count and unique name count (#632)

This commit is contained in:
Michael Quigley 2024-05-22 16:21:52 -04:00
parent 9b03765086
commit 33fdce7337
No known key found for this signature in database
GPG Key ID: 9B60314A9DD20A62
2 changed files with 21 additions and 5 deletions

View File

@ -94,7 +94,7 @@ func (a *Agent) CanCreateEnvironment(acctId int, trx *sqlx.Tx) (bool, error) {
return true, nil return true, nil
} }
func (a *Agent) CanCreateShare(acctId, envId int, trx *sqlx.Tx) (bool, error) { func (a *Agent) CanCreateShare(acctId, envId int, reserved, uniqueName bool, trx *sqlx.Tx) (bool, error) {
if a.cfg.Enforcing { if a.cfg.Enforcing {
if err := a.str.LimitCheckLock(acctId, trx); err != nil { if err := a.str.LimitCheckLock(acctId, trx); err != nil {
return false, err return false, err
@ -123,21 +123,37 @@ func (a *Agent) CanCreateShare(acctId, envId int, trx *sqlx.Tx) (bool, error) {
return false, err return false, err
} }
if a.cfg.Shares > Unlimited { if a.cfg.Shares > Unlimited || (reserved && a.cfg.ReservedShares > Unlimited) || (reserved && uniqueName && a.cfg.UniqueNames > Unlimited) {
envs, err := a.str.FindEnvironmentsForAccount(acctId, trx) envs, err := a.str.FindEnvironmentsForAccount(acctId, trx)
if err != nil { if err != nil {
return false, err return false, err
} }
total := 0 total := 0
reserveds := 0
uniqueNames := 0
for i := range envs { for i := range envs {
shrs, err := a.str.FindSharesForEnvironment(envs[i].Id, trx) shrs, err := a.str.FindSharesForEnvironment(envs[i].Id, trx)
if err != nil { if err != nil {
return false, errors.Wrapf(err, "unable to find shares for environment '%v'", envs[i].ZId) return false, errors.Wrapf(err, "unable to find shares for environment '%v'", envs[i].ZId)
} }
total += len(shrs) total += len(shrs)
for _, shr := range shrs {
if shr.Reserved {
reserveds++
}
if shr.UniqueName {
uniqueNames++
}
}
if total+1 > a.cfg.Shares { if total+1 > a.cfg.Shares {
return false, nil return false, nil
} }
if reserved && reserveds+1 > a.cfg.ReservedShares {
return false, nil
}
if reserved && uniqueName && uniqueNames+1 > a.cfg.UniqueNames {
return false, nil
}
logrus.Infof("total = %d", total) logrus.Infof("total = %d", total)
} }
} }

View File

@ -49,7 +49,7 @@ func (h *shareHandler) Handle(params share.ShareParams, principal *rest_model_zr
return share.NewShareInternalServerError() return share.NewShareInternalServerError()
} }
if err := h.checkLimits(envId, principal, trx); err != nil { if err := h.checkLimits(envId, principal, params.Body.Reserved, params.Body.UniqueName != "", trx); err != nil {
logrus.Errorf("limits error: %v", err) logrus.Errorf("limits error: %v", err)
return share.NewShareUnauthorized() return share.NewShareUnauthorized()
} }
@ -190,10 +190,10 @@ func (h *shareHandler) Handle(params share.ShareParams, principal *rest_model_zr
}) })
} }
func (h *shareHandler) checkLimits(envId int, principal *rest_model_zrok.Principal, trx *sqlx.Tx) error { func (h *shareHandler) checkLimits(envId int, principal *rest_model_zrok.Principal, reserved, uniqueName bool, trx *sqlx.Tx) error {
if !principal.Limitless { if !principal.Limitless {
if limitsAgent != nil { if limitsAgent != nil {
ok, err := limitsAgent.CanCreateShare(int(principal.ID), envId, trx) ok, err := limitsAgent.CanCreateShare(int(principal.ID), envId, reserved, uniqueName, trx)
if err != nil { if err != nil {
return errors.Wrapf(err, "error checking share limits for '%v'", principal.Email) return errors.Wrapf(err, "error checking share limits for '%v'", principal.Email)
} }