From 3e21d30f626c3120d0dd932a749de3862c74f8cc Mon Sep 17 00:00:00 2001
From: celogeek <me@celogeek.com>
Date: Sun, 8 May 2022 16:31:06 +0200
Subject: [PATCH] improve storage

---
 internal/photos/api/main.go    | 28 +++++++++++++++-------------
 internal/photos/api/storage.go | 12 ++++++------
 internal/photos/api/upload.go  | 18 ++++++++++++++----
 3 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/internal/photos/api/main.go b/internal/photos/api/main.go
index 9ac10e6..7c3fc7a 100644
--- a/internal/photos/api/main.go
+++ b/internal/photos/api/main.go
@@ -20,13 +20,14 @@ var (
 )
 
 type Service struct {
-	Gin     *gin.Engine
-	DB      *gorm.DB
-	Config  *ServiceConfig
-	Store   *photosstore.Store
-	Storage *Storage
-	LogOk   *Logger
-	LogErr  *Logger
+	Gin           *gin.Engine
+	DB            *gorm.DB
+	Config        *ServiceConfig
+	Store         *photosstore.Store
+	StorageTmp    *Storage
+	StorageUpload *Storage
+	LogOk         *Logger
+	LogErr        *Logger
 }
 
 type ServiceConfig struct {
@@ -37,12 +38,13 @@ type ServiceConfig struct {
 
 func New(config *ServiceConfig) *Service {
 	return &Service{
-		Gin:     gin.New(),
-		Config:  config,
-		Store:   &photosstore.Store{Path: config.StorePath},
-		Storage: &Storage{Path: config.StorePath},
-		LogOk:   &Logger{os.Stdout, "Photos"},
-		LogErr:  &Logger{os.Stderr, "Photos"},
+		Gin:           gin.New(),
+		Config:        config,
+		Store:         &photosstore.Store{Path: config.StorePath},
+		StorageTmp:    NewStorage(config.StorePath, "tmp"),
+		StorageUpload: NewStorage(config.StorePath, "upload"),
+		LogOk:         &Logger{os.Stdout, "Photos"},
+		LogErr:        &Logger{os.Stderr, "Photos"},
 	}
 }
 
diff --git a/internal/photos/api/storage.go b/internal/photos/api/storage.go
index 13d31cc..b6b0d3e 100644
--- a/internal/photos/api/storage.go
+++ b/internal/photos/api/storage.go
@@ -7,16 +7,16 @@ import (
 )
 
 type Storage struct {
-	Path string
+	BasePath string
+	Path     string
 }
 
-const (
-	StorageTmp    = "tmp"
-	StorageUpload = "upload"
-)
+func NewStorage(basePath, path string) *Storage {
+	return &Storage{basePath, path}
+}
 
 func (s *Storage) Join(paths ...string) string {
-	return filepath.Join(s.Path, filepath.Join(paths...))
+	return filepath.Join(s.BasePath, s.Path, filepath.Join(paths...))
 }
 
 func (s *Storage) Create(paths ...string) error {
diff --git a/internal/photos/api/upload.go b/internal/photos/api/upload.go
index fa8704d..9731e8b 100644
--- a/internal/photos/api/upload.go
+++ b/internal/photos/api/upload.go
@@ -13,8 +13,13 @@ import (
 	"github.com/google/uuid"
 )
 
+const (
+	MaxUploadPartSize = 4 << 20 // 4MB
+)
+
 var (
 	ErrUploadNotExists = errors.New("upload id doesn't exists")
+	ErrUploadTooLarge  = fmt.Errorf("upload part too large ( > %d B)", MaxUploadPartSize)
 )
 
 func (s *Service) UploadCreate(c *gin.Context) {
@@ -24,7 +29,7 @@ func (s *Service) UploadCreate(c *gin.Context) {
 		return
 	}
 
-	if err := s.Storage.Create(StorageTmp, sha.String()); err != nil {
+	if err := s.StorageTmp.Create(sha.String()); err != nil {
 		c.AbortWithError(http.StatusInternalServerError, err)
 		return
 	}
@@ -54,13 +59,18 @@ func (s *Service) UploadPart(c *gin.Context) {
 		return
 	}
 
-	if !s.Storage.Exists(StorageTmp, upload.Id) {
+	if !s.StorageTmp.Exists(upload.Id) {
 		c.AbortWithError(http.StatusNotFound, ErrUploadNotExists)
 		return
 	}
 
+	if c.Request.ContentLength > MaxUploadPartSize {
+		c.AbortWithError(http.StatusRequestEntityTooLarge, ErrUploadTooLarge)
+		return
+	}
+
 	f, err := os.Create(
-		s.Storage.Join(StorageTmp, upload.Id, fmt.Sprint(uploadPart.Part)),
+		s.StorageTmp.Join(upload.Id, fmt.Sprint(uploadPart.Part)),
 	)
 	if err != nil {
 		c.AbortWithError(http.StatusInternalServerError, err)
@@ -93,7 +103,7 @@ func (s *Service) UploadCancel(c *gin.Context) {
 		return
 	}
 
-	if err := s.Storage.Delete(StorageTmp, upload.Id); err != nil {
+	if err := s.StorageTmp.Delete(upload.Id); err != nil {
 		c.AbortWithError(http.StatusNotFound, ErrUploadNotExists)
 		return
 	}