Fix data races and a deadlock in the byte stream helpers

Race: clientStream.Send and Close both reach ttrpc's unsynchronized
localClosed flag. readByteStream.window is written by
the window goroutine and by Read, so make it atomic.

Deadlock: Read the window by bytes copied instead of bytes
received, so a short read leaves the window above windowSize.

Signed-off-by: Jing Chen <milantracy@gmail.com>

core/streaming/proxy/streaming.go

@@ -21,6 +21,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"sync"
 
 	streamingapi "github.com/containerd/containerd/api/services/streaming/v1"
 	"github.com/containerd/errdefs"
@@ -105,10 +106,13 @@ func (sc *streamCreator) Create(ctx context.Context, id string) (streaming.Strea
 }
 
 type clientStream struct {
-	s streamingapi.TTRPCStreaming_StreamClient
+	sendMu sync.Mutex
+	s      streamingapi.TTRPCStreaming_StreamClient
 }
 
 func (cs *clientStream) Send(a typeurl.Any) (err error) {
+	cs.sendMu.Lock()
+	defer cs.sendMu.Unlock()
 	err = cs.s.Send(typeurl.MarshalProto(a))
 	if !errors.Is(err, io.EOF) {
 		err = errgrpc.ToNative(err)
@@ -125,5 +129,7 @@ func (cs *clientStream) Recv() (a typeurl.Any, err error) {
 }
 
 func (cs *clientStream) Close() error {
+	cs.sendMu.Lock()
+	defer cs.sendMu.Unlock()
 	return cs.s.CloseSend()
 }

core/streaming/proxy/streaming_test.go

@@ -0,0 +1,61 @@
+/*
+   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 proxy
+
+import (
+	"io"
+	"testing"
+
+	"github.com/containerd/typeurl/v2"
+	"google.golang.org/protobuf/types/known/anypb"
+)
+
+// unsafeStreamClient writes a flag in CloseSend that Send reads, as ttrpc does.
+type unsafeStreamClient struct {
+	closed bool
+}
+
+func (c *unsafeStreamClient) Send(*anypb.Any) error {
+	if c.closed {
+		return io.EOF
+	}
+	return nil
+}
+
+func (c *unsafeStreamClient) Recv() (*anypb.Any, error) { return &anypb.Any{}, nil }
+func (c *unsafeStreamClient) CloseSend() error          { c.closed = true; return nil }
+func (c *unsafeStreamClient) SendMsg(m any) error       { return nil }
+func (c *unsafeStreamClient) RecvMsg(m any) error       { return nil }
+
+// TestClientStreamSendClose races under -race unless Send and Close are serialized.
+func TestClientStreamSendClose(t *testing.T) {
+	cs := &clientStream{s: &unsafeStreamClient{}}
+	var a typeurl.Any = &anypb.Any{}
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for range 100 {
+			cs.Send(a)
+		}
+	}()
+
+	if err := cs.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	<-done
+}

core/transfer/streaming/reader.go

@@ -18,9 +18,9 @@ package streaming
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"io"
+	"sync/atomic"
 
 	transferapi "github.com/containerd/containerd/api/types/transfer"
 	"github.com/containerd/containerd/v2/core/streaming"
@@ -30,7 +30,7 @@ import (
 type readByteStream struct {
 	ctx       context.Context
 	stream    streaming.Stream
-	window    int32
+	window    atomic.Int32
 	updated   chan struct{}
 	errCh     chan error
 	remaining []byte
@@ -40,13 +40,12 @@ func ReadByteStream(ctx context.Context, stream streaming.Stream) io.ReadCloser
 	rbs := &readByteStream{
 		ctx:     ctx,
 		stream:  stream,
-		window:  0,
-		errCh:   make(chan error),
+		errCh:   make(chan error, 1),
 		updated: make(chan struct{}, 1),
 	}
 	go func() {
 		for {
-			if rbs.window >= windowSize {
+			if rbs.window.Load() >= windowSize {
 				select {
 				case <-ctx.Done():
 					return
@@ -62,11 +61,11 @@ func ReadByteStream(ctx context.Context, stream streaming.Stream) io.ReadCloser
 				rbs.errCh <- err
 				return
 			}
-			if err := stream.Send(anyType); err == nil {
-				rbs.window += windowSize
-			} else if !errors.Is(err, io.EOF) {
-				rbs.errCh <- err
+			if err := stream.Send(anyType); err != nil {
+				// Read reports stream errors.
+				return
 			}
+			rbs.window.Add(windowSize)
 		}
 
 	}()
@@ -105,9 +104,12 @@ func (r *readByteStream) Read(p []byte) (n int, err error) {
 		if len(v.Data) > plen {
 			r.remaining = v.Data[plen:]
 		}
-		r.window = r.window - int32(n)
-		if r.window < windowSize {
-			r.updated <- struct{}{}
+		if r.window.Add(-int32(len(v.Data))) < windowSize {
+			select {
+			case r.updated <- struct{}{}:
+			default:
+				// Don't block if the window goroutine has ended.
+			}
 		}
 		return n, nil
 	default:

core/transfer/streaming/reader_test.go

@@ -0,0 +1,45 @@
+/*
+   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 streaming
+
+import (
+	"bytes"
+	"io"
+	"testing"
+)
+
+// TestReadByteStream races under -race unless the window counter is atomic.
+func TestReadByteStream(t *testing.T) {
+	expected := bytes.Repeat([]byte("containerd"), windowSize/4)
+
+	rs, ws := pipeStream()
+	w := WriteByteStream(t.Context(), ws)
+	r := ReadByteStream(t.Context(), rs)
+
+	go func() {
+		io.Copy(w, bytes.NewReader(expected))
+		w.Close()
+	}()
+
+	actual, err := io.ReadAll(r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(expected, actual) {
+		t.Fatalf("read %d bytes, want %d", len(actual), len(expected))
+	}
+}