From 10ddf50bf36d59dd3dbaddda0ead1b500b5772b1 Mon Sep 17 00:00:00 2001 From: Quentin Dufournet Date: Tue, 3 Jun 2025 15:39:55 +0200 Subject: [PATCH] fix: prevent OOM on large upload/download - Refactored Upload function to stream file content instead of loading it all into memory (fixes OOM issue for large files) - Added log.IsVerbose() to expose verbose state from log package - Updated Download function to dump HTTP response only when verbose mode is enabled, to avoid potential OOM on large responses Changelog: None Ticket: None Signed-off-by: Quentin Dufournet --- client/deviceconnect/client.go | 67 ++++++++++++++++++++++------------ licenses.csv | 1 + log/log.go | 4 ++ 3 files changed, 48 insertions(+), 24 deletions(-) diff --git a/client/deviceconnect/client.go b/client/deviceconnect/client.go index 77f3ff11..9aed66a9 100644 --- a/client/deviceconnect/client.go +++ b/client/deviceconnect/client.go @@ -14,7 +14,6 @@ package deviceconnect import ( - "bytes" "context" "crypto/tls" "encoding/json" @@ -26,6 +25,7 @@ import ( "net/http/httputil" "net/url" "os" + "path/filepath" "strconv" "strings" "sync" @@ -257,36 +257,53 @@ func NewDeviceConnectError(errCode int, r io.Reader) *DeviceConnectError { } func (c *Client) Upload(sourcePath string, deviceSpec *DeviceSpec) error { - body := &bytes.Buffer{} - writer := multipart.NewWriter(body) file, err := os.Open(sourcePath) if err != nil { return err } - fi, err := os.Stat(sourcePath) - if err != nil { - return err - } - log.Verbf("Uploading the file to %s\n", deviceSpec.DevicePath) - if err = writer.WriteField("path", deviceSpec.DevicePath); err != nil { - return err - } - part, err := writer.CreateFormFile("file", sourcePath) + defer file.Close() + + fi, err := file.Stat() if err != nil { return err } - if _, err = io.Copy(part, file); err != nil { - return err - } - if err = writer.WriteField("mode", fmt.Sprintf("%o", fi.Mode())); err != nil { - return err - } - if err = writer.Close(); err != nil { - return err - } + + pr, pw := io.Pipe() + writer := multipart.NewWriter(pw) + + go func() { + defer pw.Close() + + if err := writer.WriteField("path", deviceSpec.DevicePath); err != nil { + pw.CloseWithError(err) + return + } + + part, err := writer.CreateFormFile("file", filepath.Base(sourcePath)) + if err != nil { + pw.CloseWithError(err) + return + } + + if _, err := io.Copy(part, file); err != nil { + pw.CloseWithError(err) + return + } + + if err := writer.WriteField("mode", fmt.Sprintf("%o", fi.Mode())); err != nil { + pw.CloseWithError(err) + return + } + + if err := writer.Close(); err != nil { + pw.CloseWithError(err) + return + } + }() + req, err := http.NewRequest(http.MethodPut, c.url+fileUploadURL+"devices/"+deviceSpec.DeviceID+"/upload", - body) + pr) if err != nil { return err } @@ -344,8 +361,10 @@ func (c *Client) Download(deviceSpec *DeviceSpec, sourcePath string) error { } defer resp.Body.Close() - rspDump, _ := httputil.DumpResponse(resp, true) - log.Verbf("Response: \n%v\n", string(rspDump)) + if log.IsVerbose() { + rspDump, _ := httputil.DumpResponse(resp, true) + log.Verbf("Response: \n%v\n", string(rspDump)) + } switch resp.StatusCode { case http.StatusOK: diff --git a/licenses.csv b/licenses.csv index 12966b0e..9a561ddf 100644 --- a/licenses.csv +++ b/licenses.csv @@ -17,6 +17,7 @@ github.com/minio/sha256-simd,Apache-2.0 github.com/mitchellh/mapstructure,MIT github.com/pelletier/go-toml/v2,MIT github.com/pkg/errors,BSD-2-Clause +github.com/remyoudompheng/go-liblzma,BSD-3-Clause github.com/rivo/uniseg,MIT github.com/sagikazarmark/slog-shim,BSD-3-Clause github.com/spf13/afero,Apache-2.0 diff --git a/log/log.go b/log/log.go index 321044e4..8bd3311f 100644 --- a/log/log.go +++ b/log/log.go @@ -26,6 +26,10 @@ func Setup(verb bool) { verbose = verb } +func IsVerbose() bool { + return verbose +} + func Err(msg string) { fmt.Fprintln(os.Stderr, msg) }