Skip to content

Commit af3df84

Browse files
authored
Merge pull request #339 from mason5052/codex/issue-338-stream-zip-downloads
fix: stream ZIP downloads without buffering archives
2 parents 2ef7076 + 41524bd commit af3df84

4 files changed

Lines changed: 156 additions & 49 deletions

File tree

backend/pkg/server/services/flow_files.go

Lines changed: 11 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package services
22

33
import (
4-
"bytes"
54
"context"
65
"errors"
76
"fmt"
@@ -547,43 +546,29 @@ func (s *FlowFileService) DownloadFlowFile(c *gin.Context) {
547546
}
548547

549548
// Single directory → ZIP with paths relative to that directory (backward-compat).
550-
// The archive is buffered so an explicit Content-Length can be set; this allows
551-
// clients (including Swagger UI) to recognise and download the file correctly.
549+
// The archive is streamed directly to the response writer so the whole ZIP is
550+
// never held in memory; the response uses chunked transfer encoding.
552551
if len(entries) == 1 && entries[0].info.IsDir() {
553552
name := filepath.Base(entries[0].localPath)
554-
var buf bytes.Buffer
555-
if err := flowfiles.ZipDirectory(&buf, entries[0].localPath); err != nil {
556-
logger.FromContext(c).WithError(err).WithField("flow_id", flowID).Error("error creating ZIP archive")
557-
response.Error(c, response.ErrInternal, err)
558-
return
553+
if err := streamZipArchive(c, name+".zip", func(w io.Writer) error {
554+
return flowfiles.ZipDirectory(w, entries[0].localPath)
555+
}); err != nil {
556+
logger.FromContext(c).WithError(err).WithField("flow_id", flowID).Error("error streaming ZIP archive")
559557
}
560-
c.DataFromReader(http.StatusOK, int64(buf.Len()), "application/zip", &buf,
561-
map[string]string{
562-
"Content-Disposition": mime.FormatMediaType("attachment", map[string]string{
563-
"filename": name + ".zip",
564-
}),
565-
})
566558
return
567559
}
568560

569561
// Multiple paths (any mix of files and directories) → ZIP with cache-relative paths.
570-
// Buffered for the same reason as above.
562+
// Streamed directly to the response writer for the same reason as above.
571563
relPaths := make([]string, 0, len(entries))
572564
for _, e := range entries {
573565
relPaths = append(relPaths, filepath.ToSlash(e.reqPath))
574566
}
575-
var buf bytes.Buffer
576-
if err := flowfiles.ZipRelativePaths(&buf, s.flowDataDir(flowID), relPaths); err != nil {
577-
logger.FromContext(c).WithError(err).WithField("flow_id", flowID).Error("error creating ZIP archive")
578-
response.Error(c, response.ErrInternal, err)
579-
return
567+
if err := streamZipArchive(c, "download.zip", func(w io.Writer) error {
568+
return flowfiles.ZipRelativePaths(w, s.flowDataDir(flowID), relPaths)
569+
}); err != nil {
570+
logger.FromContext(c).WithError(err).WithField("flow_id", flowID).Error("error streaming ZIP archive")
580571
}
581-
c.DataFromReader(http.StatusOK, int64(buf.Len()), "application/zip", &buf,
582-
map[string]string{
583-
"Content-Disposition": mime.FormatMediaType("attachment", map[string]string{
584-
"filename": "download.zip",
585-
}),
586-
})
587572
}
588573

589574
// PullFlowFiles is a function to sync one or more paths from the container into the local cache

backend/pkg/server/services/flow_files_test.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2501,6 +2501,9 @@ func TestFlowFileService_DownloadFlowFileScenarios(t *testing.T) {
25012501
assert.Contains(t, w.Header().Get("Content-Disposition"), tt.wantDispContains)
25022502
}
25032503
if tt.wantZipEntries != nil {
2504+
// A streamed ZIP download must not buffer the whole archive, so it
2505+
// must not carry a Content-Length computed from a full buffer.
2506+
assert.Empty(t, w.Header().Get("Content-Length"))
25042507
zr, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len()))
25052508
require.NoError(t, err)
25062509
got := map[string]string{}

backend/pkg/server/services/resources.go

Lines changed: 53 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
package services
22

33
import (
4-
"bytes"
54
"context"
65
"errors"
76
"fmt"
7+
"io"
88
"mime"
99
"net/http"
1010
"os"
@@ -1930,18 +1930,11 @@ func (s *ResourceService) DownloadResource(c *gin.Context) {
19301930
}
19311931

19321932
dirName := path.Base(e.rec.Path)
1933-
var buf bytes.Buffer
1934-
if err := resources.ZipResources(&buf, zipEntries); err != nil {
1935-
logger.FromContext(c).WithError(err).Error("error creating zip archive for download")
1936-
response.Error(c, response.ErrInternal, err)
1937-
return
1933+
if err := streamZipArchive(c, dirName+".zip", func(w io.Writer) error {
1934+
return resources.ZipResources(w, zipEntries)
1935+
}); err != nil {
1936+
logger.FromContext(c).WithError(err).Error("error streaming zip archive for download")
19381937
}
1939-
c.DataFromReader(http.StatusOK, int64(buf.Len()), "application/zip", &buf,
1940-
map[string]string{
1941-
"Content-Disposition": mime.FormatMediaType("attachment", map[string]string{
1942-
"filename": dirName + ".zip",
1943-
}),
1944-
})
19451938
return
19461939
}
19471940

@@ -1983,18 +1976,55 @@ func (s *ResourceService) DownloadResource(c *gin.Context) {
19831976
}
19841977
}
19851978

1986-
var buf bytes.Buffer
1987-
if err := resources.ZipResources(&buf, zipEntries); err != nil {
1988-
logger.FromContext(c).WithError(err).Error("error creating zip archive for download")
1989-
response.Error(c, response.ErrInternal, err)
1990-
return
1979+
if err := streamZipArchive(c, "download.zip", func(w io.Writer) error {
1980+
return resources.ZipResources(w, zipEntries)
1981+
}); err != nil {
1982+
logger.FromContext(c).WithError(err).Error("error streaming zip archive for download")
19911983
}
1992-
c.DataFromReader(http.StatusOK, int64(buf.Len()), "application/zip", &buf,
1993-
map[string]string{
1994-
"Content-Disposition": mime.FormatMediaType("attachment", map[string]string{
1995-
"filename": "download.zip",
1996-
}),
1997-
})
1984+
}
1985+
1986+
// zipStreamWriter defers committing the ZIP response status and headers until the
1987+
// first byte is written. That lets streamZipArchive turn a build failure that
1988+
// happens before any output into a normal structured error response, while a
1989+
// mid-stream failure (after the 200 status is already on the wire) is aborted.
1990+
type zipStreamWriter struct {
1991+
c *gin.Context
1992+
filename string
1993+
started bool
1994+
}
1995+
1996+
func (zw *zipStreamWriter) Write(p []byte) (int, error) {
1997+
if !zw.started {
1998+
zw.started = true
1999+
zw.c.Header("Content-Type", "application/zip")
2000+
zw.c.Header("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{
2001+
"filename": zw.filename,
2002+
}))
2003+
zw.c.Status(http.StatusOK)
2004+
}
2005+
return zw.c.Writer.Write(p)
2006+
}
2007+
2008+
// streamZipArchive streams a ZIP archive to the client as build writes it, instead
2009+
// of buffering the entire archive in memory before sending it. Memory stays
2010+
// proportional to a single file copy buffer rather than the full archive size.
2011+
//
2012+
// Response headers and the 200 status are committed lazily, on the first byte
2013+
// build writes (see zipStreamWriter). If build fails before writing anything
2014+
// (e.g. a blob that cannot be opened), nothing has been committed and a normal
2015+
// structured error response is returned. If it fails mid-stream the 200 status is
2016+
// already on the wire, so the partial response is aborted; either way the error is
2017+
// returned for the caller to log with its own context.
2018+
func streamZipArchive(c *gin.Context, filename string, build func(w io.Writer) error) error {
2019+
if err := build(&zipStreamWriter{c: c, filename: filename}); err != nil {
2020+
if c.Writer.Written() {
2021+
c.Abort()
2022+
} else {
2023+
response.Error(c, response.ErrInternal, err)
2024+
}
2025+
return err
2026+
}
2027+
return nil
19982028
}
19992029

20002030
// ---- helper methods --------------------------------------------------------

backend/pkg/server/services/resources_test.go

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"crypto/md5"
88
"encoding/hex"
99
"encoding/json"
10+
"errors"
1011
"io"
1112
"mime/multipart"
1213
"net/http"
@@ -1351,6 +1352,9 @@ func TestResourceService_DownloadResourceScenarios(t *testing.T) {
13511352
assert.Contains(t, w.Header().Get("Content-Disposition"), tt.wantDispContains)
13521353
}
13531354
if tt.wantZipEntries != nil {
1355+
// A streamed ZIP download must not buffer the whole archive, so it
1356+
// must not carry a Content-Length computed from a full buffer.
1357+
assert.Empty(t, w.Header().Get("Content-Length"))
13541358
zr, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len()))
13551359
require.NoError(t, err)
13561360
got := map[string]string{}
@@ -1370,6 +1374,91 @@ func TestResourceService_DownloadResourceScenarios(t *testing.T) {
13701374
}
13711375
}
13721376

1377+
// TestStreamZipArchive verifies the shared streaming helper writes a valid ZIP
1378+
// straight to the response writer without buffering the whole archive (no
1379+
// Content-Length is emitted) and that a build error mid-stream propagates to the
1380+
// caller and aborts the request.
1381+
func TestStreamZipArchive(t *testing.T) {
1382+
gin.SetMode(gin.TestMode)
1383+
1384+
t.Run("streams archive without content-length", func(t *testing.T) {
1385+
w := httptest.NewRecorder()
1386+
c, _ := gin.CreateTestContext(w)
1387+
c.Request = httptest.NewRequest(http.MethodGet, "/download", nil)
1388+
1389+
err := streamZipArchive(c, "out.zip", func(zw io.Writer) error {
1390+
z := zip.NewWriter(zw)
1391+
f, createErr := z.Create("hello.txt")
1392+
require.NoError(t, createErr)
1393+
if _, writeErr := f.Write([]byte("hello world")); writeErr != nil {
1394+
return writeErr
1395+
}
1396+
return z.Close()
1397+
})
1398+
1399+
require.NoError(t, err)
1400+
assert.False(t, c.IsAborted())
1401+
assert.Equal(t, "application/zip", w.Header().Get("Content-Type"))
1402+
assert.Contains(t, w.Header().Get("Content-Disposition"), "out.zip")
1403+
// The whole archive was never buffered, so no buffer-derived length exists.
1404+
assert.Empty(t, w.Header().Get("Content-Length"))
1405+
1406+
zr, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len()))
1407+
require.NoError(t, err)
1408+
require.Len(t, zr.File, 1)
1409+
assert.Equal(t, "hello.txt", zr.File[0].Name)
1410+
rc, err := zr.File[0].Open()
1411+
require.NoError(t, err)
1412+
data, err := io.ReadAll(rc)
1413+
require.NoError(t, err)
1414+
require.NoError(t, rc.Close())
1415+
assert.Equal(t, "hello world", string(data))
1416+
})
1417+
1418+
t.Run("propagates build error after partial stream and aborts", func(t *testing.T) {
1419+
w := httptest.NewRecorder()
1420+
c, _ := gin.CreateTestContext(w)
1421+
c.Request = httptest.NewRequest(http.MethodGet, "/download", nil)
1422+
1423+
sentinel := errors.New("source reader failed mid-stream")
1424+
err := streamZipArchive(c, "out.zip", func(zw io.Writer) error {
1425+
z := zip.NewWriter(zw)
1426+
f, createErr := z.Create("partial.txt")
1427+
require.NoError(t, createErr)
1428+
if _, writeErr := f.Write([]byte("partial data")); writeErr != nil {
1429+
return writeErr
1430+
}
1431+
// Flush bytes to the response, then fail as a slow/erroring source would.
1432+
if flushErr := z.Flush(); flushErr != nil {
1433+
return flushErr
1434+
}
1435+
return sentinel
1436+
})
1437+
1438+
require.ErrorIs(t, err, sentinel)
1439+
assert.True(t, c.IsAborted())
1440+
})
1441+
1442+
t.Run("returns structured error when build fails before writing", func(t *testing.T) {
1443+
w := httptest.NewRecorder()
1444+
c, _ := gin.CreateTestContext(w)
1445+
c.Request = httptest.NewRequest(http.MethodGet, "/download", nil)
1446+
1447+
sentinel := errors.New("blob open failed before any write")
1448+
err := streamZipArchive(c, "out.zip", func(zw io.Writer) error {
1449+
// Fail immediately, before writing any archive bytes.
1450+
return sentinel
1451+
})
1452+
1453+
require.ErrorIs(t, err, sentinel)
1454+
// Nothing was streamed, so a normal (non-200, non-zip) error response is
1455+
// sent instead of a truncated 200.
1456+
assert.Equal(t, http.StatusInternalServerError, w.Code)
1457+
assert.NotEqual(t, "application/zip", w.Header().Get("Content-Type"))
1458+
assert.True(t, c.IsAborted())
1459+
})
1460+
}
1461+
13731462
func TestResourceService_DeleteResourceScenarios(t *testing.T) {
13741463
type seed struct {
13751464
userID uint64

0 commit comments

Comments
 (0)