Limit the response size of ExecSync

Kazuyoshi Kato · Kazuyoshi Kato · commit 49ca87d72700 · 2022-05-31T22:21:35.000Z
Signed-off-by: Kazuyoshi Kato &lt;katokazu@amazon.com&gt;
diff --git a/pkg/cri/server/container_execsync.go b/pkg/cri/server/container_execsync.go
@@ -19,6 +19,7 @@ package server
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"syscall"
@@ -38,14 +39,48 @@ import (
 	cioutil "github.com/containerd/containerd/pkg/ioutil"
 )
 
+type cappedWriter struct {
+	w      io.WriteCloser
+	remain int
+}
+
+var errNoRemain = errors.New("no more space to write")
+
+func (cw *cappedWriter) Write(p []byte) (int, error) {
+	if cw.remain <= 0 {
+		return 0, errNoRemain
+	}
+
+	end := cw.remain
+	if end > len(p) {
+		end = len(p)
+	}
+	written, err := cw.w.Write(p[0:end])
+	cw.remain -= written
+
+	if err != nil {
+		return written, err
+	}
+	if written < len(p) {
+		return written, errNoRemain
+	}
+	return written, nil
+}
+
+func (cw *cappedWriter) Close() error {
+	return cw.w.Close()
+}
+
 // ExecSync executes a command in the container, and returns the stdout output.
 // If command exits with a non-zero exit code, an error is returned.
 func (c *criService) ExecSync(ctx context.Context, r *runtime.ExecSyncRequest) (*runtime.ExecSyncResponse, error) {
+	const maxStreamSize = 1024 * 1024 * 16
+
 	var stdout, stderr bytes.Buffer
 	exitCode, err := c.execInContainer(ctx, r.GetContainerId(), execOptions{
 		cmd:     r.GetCmd(),
-		stdout:  cioutil.NewNopWriteCloser(&stdout),
-		stderr:  cioutil.NewNopWriteCloser(&stderr),
+		stdout:  &cappedWriter{w: cioutil.NewNopWriteCloser(&stdout), remain: maxStreamSize},
+		stderr:  &cappedWriter{w: cioutil.NewNopWriteCloser(&stderr), remain: maxStreamSize},
 		timeout: time.Duration(r.GetTimeout()) * time.Second,
 	})
 	if err != nil {
diff --git a/pkg/cri/server/container_execsync_test.go b/pkg/cri/server/container_execsync_test.go
@@ -0,0 +1,49 @@
+/*
+   Copyright The containerd Authors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package server
+
+import (
+	"bytes"
+	"testing"
+
+	cioutil "github.com/containerd/containerd/pkg/ioutil"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCWWrite(t *testing.T) {
+	var buf bytes.Buffer
+	cw := &cappedWriter{w: cioutil.NewNopWriteCloser(&buf), remain: 10}
+
+	n, err := cw.Write([]byte("hello"))
+	assert.NoError(t, err)
+	assert.Equal(t, 5, n)
+
+	n, err = cw.Write([]byte("helloworld"))
+	assert.Equal(t, []byte("hellohello"), buf.Bytes(), "partial write")
+	assert.Equal(t, 5, n)
+	assert.ErrorIs(t, err, errNoRemain)
+
+	_, err = cw.Write([]byte("world"))
+	assert.ErrorIs(t, err, errNoRemain)
+}
+
+func TestCWClose(t *testing.T) {
+	var buf bytes.Buffer
+	cw := &cappedWriter{w: cioutil.NewNopWriteCloser(&buf), remain: 5}
+	err := cw.Close()
+	assert.NoError(t, err)
+}
