From 0688050552e10969f036f658e70aa16e7833baa1 Mon Sep 17 00:00:00 2001
From: woodsaj <anthony@monkey.id.au>
Date: Mon, 20 Jul 2015 22:38:25 +0800
Subject: [PATCH] add quota middleware to enforce quotas. issue #321

Conflicts:
	pkg/api/api.go
---
 pkg/api/api.go               |  6 ++++--
 pkg/api/dashboard.go         | 12 ++++++++++++
 pkg/middleware/middleware.go | 13 +++++++++++++
 pkg/models/quotas.go         | 12 ++++++++++++
 4 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/pkg/api/api.go b/pkg/api/api.go
index fce99e733f5..2b87863d188 100644
--- a/pkg/api/api.go
+++ b/pkg/api/api.go
@@ -14,6 +14,7 @@ func Register(r *macaron.Macaron) {
 	reqGrafanaAdmin := middleware.Auth(&middleware.AuthOptions{ReqSignedIn: true, ReqGrafanaAdmin: true})
 	reqEditorRole := middleware.RoleAuth(m.ROLE_EDITOR, m.ROLE_ADMIN)
 	regOrgAdmin := middleware.RoleAuth(m.ROLE_ADMIN)
+	limitQuota := middleware.LimitQuota
 	bind := binding.Bind
 
 	// not logged in views
@@ -95,7 +96,7 @@ func Register(r *macaron.Macaron) {
 			r.Get("/", wrap(GetOrgCurrent))
 			r.Put("/", bind(dtos.UpdateOrgForm{}), wrap(UpdateOrgCurrent))
 			r.Put("/address", bind(dtos.UpdateOrgAddressForm{}), wrap(UpdateOrgAddressCurrent))
-			r.Post("/users", bind(m.AddOrgUserCommand{}), wrap(AddOrgUserToCurrentOrg))
+			r.Post("/users", limitQuota(m.QUOTA_USER), bind(m.AddOrgUserCommand{}), wrap(AddOrgUserToCurrentOrg))
 			r.Get("/users", wrap(GetOrgUsersForCurrentOrg))
 			r.Patch("/users/:userId", bind(m.UpdateOrgUserCommand{}), wrap(UpdateOrgUserForCurrentOrg))
 			r.Delete("/users/:userId", wrap(RemoveOrgUserForCurrentOrg))
@@ -136,7 +137,7 @@ func Register(r *macaron.Macaron) {
 		// Data sources
 		r.Group("/datasources", func() {
 			r.Get("/", GetDataSources)
-			r.Post("/", bind(m.AddDataSourceCommand{}), AddDataSource)
+			r.Post("/", limitQuota(m.QUOTA_DATASOURCE), bind(m.AddDataSourceCommand{}), AddDataSource)
 			r.Put("/:id", bind(m.UpdateDataSourceCommand{}), UpdateDataSource)
 			r.Delete("/:id", DeleteDataSource)
 			r.Get("/:id", GetDataSourceById)
@@ -161,6 +162,7 @@ func Register(r *macaron.Macaron) {
 
 		// metrics
 		r.Get("/metrics/test", GetTestMetrics)
+
 	}, reqSignedIn)
 
 	// admin api
diff --git a/pkg/api/dashboard.go b/pkg/api/dashboard.go
index a10c3c92f96..21b69265da1 100644
--- a/pkg/api/dashboard.go
+++ b/pkg/api/dashboard.go
@@ -86,6 +86,18 @@ func DeleteDashboard(c *middleware.Context) {
 func PostDashboard(c *middleware.Context, cmd m.SaveDashboardCommand) {
 	cmd.OrgId = c.OrgId
 
+	dash := cmd.GetDashboardModel()
+	if dash.Id == 0 {
+		limitReached, err := m.QuotaReached(cmd.OrgId, m.QUOTA_DASHBOARD)
+		if err != nil {
+			c.JsonApiErr(500, "failed to get quota", err)
+		}
+		if limitReached {
+			c.JsonApiErr(403, "Quota reached", nil)
+			return
+		}
+	}
+
 	err := bus.Dispatch(&cmd)
 	if err != nil {
 		if err == m.ErrDashboardWithSameNameExists {
diff --git a/pkg/middleware/middleware.go b/pkg/middleware/middleware.go
index 8704ec5a787..e10b2f76781 100644
--- a/pkg/middleware/middleware.go
+++ b/pkg/middleware/middleware.go
@@ -253,3 +253,16 @@ func (ctx *Context) JsonApiErr(status int, message string, err error) {
 
 	ctx.JSON(status, resp)
 }
+
+func LimitQuota(target m.QuotaTarget) macaron.Handler {
+	return func(c *Context) {
+		limitReached, err := m.QuotaReached(c.OrgId, target)
+		if err != nil {
+			c.JsonApiErr(500, "failed to get quota", err)
+		}
+		if limitReached {
+			c.JsonApiErr(403, "Quota reached", nil)
+			return
+		}
+	}
+}
diff --git a/pkg/models/quotas.go b/pkg/models/quotas.go
index e296ddedb8e..d16bd0325cc 100644
--- a/pkg/models/quotas.go
+++ b/pkg/models/quotas.go
@@ -1,6 +1,7 @@
 package models
 
 import (
+	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/setting"
 	"time"
 )
@@ -61,3 +62,14 @@ type UpdateQuotaCmd struct {
 	Limit  int64       `json:"limit"`
 	OrgId  int64       `json:"-"`
 }
+
+func QuotaReached(org_id int64, target QuotaTarget) (bool, error) {
+	query := GetQuotaByTargetQuery{OrgId: org_id, Target: target}
+	if err := bus.Dispatch(&query); err != nil {
+		return true, err
+	}
+	if query.Result.Used >= query.Result.Limit {
+		return true, nil
+	}
+	return false, nil
+}