From 0d817987ec53bd58b2da9c447353f57bca91c951 Mon Sep 17 00:00:00 2001
From: Artur Wierzbicki <wierzbicki.artur.94@gmail.com>
Date: Mon, 1 Aug 2022 22:53:44 +0400
Subject: [PATCH] Storage: move file quota check to service (#53058)

* Storage: move file quota check to service

* Fix tests
---
 pkg/services/store/http.go         | 13 +++----------
 pkg/services/store/service.go      | 25 +++++++++++++++++++++++++
 pkg/services/store/service_test.go |  2 ++
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/pkg/services/store/http.go b/pkg/services/store/http.go
index 8246b5a3a46..b68df21b368 100644
--- a/pkg/services/store/http.go
+++ b/pkg/services/store/http.go
@@ -28,6 +28,9 @@ func UploadErrorToStatusCode(err error) int {
 	case errors.Is(err, ErrValidationFailed):
 		return 400
 
+	case errors.Is(err, ErrQuotaReached):
+		return 400
+
 	case errors.Is(err, ErrFileAlreadyExists):
 		return 400
 
@@ -70,16 +73,6 @@ func (s *standardStorageService) doWrite(c *models.ReqContext) response.Response
 }
 
 func (s *standardStorageService) doUpload(c *models.ReqContext) response.Response {
-	// assumes we are only uploading to the SQL database - TODO: refactor once we introduce object stores
-	quotaReached, err := s.quotaService.CheckQuotaReached(c.Req.Context(), "file", nil)
-	if err != nil {
-		return response.Error(500, "Internal server error", err)
-	}
-
-	if quotaReached {
-		return response.Error(400, "File quota reached", errors.New("file quota reached"))
-	}
-
 	type rspInfo struct {
 		Message string `json:"message,omitempty"`
 		Path    string `json:"path,omitempty"`
diff --git a/pkg/services/store/service.go b/pkg/services/store/service.go
index ad69a4d65ad..93de53853e6 100644
--- a/pkg/services/store/service.go
+++ b/pkg/services/store/service.go
@@ -25,6 +25,7 @@ var grafanaStorageLogger = log.New("grafanaStorageLogger")
 
 var ErrUnsupportedStorage = errors.New("storage does not support this operation")
 var ErrUploadInternalError = errors.New("upload internal error")
+var ErrQuotaReached = errors.New("file quota reached")
 var ErrValidationFailed = errors.New("request validation failed")
 var ErrFileAlreadyExists = errors.New("file exists")
 var ErrStorageNotFound = errors.New("storage not found")
@@ -300,6 +301,10 @@ type UploadRequest struct {
 }
 
 func (s *standardStorageService) Upload(ctx context.Context, user *models.SignedInUser, req *UploadRequest) error {
+	if err := s.checkFileQuota(ctx, req.Path); err != nil {
+		return err
+	}
+
 	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(req.Path))
 	if !guardian.canWrite(req.Path) {
 		return ErrAccessDenied
@@ -348,6 +353,22 @@ func (s *standardStorageService) Upload(ctx context.Context, user *models.Signed
 	return nil
 }
 
+func (s *standardStorageService) checkFileQuota(ctx context.Context, path string) error {
+	// assumes we are only uploading to the SQL database - TODO: refactor once we introduce object stores
+	quotaReached, err := s.quotaService.CheckQuotaReached(ctx, "file", nil)
+	if err != nil {
+		grafanaStorageLogger.Error("failed while checking upload quota", "path", path, "error", err)
+		return ErrUploadInternalError
+	}
+
+	if quotaReached {
+		grafanaStorageLogger.Info("reached file quota", "path", path)
+		return ErrQuotaReached
+	}
+
+	return nil
+}
+
 func (s *standardStorageService) DeleteFolder(ctx context.Context, user *models.SignedInUser, cmd *DeleteFolderCmd) error {
 	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(cmd.Path))
 	if !guardian.canDelete(cmd.Path) {
@@ -370,6 +391,10 @@ func (s *standardStorageService) DeleteFolder(ctx context.Context, user *models.
 }
 
 func (s *standardStorageService) CreateFolder(ctx context.Context, user *models.SignedInUser, cmd *CreateFolderCmd) error {
+	if err := s.checkFileQuota(ctx, cmd.Path); err != nil {
+		return err
+	}
+
 	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(cmd.Path))
 	if !guardian.canWrite(cmd.Path) {
 		return ErrAccessDenied
diff --git a/pkg/services/store/service_test.go b/pkg/services/store/service_test.go
index 871c81c3f86..647cf9cdfb6 100644
--- a/pkg/services/store/service_test.go
+++ b/pkg/services/store/service_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/experimental"
 	"github.com/grafana/grafana/pkg/infra/filestorage"
 	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/quota/quotatest"
 	"github.com/grafana/grafana/pkg/services/sqlstore"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/tsdb/testdatasource"
@@ -117,6 +118,7 @@ func setupUploadStore(t *testing.T, authService storageAuthService) (StorageServ
 	store.cfg = &GlobalStorageConfig{
 		AllowUnsanitizedSvgUpload: true,
 	}
+	store.quotaService = quotatest.NewQuotaServiceFake()
 
 	return store, mockStorage, storageName
 }