zrok/controller/limits/userLimits.go

package limits

import (
	"github.com/jmoiron/sqlx"
	"github.com/openziti/zrok/controller/store"
	"github.com/openziti/zrok/sdk/golang/sdk"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

type userLimits struct {
	resource  store.ResourceCountClass
	bandwidth []store.BandwidthClass
	scopes    map[sdk.BackendMode]*store.LimitClass
}

func (ul *userLimits) toBandwidthArray(backendMode sdk.BackendMode) []store.BandwidthClass {
	if scopedBwc, found := ul.scopes[backendMode]; found {
		out := make([]store.BandwidthClass, 0)
		for _, bwc := range ul.bandwidth {
			out = append(out, bwc)
		}
		out = append(out, scopedBwc)
		return out
	}
	return ul.bandwidth
}

func (ul *userLimits) ignoreBackends(bwc store.BandwidthClass) map[sdk.BackendMode]bool {
	if bwc.IsScoped() {
		ignoreBackends := make(map[sdk.BackendMode]bool)
		for backendMode := range ul.scopes {
			if backendMode != bwc.GetBackendMode() {
				ignoreBackends[backendMode] = true
			}
		}
		return ignoreBackends
	} else {
		ignoreBackends := make(map[sdk.BackendMode]bool)
		for backendMode := range ul.scopes {
			ignoreBackends[backendMode] = true
		}
		return ignoreBackends
	}
}

func (a *Agent) getUserLimits(acctId int, trx *sqlx.Tx) (*userLimits, error) {
	resource := newConfigResourceCountClass(a.cfg)
	cfgBwcs := newConfigBandwidthClasses(a.cfg.Bandwidth)
	bwWarning := cfgBwcs[0]
	bwLimit := cfgBwcs[1]
	scopes := make(map[sdk.BackendMode]*store.LimitClass)

	alcs, err := a.str.FindAppliedLimitClassesForAccount(acctId, trx)
	if err != nil {
		return nil, errors.Wrapf(err, "error finding applied limit classes for account '%d'", acctId)
	}
	for _, alc := range alcs {
		if a.isResourceCountClass(alc) {
			resource = alc
		} else if a.isUnscopedBandwidthClass(alc) {
			if alc.LimitAction == store.WarningLimitAction {
				bwWarning = alc
			} else {
				bwLimit = alc
			}
		} else if a.isScopedLimitClass(alc) {
			scopes[*alc.BackendMode] = alc
		} else {
			logrus.Warnf("unknown type of limit class '%v' for account '#%d'", alc, acctId)
		}
	}

	userLimits := &userLimits{
		resource:  resource,
		bandwidth: []store.BandwidthClass{bwWarning, bwLimit},
		scopes:    scopes,
	}

	return userLimits, nil
}

func (a *Agent) isResourceCountClass(alc *store.LimitClass) bool {
	if alc.BackendMode != nil {
		return false
	}
	if alc.Environments == store.Unlimited && alc.Shares == store.Unlimited && alc.ReservedShares == store.Unlimited && alc.UniqueNames == store.Unlimited && alc.ShareFrontends == store.Unlimited {
		return false
	}
	return true
}

func (a *Agent) isUnscopedBandwidthClass(alc *store.LimitClass) bool {
	if alc.BackendMode != nil {
		return false
	}
	if alc.Environments > store.Unlimited || alc.Shares > store.Unlimited || alc.ReservedShares > store.Unlimited || alc.UniqueNames > store.Unlimited || alc.ShareFrontends > store.Unlimited {
		return false
	}
	if alc.PeriodMinutes < 1 {
		return false
	}
	if alc.RxBytes == store.Unlimited && alc.TxBytes == store.Unlimited && alc.TotalBytes == store.Unlimited {
		return false
	}
	return true
}

func (a *Agent) isScopedLimitClass(alc *store.LimitClass) bool {
	if alc.BackendMode == nil {
		return false
	}
	if alc.Environments > store.Unlimited {
		return false
	}
	if alc.PeriodMinutes < 1 {
		return false
	}
	if alc.RxBytes == store.Unlimited && alc.TxBytes == store.Unlimited && alc.TotalBytes == store.Unlimited {
		return false
	}
	return true
}