Implicitly discard the input to drain the reader

Kazuyoshi Kato · dmcgowan · commit b40a356cfd9f · 2022-06-06T10:03:15.000-07:00
Signed-off-by: Derek McGowan &lt;derek@mcg.dev&gt;
diff --git a/pkg/cri/server/container_execsync.go b/pkg/cri/server/container_execsync.go
@@ -43,11 +43,9 @@ type cappedWriter struct {
 	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
+		return len(p), nil
 	}
 
 	end := cw.remain
@@ -60,26 +58,35 @@ func (cw *cappedWriter) Write(p []byte) (int, error) {
 	if err != nil {
 		return written, err
 	}
-	if written < len(p) {
-		return written, errNoRemain
-	}
-	return written, nil
+	return len(p), nil
 }
 
 func (cw *cappedWriter) Close() error {
 	return cw.w.Close()
 }
 
+func (cw *cappedWriter) isFull() bool {
+	return cw.remain <= 0
+}
+
 // 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
+
+	// cappedWriter truncates the output. In that case, the size of
+	// the ExecSyncResponse will hit the CRI plugin's gRPC response limit.
+	// Thus the callers outside of the containerd process (e.g. Kubelet) never see
+	// the truncated output.
+	cout := &cappedWriter{w: cioutil.NewNopWriteCloser(&stdout), remain: maxStreamSize}
+	cerr := &cappedWriter{w: cioutil.NewNopWriteCloser(&stderr), remain: maxStreamSize}
+
 	exitCode, err := c.execInContainer(ctx, r.GetContainerId(), execOptions{
 		cmd:     r.GetCmd(),
-		stdout:  &cappedWriter{w: cioutil.NewNopWriteCloser(&stdout), remain: maxStreamSize},
-		stderr:  &cappedWriter{w: cioutil.NewNopWriteCloser(&stderr), remain: maxStreamSize},
+		stdout:  cout,
+		stderr:  cerr,
 		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
@@ -33,12 +33,15 @@ func TestCWWrite(t *testing.T) {
 	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)
+	assert.NoError(t, err, "no errors even it hits the cap")
+	assert.Equal(t, 10, n, "no indication of partial write")
+	assert.True(t, cw.isFull())
+	assert.Equal(t, []byte("hellohello"), buf.Bytes(), "the underlying writer is capped")
 
 	_, err = cw.Write([]byte("world"))
-	assert.ErrorIs(t, err, errNoRemain)
+	assert.NoError(t, err)
+	assert.True(t, cw.isFull())
+	assert.Equal(t, []byte("hellohello"), buf.Bytes(), "the underlying writer is capped")
 }
 
 func TestCWClose(t *testing.T) {
