From 88dd6019419488b7efd0cd41d47fb29eb465c1a0 Mon Sep 17 00:00:00 2001
From: Brad Davidson <brad.davidson@rancher.com>
Date: Tue, 16 Feb 2021 15:15:16 -0800
Subject: [PATCH] Limit zstd decoder memory

Signed-off-by: Brad Davidson <brad.davidson@rancher.com>
---
 pkg/agent/containerd/containerd.go |  3 ++-
 pkg/untar/untar.go                 | 15 +++++++++++++--
 scripts/package-cli                |  2 +-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/pkg/agent/containerd/containerd.go b/pkg/agent/containerd/containerd.go
index 7fd18c3f87..c77f961c3b 100644
--- a/pkg/agent/containerd/containerd.go
+++ b/pkg/agent/containerd/containerd.go
@@ -27,6 +27,7 @@ import (
 	"github.com/rancher/k3s/pkg/agent/templates"
 	util2 "github.com/rancher/k3s/pkg/agent/util"
 	"github.com/rancher/k3s/pkg/daemons/config"
+	"github.com/rancher/k3s/pkg/untar"
 	"github.com/rancher/k3s/pkg/version"
 	"github.com/rancher/wrangler/pkg/merr"
 	"github.com/sirupsen/logrus"
@@ -229,7 +230,7 @@ func preloadFile(ctx context.Context, cfg *config.Node, client *containerd.Clien
 		defer zr.Close()
 		imageReader = zr
 	case util2.HasSuffixI(filePath, "tar.zst", ".tzst"):
-		zr, err := zstd.NewReader(file)
+		zr, err := zstd.NewReader(file, zstd.WithDecoderMaxMemory(untar.MaxDecoderMemory))
 		if err != nil {
 			return err
 		}
diff --git a/pkg/untar/untar.go b/pkg/untar/untar.go
index 793259f7ee..16b03c1262 100644
--- a/pkg/untar/untar.go
+++ b/pkg/untar/untar.go
@@ -18,6 +18,17 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+const (
+	// The zstd decoder will attempt to use up to 1GB memory for streaming operations by default,
+	// which is excessive and will OOM low-memory devices.
+	// NOTE: This must be at least as large as the window size used when compressing tarballs, or you
+	// will see a "window size exceeded" error when decompressing. The zstd CLI tool uses 4MB by
+	// default; the --long option defaults to 27 or 128M, which is still too much for a Pi3. 32MB
+	// (--long=25) has been tested to work acceptably while still compressing by an additional 3-6% on
+	// our datasets.
+	MaxDecoderMemory = 1 << 25
+)
+
 // TODO(bradfitz): this was copied from x/build/cmd/buildlet/buildlet.go
 // but there were some buildlet-specific bits in there, so the code is
 // forked for now.  Unfork and add some opts arguments here, so the
@@ -38,9 +49,9 @@ func untar(r io.Reader, dir string) (err error) {
 			logrus.Printf("error extracting tarball into %s after %d files, %d dirs, %v: %v", dir, nFiles, len(madeDir), td, err)
 		}
 	}()
-	zr, err := zstd.NewReader(r)
+	zr, err := zstd.NewReader(r, zstd.WithDecoderMaxMemory(MaxDecoderMemory))
 	if err != nil {
-		return fmt.Errorf("requires zstd-compressed body: %v", err)
+		return fmt.Errorf("error extracting zstd-compressed body: %v", err)
 	}
 	defer zr.Close()
 	tr := tar.NewReader(zr)
diff --git a/scripts/package-cli b/scripts/package-cli
index cb2839414c..776ae8c5ec 100755
--- a/scripts/package-cli
+++ b/scripts/package-cli
@@ -36,7 +36,7 @@ mkdir -p dist/artifacts
 )
 
 tar cvf ./build/out/data.tar --exclude ./bin/hyperkube ./bin ./etc 
-zstd -v -T0 -16 -f --long --rm ./build/out/data.tar -o ./build/out/data.tar.zst
+zstd -v -T0 -16 -f --long=25 --rm ./build/out/data.tar -o ./build/out/data.tar.zst
 HASH=$(sha256sum ./build/out/data.tar.zst | awk '{print $1}')
 
 cp ./build/out/data.tar.zst ./build/data/${HASH}.tar.zst