From f3be950ba76274ad0a3542d1a8f1f6cc532017df Mon Sep 17 00:00:00 2001 From: Derrick Hammer 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