Merge pull request #981 from trheyi/main

Refactor file upload handling in neo package

neo/api.go

@@ -248,15 +248,8 @@ func (neo *DSL) handleUpload(c *gin.Context) {
 		os.Remove(file.Filename)
 	}()
 
-	// Convert the header to a FileHeader
-	header, err := attachment.ToFileHeader(file.Header)
-	if err != nil {
-		c.JSON(400, gin.H{"message": err.Error(), "code": 400})
-		c.Done()
-		return
-	}
-
 	// Upload the file
+	header := attachment.GetHeader(c.Request.Header, file.Header)
 	res, err := manager.Upload(c.Request.Context(), header, reader, option)
 	if err != nil {
 		c.JSON(400, gin.H{"message": err.Error(), "code": 500})

neo/attachment/local/storage.go

@@ -1,6 +1,7 @@
 package local
 
 import (
+	"compress/gzip"
 	"context"
 	"crypto/sha256"
 	"fmt"
@@ -142,10 +143,21 @@ func (storage *Storage) MergeChunks(ctx context.Context, fileID string, totalChu
 // Reader read file from local storage
 func (storage *Storage) Reader(ctx context.Context, fileID string) (io.ReadCloser, error) {
 	fullpath := filepath.Join(storage.Path, fileID)
+
 	reader, err := os.Open(fullpath)
 	if err != nil {
 		return nil, err
 	}
+
+	// If the file is a gzip file, decompress it
+	if strings.HasSuffix(fileID, ".gz") {
+		reader, err := gzip.NewReader(reader)
+		if err != nil {
+			return nil, err
+		}
+		return reader, nil
+	}
+
 	return reader, nil
 }
 
@@ -159,7 +171,7 @@ func (storage *Storage) Download(ctx context.Context, fileID string) (io.ReadClo
 
 	// Try to detect content type from file extension
 	contentType := "application/octet-stream"
-	ext := filepath.Ext(path)
+	ext := filepath.Ext(strings.TrimSuffix(fileID, ".gz"))
 	switch strings.ToLower(ext) {
 	case ".txt":
 		contentType = "text/plain"
@@ -179,6 +191,28 @@ func (storage *Storage) Download(ctx context.Context, fileID string) (io.ReadClo
 		contentType = "image/gif"
 	case ".pdf":
 		contentType = "application/pdf"
+	case ".mp4":
+		contentType = "video/mp4"
+	case ".mp3":
+		contentType = "audio/mpeg"
+	case ".wav":
+		contentType = "audio/wav"
+	case ".ogg":
+		contentType = "audio/ogg"
+	case ".webm":
+		contentType = "video/webm"
+	case ".webp":
+		contentType = "image/webp"
+	case ".zip":
+	}
+
+	// If the file is a gzip file, decompress it
+	if strings.HasSuffix(fileID, ".gz") {
+		reader, err := gzip.NewReader(reader)
+		if err != nil {
+			return nil, "", err
+		}
+		return reader, contentType, nil
 	}
 
 	return reader, contentType, nil

neo/attachment/manager.go

@@ -9,11 +9,13 @@ import (
 	"io"
 	"mime"
 	"mime/multipart"
+	"net/http"
 	"net/textproto"
 	"os"
 	"path/filepath"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/yaoapp/yao/config"
@@ -23,6 +25,45 @@ import (
 
 // Managers the managers
 var Managers = map[string]*Manager{}
+var uploadChunks = sync.Map{}
+
+// UploadChunk is the chunk data
+type UploadChunk struct {
+	Last        int
+	Total       int64
+	Chunksize   int64
+	TotalChunks int64
+}
+
+// GetHeader gets the header from the file header and request header
+func GetHeader(requestHeader http.Header, fileHeader textproto.MIMEHeader) *FileHeader {
+
+	// Convert the header to a FileHeader
+	header := &FileHeader{FileHeader: &multipart.FileHeader{Header: make(map[string][]string)}}
+
+	for key, values := range fileHeader {
+		for _, value := range values {
+			header.Header.Set(key, value)
+		}
+	}
+
+	// Set Content-Sync, Content-Uid, Content-Range
+	if requestHeader.Get("Content-Sync") != "" {
+		header.Header.Set("Content-Sync", requestHeader.Get("Content-Sync"))
+	}
+
+	// Set Content-Uid
+	if requestHeader.Get("Content-Uid") != "" {
+		header.Header.Set("Content-Uid", requestHeader.Get("Content-Uid"))
+	}
+
+	// Set Content-Range
+	if requestHeader.Get("Content-Range") != "" {
+		header.Header.Set("Content-Range", requestHeader.Get("Content-Range"))
+	}
+
+	return header
+}
 
 // Register registers a global attachment manager
 func Register(name string, driver string, option ManagerOption) (*Manager, error) {
@@ -38,25 +79,6 @@ func Register(name string, driver string, option ManagerOption) (*Manager, error
 	return manager, nil
 }
 
-// ToFileHeader converts a multipart.FileHeader or textproto.MIMEHeader to a FileHeader
-func ToFileHeader(header interface{}) (*FileHeader, error) {
-
-	switch header := header.(type) {
-	case *multipart.FileHeader:
-		return &FileHeader{
-			FileHeader: header,
-		}, nil
-	case textproto.MIMEHeader:
-		return &FileHeader{
-			FileHeader: &multipart.FileHeader{
-				Header: header,
-			},
-		}, nil
-	default:
-		return nil, fmt.Errorf("invalid header type: %T", header)
-	}
-}
-
 // RegisterDefault registers a default attachment manager
 func RegisterDefault(name string) (*Manager, error) {
 
@@ -186,7 +208,7 @@ func New(option ManagerOption) (*Manager, error) {
 	return manager, nil
 }
 
-// Upload uploads a file
+// Upload uploads a file, Content-Sync must be true for chunked upload
 func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reader io.Reader, option UploadOption) (*File, error) {
 
 	file, err := manager.makeFile(fileheader, option)
@@ -196,36 +218,38 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade
 
 	// Handle chunked upload
 	if fileheader.IsChunk() {
-		start, _, total, err := fileheader.GetChunkInfo()
+		start, end, total, err := fileheader.GetChunkInfo()
 		if err != nil {
 			return nil, fmt.Errorf("invalid chunk info: %w", err)
 		}
 
-		// Calculate chunk index based on start position and a standard chunk size
+		// Store the chunk info
-		// We need to determine the standard chunk size from the first few chunks
+		chunkIndex := 0
-		standardChunkSize := int64(1024) // Default chunk size
+		if start == 0 {
-		if start > 0 {
+			chunksize := end - start + 1
-			// For non-first chunks, we can infer the standard chunk size
+			totalChunks := (total + chunksize - 1) / chunksize
-			// by looking at the start position
+			uploadChunks.LoadOrStore(file.ID, &UploadChunk{
-			if start%1024 == 0 {
+				Last:        chunkIndex,
-				standardChunkSize = 1024
+				Total:       total,
-			} else if start%2048 == 0 {
+				Chunksize:   chunksize,
-				standardChunkSize = 2048
+				TotalChunks: totalChunks,
-			} else if start%4096 == 0 {
+			})
-				standardChunkSize = 4096
-			} else {
-				// Try to infer from the start position
-				for size := int64(512); size <= 8192; size *= 2 {
-					if start%size == 0 {
-						standardChunkSize = size
-						break
-					}
-				}
-			}
 		}
 
-		chunkIndex := int(start / standardChunkSize)
+		// Update the chunk index
-		totalChunks := int((total + standardChunkSize - 1) / standardChunkSize) // Ceiling division
+		v, ok := uploadChunks.Load(file.ID)
+		if !ok {
+			return nil, fmt.Errorf("chunk data not found")
+		}
+
+		chunkdata := v.(*UploadChunk)
+
+		// Update the chunk index
+		if start != 0 {
+			chunkIndex = chunkdata.Last + 1
+			chunkdata.Last = chunkIndex
+			uploadChunks.Store(file.ID, chunkdata)
+		}
 
 		// Apply gzip compression if requested
 		if option.Gzip {
@@ -234,6 +258,7 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade
 				return nil, fmt.Errorf("failed to gzip chunk: %w", err)
 			}
 			reader = bytes.NewReader(compressed)
+
 		}
 
 		// Upload chunk
@@ -244,7 +269,7 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade
 
 		// If this is the last chunk, merge all chunks
 		if fileheader.Complete() {
-			err = manager.storage.MergeChunks(ctx, file.ID, totalChunks)
+			err = manager.storage.MergeChunks(ctx, file.ID, int(chunkdata.TotalChunks))
 			if err != nil {
 				return nil, err
 			}
@@ -256,6 +281,9 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade
 					return nil, err
 				}
 			}
+
+			// Remove the chunk data
+			uploadChunks.Delete(file.ID)
 		}
 
 		return file, nil
@@ -514,7 +542,11 @@ func (manager Manager) generateFileID(file *FileHeader, extension string, option
 		path = filepath.Join(path, option.AssistantID)
 	}
 
-	return filepath.Join(path, hash[:2], hash[2:4], hash) + extension, nil
+	id := filepath.Join(path, hash[:2], hash[2:4], hash) + extension
+	if option.Gzip {
+		id = id + ".gz"
+	}
+	return id, nil
 }
 
 // getSize converts the size to bytes

neo/attachment/manager_test.go

@@ -118,14 +118,8 @@ func TestManagerUpload(t *testing.T) {
 			t.Fatalf("Failed to read gzipped file: %v", err)
 		}
 
-		// Since we're storing compressed data, we need to decompress it
+		if string(downloadedContent) != content {
-		decompressed, err := Gunzip(downloadedContent)
+			t.Errorf("Expected content '%s', got '%s'", content, string(downloadedContent))
-		if err != nil {
-			t.Fatalf("Failed to decompress file: %v", err)
-		}
-
-		if string(decompressed) != content {
-			t.Errorf("Expected content '%s', got '%s'", content, string(decompressed))
 		}
 	})
 
@@ -158,10 +152,7 @@ func TestManagerUpload(t *testing.T) {
 				fmt.Sprintf("bytes %d-%d/%d", start, end, totalSize))
 			fileHeader.Header.Set("Content-Uid", "unique-file-id-123")
 
-			option := UploadOption{
+			option := UploadOption{UserID: "user123"}
-				UserID: "user123",
-			}
-
 			file, err := manager.Upload(context.Background(), fileHeader, bytes.NewReader(chunk), option)
 			if err != nil {
 				t.Fatalf("Failed to upload chunk starting at %d: %v", start, err)

neo/attachment/s3/storage.go

@@ -2,6 +2,7 @@ package s3
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"fmt"
 	"image"
@@ -229,6 +230,15 @@ func (storage *Storage) Reader(ctx context.Context, fileID string) (io.ReadClose
 		return nil, fmt.Errorf("failed to get file: %w", err)
 	}
 
+	// If the file is a gzip file, decompress it
+	if strings.HasSuffix(fileID, ".gz") {
+		reader, err := gzip.NewReader(result.Body)
+		if err != nil {
+			return nil, err
+		}
+		return reader, nil
+	}
+
 	return result.Body, nil
 }
 
@@ -254,6 +264,51 @@ func (storage *Storage) Download(ctx context.Context, fileID string) (io.ReadClo
 		contentType = *result.ContentType
 	}
 
+	// Try to detect content type from file extension
+	ext := filepath.Ext(strings.TrimSuffix(fileID, ".gz"))
+	switch strings.ToLower(ext) {
+	case ".txt":
+		contentType = "text/plain"
+	case ".html":
+		contentType = "text/html"
+	case ".css":
+		contentType = "text/css"
+	case ".js":
+		contentType = "application/javascript"
+	case ".json":
+		contentType = "application/json"
+	case ".jpg", ".jpeg":
+		contentType = "image/jpeg"
+	case ".png":
+		contentType = "image/png"
+	case ".gif":
+		contentType = "image/gif"
+	case ".pdf":
+		contentType = "application/pdf"
+	case ".mp4":
+		contentType = "video/mp4"
+	case ".mp3":
+		contentType = "audio/mpeg"
+	case ".wav":
+		contentType = "audio/wav"
+	case ".ogg":
+		contentType = "audio/ogg"
+	case ".webm":
+		contentType = "video/webm"
+	case ".webp":
+		contentType = "image/webp"
+	case ".zip":
+	}
+
+	// If the file is a gzip file, decompress it
+	if strings.HasSuffix(fileID, ".gz") {
+		reader, err := gzip.NewReader(result.Body)
+		if err != nil {
+			return nil, "", err
+		}
+		return reader, contentType, nil
+	}
+
 	return result.Body, contentType, nil
 }