diff --git a/pkg/azurestore/azurestore.go b/pkg/azurestore/azurestore.go index 1a7051ba0..0768bb421 100644 --- a/pkg/azurestore/azurestore.go +++ b/pkg/azurestore/azurestore.go @@ -1,13 +1,14 @@ package azurestore import ( - "bufio" "bytes" "context" - "encoding/binary" "encoding/json" + "errors" "fmt" "io" + "io/fs" + "os" "strings" "github.com/tus/tusd/v2/internal/uid" @@ -18,6 +19,11 @@ type AzureStore struct { Service AzService ObjectPrefix string Container string + + // TemporaryDirectory is the path where AzureStore will create temporary files + // on disk during the upload. An empty string ("", the default value) will + // cause AzureStore to use the operating system's default temporary directory. + TemporaryDirectory string } type AzUpload struct { @@ -25,6 +31,8 @@ type AzUpload struct { InfoBlob AzBlob BlockBlob AzBlob InfoHandler *handler.FileInfo + + tempDir string } func New(service AzService) *AzureStore { @@ -73,6 +81,7 @@ func (store AzureStore) NewUpload(ctx context.Context, info handler.FileInfo) (h InfoHandler: &info, InfoBlob: infoBlob, BlockBlob: blockBlob, + tempDir: store.TemporaryDirectory, } err = azUpload.writeInfo(ctx) @@ -128,6 +137,7 @@ func (store AzureStore) GetUpload(ctx context.Context, id string) (handler.Uploa InfoHandler: &info, InfoBlob: infoBlob, BlockBlob: blockBlob, + tempDir: store.TemporaryDirectory, }, nil } @@ -140,21 +150,38 @@ func (store AzureStore) AsLengthDeclarableUpload(upload handler.Upload) handler. } func (upload *AzUpload) WriteChunk(ctx context.Context, offset int64, src io.Reader) (int64, error) { - r := bufio.NewReader(src) - buf := new(bytes.Buffer) - n, err := r.WriteTo(buf) + // Create a temporary file for holding the uploaded data + file, err := os.CreateTemp(upload.tempDir, "tusd-az-tmp-") + if err != nil { + return 0, err + } + defer os.Remove(file.Name()) + + // Copy the entire request body into the file + n, err := io.Copy(file, src) if err != nil { + file.Close() return 0, err } - chunkSize := int64(binary.Size(buf.Bytes())) - if chunkSize > MaxBlockBlobChunkSize { - return 0, fmt.Errorf("azurestore: Chunk of size %v too large. Max chunk size is %v", chunkSize, MaxBlockBlobChunkSize) + // Seek to the beginning + if _, err := file.Seek(0, 0); err != nil { + file.Close() + return 0, err + } + + if n > MaxBlockBlobChunkSize { + file.Close() + return 0, fmt.Errorf("azurestore: Chunk of size %v too large. Max chunk size is %v", n, MaxBlockBlobChunkSize) } - re := bytes.NewReader(buf.Bytes()) - err = upload.BlockBlob.Upload(ctx, re) + err = upload.BlockBlob.Upload(ctx, file) if err != nil { + file.Close() + return 0, err + } + + if err := file.Close(); err != nil && !errors.Is(err, fs.ErrClosed) { return 0, err } diff --git a/pkg/azurestore/azurestore_test.go b/pkg/azurestore/azurestore_test.go index 3c35ee256..023d59801 100644 --- a/pkg/azurestore/azurestore_test.go +++ b/pkg/azurestore/azurestore_test.go @@ -290,7 +290,12 @@ func TestWriteChunk(t *testing.T) { infoBlob.EXPECT().Download(ctx).Return(newReadCloser(data), nil).Times(1), service.EXPECT().NewBlob(ctx, mockID).Return(blockBlob, nil).Times(1), blockBlob.EXPECT().GetOffset(ctx).Return(offset, nil).Times(1), - blockBlob.EXPECT().Upload(ctx, bytes.NewReader([]byte(mockReaderData))).Return(nil).Times(1), + blockBlob.EXPECT().Upload(ctx, gomock.Any()).DoAndReturn(func(_ context.Context, reader io.ReadSeeker) error { + actual, err := io.ReadAll(reader) + assert.Nil(err) + assert.Equal(mockReaderData, string(actual)) + return nil + }).Times(1), ) upload, err := store.GetUpload(ctx, mockID)