From 5e949b05646058ce1186a388db4f756ea5448521 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torkel=20=C3=96degaard?= Date: Fri, 18 Sep 2015 08:36:58 +0200 Subject: [PATCH] fix(quota): fixed failing quota unit tests --- pkg/middleware/middleware.go | 94 ------------------------------- pkg/middleware/quota.go | 106 +++++++++++++++++++++++++++++++++++ pkg/middleware/quota_test.go | 8 ++- pkg/middleware/session.go | 4 ++ 4 files changed, 117 insertions(+), 95 deletions(-) create mode 100644 pkg/middleware/quota.go diff --git a/pkg/middleware/middleware.go b/pkg/middleware/middleware.go index 8a61bdb3a0b..8704ec5a787 100644 --- a/pkg/middleware/middleware.go +++ b/pkg/middleware/middleware.go @@ -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 -} diff --git a/pkg/middleware/quota.go b/pkg/middleware/quota.go new file mode 100644 index 00000000000..f6ba74d77df --- /dev/null +++ b/pkg/middleware/quota.go @@ -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 +} diff --git a/pkg/middleware/quota_test.go b/pkg/middleware/quota_test.go index ae6b414b248..b68aa485fa7 100644 --- a/pkg/middleware/quota_test.go +++ b/pkg/middleware/quota_test.go @@ -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{ diff --git a/pkg/middleware/session.go b/pkg/middleware/session.go index 7b036b9790e..fc1512e5bd6 100644 --- a/pkg/middleware/session.go +++ b/pkg/middleware/session.go @@ -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 {