From fb876c4f81f9a6de5147728eea1ee433f922ec0a Mon Sep 17 00:00:00 2001
From: Sai Sridhar <tarrasridhar1154@gmail.com>
Date: Fri, 10 Jul 2026 00:38:32 +0530
Subject: [PATCH] fix: stream ZIP downloads directly to response to prevent OOM
 (#338)

Replace bytes.Buffer accumulation in DownloadFlowFile with direct streaming
to c.Writer for both single-directory and multi-path ZIP cases. Headers are
set before writing begins; Content-Length is omitted since the archive size
is not known upfront. Removes the bytes import that is no longer needed.
---
 backend/pkg/server/services/flow_files.go | 37 +++++++++--------------
 1 file changed, 15 insertions(+), 22 deletions(-)

diff --git a/backend/pkg/server/services/flow_files.go b/backend/pkg/server/services/flow_files.go
index 10cb5d5a8..3bdbf0a25 100644
--- a/backend/pkg/server/services/flow_files.go
+++ b/backend/pkg/server/services/flow_files.go
@@ -1,7 +1,6 @@
 package services
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -547,43 +546,37 @@ func (s *FlowFileService) DownloadFlowFile(c *gin.Context) {
 	}
 
 	// Single directory → ZIP with paths relative to that directory (backward-compat).
-	// The archive is buffered so an explicit Content-Length can be set; this allows
-	// clients (including Swagger UI) to recognise and download the file correctly.
+	// Written directly to the response writer to avoid buffering the entire archive
+	// in memory, which can OOM the server for large directories.
 	if len(entries) == 1 && entries[0].info.IsDir() {
 		name := filepath.Base(entries[0].localPath)
-		var buf bytes.Buffer
-		if err := flowfiles.ZipDirectory(&buf, entries[0].localPath); err != nil {
+		c.Header("Content-Type", "application/zip")
+		c.Header("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{
+			"filename": name + ".zip",
+		}))
+		c.Status(http.StatusOK)
+		if err := flowfiles.ZipDirectory(c.Writer, entries[0].localPath); err != nil {
 			logger.FromContext(c).WithError(err).WithField("flow_id", flowID).Error("error creating ZIP archive")
-			response.Error(c, response.ErrInternal, err)
 			return
 		}
-		c.DataFromReader(http.StatusOK, int64(buf.Len()), "application/zip", &buf,
-			map[string]string{
-				"Content-Disposition": mime.FormatMediaType("attachment", map[string]string{
-					"filename": name + ".zip",
-				}),
-			})
 		return
 	}
 
 	// Multiple paths (any mix of files and directories) → ZIP with cache-relative paths.
-	// Buffered for the same reason as above.
+	// Streamed directly to the response writer for the same reason as above.
 	relPaths := make([]string, 0, len(entries))
 	for _, e := range entries {
 		relPaths = append(relPaths, filepath.ToSlash(e.reqPath))
 	}
-	var buf bytes.Buffer
-	if err := flowfiles.ZipRelativePaths(&buf, s.flowDataDir(flowID), relPaths); err != nil {
+	c.Header("Content-Type", "application/zip")
+	c.Header("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{
+		"filename": "download.zip",
+	}))
+	c.Status(http.StatusOK)
+	if err := flowfiles.ZipRelativePaths(c.Writer, s.flowDataDir(flowID), relPaths); err != nil {
 		logger.FromContext(c).WithError(err).WithField("flow_id", flowID).Error("error creating ZIP archive")
-		response.Error(c, response.ErrInternal, err)
 		return
 	}
-	c.DataFromReader(http.StatusOK, int64(buf.Len()), "application/zip", &buf,
-		map[string]string{
-			"Content-Disposition": mime.FormatMediaType("attachment", map[string]string{
-				"filename": "download.zip",
-			}),
-		})
 }
 
 // PullFlowFiles is a function to sync one or more paths from the container into the local cache