limit class filtering (#606)

This commit is contained in:
Michael Quigley 2024-05-23 14:08:14 -04:00
parent 58eb8bfcca
commit 896a4a7845
No known key found for this signature in database
GPG Key ID: 9B60314A9DD20A62
4 changed files with 76 additions and 12 deletions

View File

@ -7,6 +7,7 @@ import (
"github.com/openziti/zrok/controller/metrics" "github.com/openziti/zrok/controller/metrics"
"github.com/openziti/zrok/controller/store" "github.com/openziti/zrok/controller/store"
"github.com/openziti/zrok/controller/zrokEdgeSdk" "github.com/openziti/zrok/controller/zrokEdgeSdk"
"github.com/openziti/zrok/sdk/golang/sdk"
"github.com/openziti/zrok/util" "github.com/openziti/zrok/util"
"github.com/pkg/errors" "github.com/pkg/errors"
"github.com/sirupsen/logrus" "github.com/sirupsen/logrus"
@ -94,7 +95,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, reserved, uniqueName bool, trx *sqlx.Tx) (bool, error) { func (a *Agent) CanCreateShare(acctId, envId int, reserved, uniqueName bool, shareMode sdk.ShareMode, backendMode sdk.BackendMode, 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,6 +124,16 @@ func (a *Agent) CanCreateShare(acctId, envId int, reserved, uniqueName bool, trx
return false, err return false, err
} }
alc, err := a.str.FindLimitClassesForAccount(acctId, trx)
if err != nil {
logrus.Errorf("error finding limit classes for account with id '%d': %v", acctId, err)
return false, err
}
sortLimitClasses(alc)
if len(alc) > 0 {
logrus.Infof("selected limit class: %v", alc[0])
}
if a.cfg.Shares > Unlimited || (reserved && a.cfg.ReservedShares > Unlimited) || (reserved && uniqueName && a.cfg.UniqueNames > 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 {

View File

@ -0,0 +1,38 @@
package limits
import (
"github.com/openziti/zrok/controller/store"
"sort"
)
func sortLimitClasses(lcs []*store.LimitClass) {
sort.Slice(lcs, func(i, j int) bool {
ipoints := limitScopePoints(lcs[i]) + modePoints(lcs[i])
jpoints := limitScopePoints(lcs[j]) + modePoints(lcs[j])
return ipoints > jpoints
})
}
func limitScopePoints(lc *store.LimitClass) int {
points := 0
switch lc.LimitScope {
case store.AccountLimitScope:
points += 1000
case store.EnvironmentLimitScope:
points += 100
case store.ShareLimitScope:
points += 10
}
return points
}
func modePoints(lc *store.LimitClass) int {
points := 0
if lc.BackendMode != "" {
points += 1
}
if lc.ShareMode != "" {
points += 1
}
return points
}

View File

@ -49,7 +49,9 @@ func (h *shareHandler) Handle(params share.ShareParams, principal *rest_model_zr
return share.NewShareInternalServerError() return share.NewShareInternalServerError()
} }
if err := h.checkLimits(envId, principal, params.Body.Reserved, params.Body.UniqueName != "", trx); err != nil { shareMode := sdk.ShareMode(params.Body.ShareMode)
backendMode := sdk.BackendMode(params.Body.BackendMode)
if err := h.checkLimits(envId, principal, params.Body.Reserved, params.Body.UniqueName != "", shareMode, backendMode, trx); err != nil {
logrus.Errorf("limits error: %v", err) logrus.Errorf("limits error: %v", err)
return share.NewShareUnauthorized() return share.NewShareUnauthorized()
} }
@ -190,10 +192,10 @@ func (h *shareHandler) Handle(params share.ShareParams, principal *rest_model_zr
}) })
} }
func (h *shareHandler) checkLimits(envId int, principal *rest_model_zrok.Principal, reserved, uniqueName bool, trx *sqlx.Tx) error { func (h *shareHandler) checkLimits(envId int, principal *rest_model_zrok.Principal, reserved, uniqueName bool, shareMode sdk.ShareMode, backendMode sdk.BackendMode, trx *sqlx.Tx) error {
if !principal.Limitless { if !principal.Limitless {
if limitsAgent != nil { if limitsAgent != nil {
ok, err := limitsAgent.CanCreateShare(int(principal.ID), envId, reserved, uniqueName, trx) ok, err := limitsAgent.CanCreateShare(int(principal.ID), envId, reserved, uniqueName, shareMode, backendMode, 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)
} }

View File

@ -1,6 +1,7 @@
package store package store
import ( import (
"encoding/json"
"github.com/jmoiron/sqlx" "github.com/jmoiron/sqlx"
"github.com/openziti/zrok/sdk/golang/sdk" "github.com/openziti/zrok/sdk/golang/sdk"
"github.com/pkg/errors" "github.com/pkg/errors"
@ -8,14 +9,26 @@ import (
type LimitClass struct { type LimitClass struct {
Model Model
LimitScope LimitScope LimitScope LimitScope
LimitAction LimitAction LimitAction LimitAction
ShareMode sdk.ShareMode ShareMode sdk.ShareMode
BackendMode sdk.BackendMode BackendMode sdk.BackendMode
PeriodMinutes int Shares int
RxBytes int64 ReservedShares int
TxBytes int64 UniqueNames int
TotalBytes int64 PeriodMinutes int
RxBytes int64
TxBytes int64
TotalBytes int64
}
func (lc LimitClass) String() string {
out, err := json.MarshalIndent(&lc, "", " ")
if err != nil {
return ""
}
return string(out)
} }
func (str *Store) CreateLimitClass(lc *LimitClass, trx *sqlx.Tx) (int, error) { func (str *Store) CreateLimitClass(lc *LimitClass, trx *sqlx.Tx) (int, error) {