From 9020415e74274b795a1c1a014f5aba90219248e5 Mon Sep 17 00:00:00 2001
From: celogeek <me@celogeek.com>
Date: Sun, 8 May 2022 17:03:46 +0200
Subject: [PATCH] check checksum, cleanup

---
 internal/photos/api/upload.go | 35 +++++++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/internal/photos/api/upload.go b/internal/photos/api/upload.go
index 9731e8b..c0b7e8a 100644
--- a/internal/photos/api/upload.go
+++ b/internal/photos/api/upload.go
@@ -18,8 +18,9 @@ const (
 )
 
 var (
-	ErrUploadNotExists = errors.New("upload id doesn't exists")
-	ErrUploadTooLarge  = fmt.Errorf("upload part too large ( > %d B)", MaxUploadPartSize)
+	ErrUploadNotExists       = errors.New("upload id doesn't exists")
+	ErrUploadPartTooLarge    = fmt.Errorf("upload part too large (> %d B)", MaxUploadPartSize)
+	ErrUploadPartWrongSha256 = errors.New("upload part wrong sha256")
 )
 
 func (s *Service) UploadCreate(c *gin.Context) {
@@ -65,25 +66,39 @@ func (s *Service) UploadPart(c *gin.Context) {
 	}
 
 	if c.Request.ContentLength > MaxUploadPartSize {
-		c.AbortWithError(http.StatusRequestEntityTooLarge, ErrUploadTooLarge)
+		c.AbortWithError(http.StatusRequestEntityTooLarge, ErrUploadPartTooLarge)
 		return
 	}
 
-	f, err := os.Create(
-		s.StorageTmp.Join(upload.Id, fmt.Sprint(uploadPart.Part)),
-	)
+	tmp_file := s.StorageTmp.Join(upload.Id, fmt.Sprintf("._tmp_%d", uploadPart.Part))
+	file := s.StorageTmp.Join(upload.Id, fmt.Sprint(uploadPart.Part))
+
+	f, err := os.Create(tmp_file)
 	if err != nil {
 		c.AbortWithError(http.StatusInternalServerError, err)
 		return
 	}
 
-	defer f.Close()
-	defer c.Request.Body.Close()
-
 	sha := sha256.New()
 	t := io.TeeReader(c.Request.Body, sha)
 	w, err := io.Copy(f, t)
 	if err != nil {
+		f.Close()
+		os.Remove(tmp_file)
+		c.AbortWithError(http.StatusInternalServerError, err)
+		return
+	}
+	f.Close()
+
+	shastr := hex.EncodeToString(sha.Sum(nil))
+	if shastr != uploadPart.PartSha256 {
+		os.Remove(tmp_file)
+		c.AbortWithError(http.StatusBadRequest, ErrUploadPartWrongSha256)
+		return
+	}
+
+	if err = os.Rename(tmp_file, file); err != nil {
+		os.Remove(tmp_file)
 		c.AbortWithError(http.StatusInternalServerError, err)
 		return
 	}
@@ -92,7 +107,7 @@ func (s *Service) UploadPart(c *gin.Context) {
 		"upload_id": upload.Id,
 		"part":      uploadPart.Part,
 		"size":      w,
-		"sha256":    hex.EncodeToString(sha.Sum(nil)),
+		"sha256":    uploadPart.PartSha256,
 		"status":    "success",
 	})
 }