From f3be950ba76274ad0a3542d1a8f1f6cc532017df Mon Sep 17 00:00:00 2001
From: Derrick Hammer <derrick@derrickhammer.com>
Date: Thu, 25 Jan 2024 19:05:52 -0500
Subject: [PATCH] feat: detect and add mime type to upload so we don't need to
 make extra requests on runtime

---
 interfaces/storage.go |  2 +-
 storage/storage.go    | 30 ++++++++++++++++++++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/interfaces/storage.go b/interfaces/storage.go
index fb36313..e9737fc 100644
--- a/interfaces/storage.go
+++ b/interfaces/storage.go
@@ -19,7 +19,7 @@ type StorageService interface {
 	GetHashSmall(file io.ReadSeeker) ([]byte, error)
 	GetHash(file io.Reader) ([]byte, int64, error)
 	GetFile(hash []byte, start int64) (io.ReadCloser, int64, error)
-	CreateUpload(hash []byte, uploaderID uint, uploaderIP string, size uint64, protocol string) (*models.Upload, error)
+	CreateUpload(hash []byte, mime string, uploaderID uint, uploaderIP string, size uint64, protocol string) (*models.Upload, error)
 	TusUploadExists(hash []byte) (bool, models.TusUpload)
 	CreateTusUpload(hash []byte, uploadID string, uploaderID uint, uploaderIP string, protocol string) (*models.TusUpload, error)
 	TusUploadProgress(uploadID string) error
diff --git a/storage/storage.go b/storage/storage.go
index 27c8748..f4e8c2b 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -26,6 +26,7 @@ import (
 	"go.uber.org/zap"
 	"io"
 	"lukechampine.com/blake3"
+	"net/http"
 	"net/url"
 	"strings"
 	"time"
@@ -280,11 +281,12 @@ func (s *StorageServiceImpl) GetHash(file io.Reader) ([]byte, int64, error) {
 	return hash[:32], totalBytes, nil
 }
 
-func (s *StorageServiceImpl) CreateUpload(hash []byte, uploaderID uint, uploaderIP string, size uint64, protocol string) (*models.Upload, error) {
+func (s *StorageServiceImpl) CreateUpload(hash []byte, mime string, uploaderID uint, uploaderIP string, size uint64, protocol string) (*models.Upload, error) {
 	hashStr := hex.EncodeToString(hash)
 
 	upload := &models.Upload{
 		Hash:       hashStr,
+		MimeType:   mime,
 		UserID:     uploaderID,
 		UploaderIP: uploaderIP,
 		Protocol:   protocol,
@@ -522,6 +524,30 @@ func (s *StorageServiceImpl) buildNewTusUploadTask(upload *models.TusUpload) (jo
 				return err
 			}
 
+			var mimeBuf [512]byte
+
+			_, err = reader.Read(mimeBuf[:])
+
+			if err != nil {
+				s.portal.Logger().Error("Could not read mime", zap.Error(err))
+				return err
+			}
+
+			mimeType := http.DetectContentType(mimeBuf[:])
+
+			upload.MimeType = mimeType
+
+			if tx := s.Portal().Database().Save(upload); tx.Error != nil {
+				s.portal.Logger().Error("Could not update tus upload", zap.Error(tx.Error))
+				return tx.Error
+			}
+
+			reader, err = tusUpload.GetReader(ctx)
+			if err != nil {
+				s.portal.Logger().Error("Could not get tus file", zap.Error(err))
+				return err
+			}
+
 			err = s.PutFile(reader, upload.Protocol, dbHash)
 
 			if err != nil {
@@ -551,7 +577,7 @@ func (s *StorageServiceImpl) buildNewTusUploadTask(upload *models.TusUpload) (jo
 				return err
 			}
 
-			newUpload, err := s.CreateUpload(dbHash, upload.UploaderID, upload.UploaderIP, uint64(byteCount), upload.Protocol)
+			newUpload, err := s.CreateUpload(dbHash, mimeType, upload.UploaderID, upload.UploaderIP, uint64(byteCount), upload.Protocol)
 			if err != nil {
 				s.portal.Logger().Error("Could not create upload", zap.Error(err))
 				return err