mirror of
https://github.com/grafana/grafana.git
synced 2024-12-23 15:40:19 -06:00
fix(quota): fixed failing quota unit tests
This commit is contained in:
parent
1f959272c5
commit
5e949b0564
@ -1,7 +1,6 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
@ -254,96 +253,3 @@ func (ctx *Context) JsonApiErr(status int, message string, err error) {
|
||||
|
||||
ctx.JSON(status, resp)
|
||||
}
|
||||
|
||||
func Quota(target string) macaron.Handler {
|
||||
return func(c *Context) {
|
||||
limitReached, err := QuotaReached(c, target)
|
||||
if err != nil {
|
||||
c.JsonApiErr(500, "failed to get quota", err)
|
||||
return
|
||||
}
|
||||
if limitReached {
|
||||
c.JsonApiErr(403, fmt.Sprintf("%s Quota reached", target), nil)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func QuotaReached(c *Context, target string) (bool, error) {
|
||||
if !setting.Quota.Enabled {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// get the list of scopes that this target is valid for. Org, User, Global
|
||||
scopes, err := m.GetQuotaScopes(target)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
log.Info(fmt.Sprintf("checking quota for %s in scopes %v", target, scopes))
|
||||
|
||||
for _, scope := range scopes {
|
||||
log.Info(fmt.Sprintf("checking scope %s", scope.Name))
|
||||
switch scope.Name {
|
||||
case "global":
|
||||
if scope.DefaultLimit < 0 {
|
||||
continue
|
||||
}
|
||||
if scope.DefaultLimit == 0 {
|
||||
return true, nil
|
||||
}
|
||||
if target == "session" {
|
||||
usedSessions := sessionManager.Count()
|
||||
if int64(usedSessions) > scope.DefaultLimit {
|
||||
log.Info(fmt.Sprintf("%d sessions active, limit is %d", usedSessions, scope.DefaultLimit))
|
||||
return true, nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
query := m.GetGlobalQuotaByTargetQuery{Target: scope.Target}
|
||||
if err := bus.Dispatch(&query); err != nil {
|
||||
return true, err
|
||||
}
|
||||
if query.Result.Used >= scope.DefaultLimit {
|
||||
return true, nil
|
||||
}
|
||||
case "org":
|
||||
if !c.IsSignedIn {
|
||||
continue
|
||||
}
|
||||
query := m.GetOrgQuotaByTargetQuery{OrgId: c.OrgId, Target: scope.Target, Default: scope.DefaultLimit}
|
||||
if err := bus.Dispatch(&query); err != nil {
|
||||
return true, err
|
||||
}
|
||||
if query.Result.Limit < 0 {
|
||||
continue
|
||||
}
|
||||
if query.Result.Limit == 0 {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
if query.Result.Used >= query.Result.Limit {
|
||||
return true, nil
|
||||
}
|
||||
case "user":
|
||||
if !c.IsSignedIn || c.UserId == 0 {
|
||||
continue
|
||||
}
|
||||
query := m.GetUserQuotaByTargetQuery{UserId: c.UserId, Target: scope.Target, Default: scope.DefaultLimit}
|
||||
if err := bus.Dispatch(&query); err != nil {
|
||||
return true, err
|
||||
}
|
||||
if query.Result.Limit < 0 {
|
||||
continue
|
||||
}
|
||||
if query.Result.Limit == 0 {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
if query.Result.Used >= query.Result.Limit {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
||||
|
106
pkg/middleware/quota.go
Normal file
106
pkg/middleware/quota.go
Normal file
@ -0,0 +1,106 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/Unknwon/macaron"
|
||||
"github.com/grafana/grafana/pkg/bus"
|
||||
"github.com/grafana/grafana/pkg/log"
|
||||
m "github.com/grafana/grafana/pkg/models"
|
||||
"github.com/grafana/grafana/pkg/setting"
|
||||
)
|
||||
|
||||
func Quota(target string) macaron.Handler {
|
||||
return func(c *Context) {
|
||||
limitReached, err := QuotaReached(c, target)
|
||||
if err != nil {
|
||||
c.JsonApiErr(500, "failed to get quota", err)
|
||||
return
|
||||
}
|
||||
if limitReached {
|
||||
c.JsonApiErr(403, fmt.Sprintf("%s Quota reached", target), nil)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func QuotaReached(c *Context, target string) (bool, error) {
|
||||
if !setting.Quota.Enabled {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// get the list of scopes that this target is valid for. Org, User, Global
|
||||
scopes, err := m.GetQuotaScopes(target)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
log.Debug(fmt.Sprintf("checking quota for %s in scopes %v", target, scopes))
|
||||
|
||||
for _, scope := range scopes {
|
||||
log.Debug(fmt.Sprintf("checking scope %s", scope.Name))
|
||||
|
||||
switch scope.Name {
|
||||
case "global":
|
||||
if scope.DefaultLimit < 0 {
|
||||
continue
|
||||
}
|
||||
if scope.DefaultLimit == 0 {
|
||||
return true, nil
|
||||
}
|
||||
if target == "session" {
|
||||
usedSessions := getSessionCount()
|
||||
if int64(usedSessions) > scope.DefaultLimit {
|
||||
log.Debug(fmt.Sprintf("%d sessions active, limit is %d", usedSessions, scope.DefaultLimit))
|
||||
return true, nil
|
||||
}
|
||||
continue
|
||||
}
|
||||
query := m.GetGlobalQuotaByTargetQuery{Target: scope.Target}
|
||||
if err := bus.Dispatch(&query); err != nil {
|
||||
return true, err
|
||||
}
|
||||
if query.Result.Used >= scope.DefaultLimit {
|
||||
return true, nil
|
||||
}
|
||||
case "org":
|
||||
if !c.IsSignedIn {
|
||||
continue
|
||||
}
|
||||
query := m.GetOrgQuotaByTargetQuery{OrgId: c.OrgId, Target: scope.Target, Default: scope.DefaultLimit}
|
||||
if err := bus.Dispatch(&query); err != nil {
|
||||
return true, err
|
||||
}
|
||||
if query.Result.Limit < 0 {
|
||||
continue
|
||||
}
|
||||
if query.Result.Limit == 0 {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
if query.Result.Used >= query.Result.Limit {
|
||||
return true, nil
|
||||
}
|
||||
case "user":
|
||||
if !c.IsSignedIn || c.UserId == 0 {
|
||||
continue
|
||||
}
|
||||
query := m.GetUserQuotaByTargetQuery{UserId: c.UserId, Target: scope.Target, Default: scope.DefaultLimit}
|
||||
if err := bus.Dispatch(&query); err != nil {
|
||||
return true, err
|
||||
}
|
||||
if query.Result.Limit < 0 {
|
||||
continue
|
||||
}
|
||||
if query.Result.Limit == 0 {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
if query.Result.Used >= query.Result.Limit {
|
||||
return true, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return false, nil
|
||||
}
|
@ -1,16 +1,22 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/grafana/grafana/pkg/bus"
|
||||
m "github.com/grafana/grafana/pkg/models"
|
||||
"github.com/grafana/grafana/pkg/setting"
|
||||
. "github.com/smartystreets/goconvey/convey"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMiddlewareQuota(t *testing.T) {
|
||||
|
||||
Convey("Given the grafana quota middleware", t, func() {
|
||||
getSessionCount = func() int {
|
||||
return 4
|
||||
}
|
||||
|
||||
setting.AnonymousEnabled = false
|
||||
setting.Quota = setting.QuotaSettings{
|
||||
Enabled: true,
|
||||
Org: &setting.OrgQuota{
|
||||
|
@ -18,12 +18,16 @@ const (
|
||||
var sessionManager *session.Manager
|
||||
var sessionOptions *session.Options
|
||||
var startSessionGC func()
|
||||
var getSessionCount func() int
|
||||
|
||||
func init() {
|
||||
startSessionGC = func() {
|
||||
sessionManager.GC()
|
||||
time.AfterFunc(time.Duration(sessionOptions.Gclifetime)*time.Second, startSessionGC)
|
||||
}
|
||||
getSessionCount = func() int {
|
||||
return sessionManager.Count()
|
||||
}
|
||||
}
|
||||
|
||||
func prepareOptions(opt *session.Options) *session.Options {
|
||||
|
Loading…
Reference in New Issue
Block a user