diff --git a/beamsync/server.go b/beamsync/server.go index 169ecd5..01d7039 100644 --- a/beamsync/server.go +++ b/beamsync/server.go @@ -16,6 +16,7 @@ import ( "mime/multipart" "net" "net/http" + "net/textproto" "os" "path/filepath" "runtime/debug" @@ -428,6 +429,25 @@ const largeFileThreshold = 64 * 1024 * 1024 // 64 MB const uploadIntegrityHeader = "X-BeamSync-File-SHA256" +func uploadBufferInitialCapacity(filename string, fileSizes map[string]int64, partHeader textproto.MIMEHeader, requestContentLength int64) int { + if size, ok := fileSizes[filename]; ok && size > 0 { + if size < largeFileThreshold { + return int(size) + } + return largeFileThreshold + } + if cl, _ := strconv.ParseInt(partHeader.Get("Content-Length"), 10, 64); cl > 0 { + if cl < largeFileThreshold { + return int(cl) + } + return largeFileThreshold + } + if requestContentLength > 0 && requestContentLength < largeFileThreshold { + return int(requestContentLength) + } + return 0 +} + func sha256Hex(data []byte) string { sum := sha256.Sum256(data) return hex.EncodeToString(sum[:]) @@ -1031,7 +1051,9 @@ func StartServer(uploadDir string, startPort int, settings TransferSettings, cal // Read up to largeFileThreshold bytes to determine dispatch strategy. var buf bytes.Buffer - buf.Grow(largeFileThreshold) + if initialCapacity := uploadBufferInitialCapacity(filename, fileSizes, part.Header, r.ContentLength); initialCapacity > 0 { + buf.Grow(initialCapacity) + } readLimit := int64(largeFileThreshold) n, readErr := io.CopyN(&buf, part, readLimit) diff --git a/beamsync/server_test.go b/beamsync/server_test.go index 30b885b..e523e51 100644 --- a/beamsync/server_test.go +++ b/beamsync/server_test.go @@ -11,6 +11,7 @@ import ( "net" "net/http" "net/http/httptest" + "net/textproto" "os" "path/filepath" "runtime" @@ -119,6 +120,48 @@ func TestSetCORSHeaders(t *testing.T) { } } +func TestUploadBufferInitialCapacityUsesSmallManifestSize(t *testing.T) { + got := uploadBufferInitialCapacity( + "tiny.txt", + map[string]int64{"tiny.txt": 1024}, + textproto.MIMEHeader{}, + largeFileThreshold+1024, + ) + + if got != 1024 { + t.Fatalf("initial capacity = %d, want manifest size 1024", got) + } +} + +func TestUploadBufferInitialCapacityCapsLargeManifestAtThreshold(t *testing.T) { + got := uploadBufferInitialCapacity( + "movie.mp4", + map[string]int64{"movie.mp4": largeFileThreshold + 1}, + textproto.MIMEHeader{}, + 0, + ) + + if got != largeFileThreshold { + t.Fatalf("initial capacity = %d, want largeFileThreshold", got) + } +} + +func TestUploadBufferInitialCapacityFallsBackToRequestLength(t *testing.T) { + got := uploadBufferInitialCapacity("unknown.txt", nil, textproto.MIMEHeader{}, 4096) + + if got != 4096 { + t.Fatalf("initial capacity = %d, want request content length 4096", got) + } +} + +func TestUploadBufferInitialCapacityAvoidsUnknownLargePreallocation(t *testing.T) { + got := uploadBufferInitialCapacity("unknown.bin", nil, textproto.MIMEHeader{}, largeFileThreshold+1024) + + if got != 0 { + t.Fatalf("initial capacity = %d, want 0 for unknown large upload", got) + } +} + func TestCopyChunkedCopiesDataAndReportsCount(t *testing.T) { payload := strings.Repeat("beam-sync", 1024) var dst bytes.Buffer