@@ -682,6 +682,7 @@ type fakeUserAuthTokenService struct {
|
||||
tryRotateTokenProvider func(token *m.UserToken, clientIP, userAgent string) (bool, error)
|
||||
lookupTokenProvider func(unhashedToken string) (*m.UserToken, error)
|
||||
revokeTokenProvider func(token *m.UserToken) error
|
||||
activeAuthTokenCount func() (int64, error)
|
||||
}
|
||||
|
||||
func newFakeUserAuthTokenService() *fakeUserAuthTokenService {
|
||||
@@ -704,6 +705,9 @@ func newFakeUserAuthTokenService() *fakeUserAuthTokenService {
|
||||
revokeTokenProvider: func(token *m.UserToken) error {
|
||||
return nil
|
||||
},
|
||||
activeAuthTokenCount: func() (int64, error) {
|
||||
return 10, nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -722,3 +726,7 @@ func (s *fakeUserAuthTokenService) TryRotateToken(token *m.UserToken, clientIP,
|
||||
func (s *fakeUserAuthTokenService) RevokeToken(token *m.UserToken) error {
|
||||
return s.revokeTokenProvider(token)
|
||||
}
|
||||
|
||||
func (s *fakeUserAuthTokenService) ActiveTokenCount() (int64, error) {
|
||||
return s.activeAuthTokenCount()
|
||||
}
|
||||
|
||||
+14
-10
@@ -9,16 +9,20 @@ import (
|
||||
"github.com/grafana/grafana/pkg/services/quota"
|
||||
)
|
||||
|
||||
func Quota(target string) macaron.Handler {
|
||||
return func(c *m.ReqContext) {
|
||||
limitReached, err := quota.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
|
||||
// Quota returns a function that returns a function used to call quotaservice based on target name
|
||||
func Quota(quotaService *quota.QuotaService) func(target string) macaron.Handler {
|
||||
//https://open.spotify.com/track/7bZSoBEAEEUsGEuLOf94Jm?si=T1Tdju5qRSmmR0zph_6RBw fuuuuunky
|
||||
return func(target string) macaron.Handler {
|
||||
return func(c *m.ReqContext) {
|
||||
limitReached, err := quotaService.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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,9 +3,10 @@ package middleware
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/grafana/grafana/pkg/services/quota"
|
||||
|
||||
"github.com/grafana/grafana/pkg/bus"
|
||||
m "github.com/grafana/grafana/pkg/models"
|
||||
"github.com/grafana/grafana/pkg/services/session"
|
||||
"github.com/grafana/grafana/pkg/setting"
|
||||
. "github.com/smartystreets/goconvey/convey"
|
||||
)
|
||||
@@ -13,10 +14,6 @@ import (
|
||||
func TestMiddlewareQuota(t *testing.T) {
|
||||
|
||||
Convey("Given the grafana quota middleware", t, func() {
|
||||
session.GetSessionCount = func() int {
|
||||
return 4
|
||||
}
|
||||
|
||||
setting.AnonymousEnabled = false
|
||||
setting.Quota = setting.QuotaSettings{
|
||||
Enabled: true,
|
||||
@@ -39,6 +36,12 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
fakeAuthTokenService := newFakeUserAuthTokenService()
|
||||
qs := "a.QuotaService{
|
||||
AuthTokenService: fakeAuthTokenService,
|
||||
}
|
||||
QuotaFn := Quota(qs)
|
||||
|
||||
middlewareScenario("with user not logged in", func(sc *scenarioContext) {
|
||||
bus.AddHandler("globalQuota", func(query *m.GetGlobalQuotaByTargetQuery) error {
|
||||
query.Result = &m.GlobalQuotaDTO{
|
||||
@@ -48,26 +51,30 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
Convey("global quota not reached", func() {
|
||||
sc.m.Get("/user", Quota("user"), sc.defaultHandler)
|
||||
sc.m.Get("/user", QuotaFn("user"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/user").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 200)
|
||||
})
|
||||
|
||||
Convey("global quota reached", func() {
|
||||
setting.Quota.Global.User = 4
|
||||
sc.m.Get("/user", Quota("user"), sc.defaultHandler)
|
||||
sc.m.Get("/user", QuotaFn("user"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/user").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 403)
|
||||
})
|
||||
|
||||
Convey("global session quota not reached", func() {
|
||||
setting.Quota.Global.Session = 10
|
||||
sc.m.Get("/user", Quota("session"), sc.defaultHandler)
|
||||
sc.m.Get("/user", QuotaFn("session"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/user").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 200)
|
||||
})
|
||||
|
||||
Convey("global session quota reached", func() {
|
||||
setting.Quota.Global.Session = 1
|
||||
sc.m.Get("/user", Quota("session"), sc.defaultHandler)
|
||||
sc.m.Get("/user", QuotaFn("session"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/user").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 403)
|
||||
})
|
||||
@@ -95,6 +102,7 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
bus.AddHandler("userQuota", func(query *m.GetUserQuotaByTargetQuery) error {
|
||||
query.Result = &m.UserQuotaDTO{
|
||||
Target: query.Target,
|
||||
@@ -103,6 +111,7 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
bus.AddHandler("orgQuota", func(query *m.GetOrgQuotaByTargetQuery) error {
|
||||
query.Result = &m.OrgQuotaDTO{
|
||||
Target: query.Target,
|
||||
@@ -111,45 +120,49 @@ func TestMiddlewareQuota(t *testing.T) {
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
Convey("global datasource quota reached", func() {
|
||||
setting.Quota.Global.DataSource = 4
|
||||
sc.m.Get("/ds", Quota("data_source"), sc.defaultHandler)
|
||||
sc.m.Get("/ds", QuotaFn("data_source"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/ds").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 403)
|
||||
})
|
||||
|
||||
Convey("user Org quota not reached", func() {
|
||||
setting.Quota.User.Org = 5
|
||||
sc.m.Get("/org", Quota("org"), sc.defaultHandler)
|
||||
sc.m.Get("/org", QuotaFn("org"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/org").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 200)
|
||||
})
|
||||
|
||||
Convey("user Org quota reached", func() {
|
||||
setting.Quota.User.Org = 4
|
||||
sc.m.Get("/org", Quota("org"), sc.defaultHandler)
|
||||
sc.m.Get("/org", QuotaFn("org"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/org").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 403)
|
||||
})
|
||||
|
||||
Convey("org dashboard quota not reached", func() {
|
||||
setting.Quota.Org.Dashboard = 10
|
||||
sc.m.Get("/dashboard", Quota("dashboard"), sc.defaultHandler)
|
||||
sc.m.Get("/dashboard", QuotaFn("dashboard"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/dashboard").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 200)
|
||||
})
|
||||
|
||||
Convey("org dashboard quota reached", func() {
|
||||
setting.Quota.Org.Dashboard = 4
|
||||
sc.m.Get("/dashboard", Quota("dashboard"), sc.defaultHandler)
|
||||
sc.m.Get("/dashboard", QuotaFn("dashboard"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/dashboard").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 403)
|
||||
})
|
||||
|
||||
Convey("org dashboard quota reached but quotas disabled", func() {
|
||||
setting.Quota.Org.Dashboard = 4
|
||||
setting.Quota.Enabled = false
|
||||
sc.m.Get("/dashboard", Quota("dashboard"), sc.defaultHandler)
|
||||
sc.m.Get("/dashboard", QuotaFn("dashboard"), sc.defaultHandler)
|
||||
sc.fakeReq("GET", "/dashboard").exec()
|
||||
So(sc.resp.Code, ShouldEqual, 200)
|
||||
})
|
||||
|
||||
})
|
||||
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user