Add binary stream functionality and helpers

Signed-off-by: Derek McGowan <derek@mcg.dev>

pkg/transfer/streaming/stream.go

@@ -0,0 +1,210 @@
+/*
+   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 (
+	"context"
+	"crypto/rand"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"io"
+	"sync"
+	"time"
+
+	transferapi "github.com/containerd/containerd/api/types/transfer"
+	"github.com/containerd/containerd/log"
+	"github.com/containerd/containerd/pkg/streaming"
+	"github.com/containerd/typeurl"
+)
+
+const maxRead = 32 * 1024
+const windowSize = 2 * maxRead
+
+var bufPool = &sync.Pool{
+	New: func() interface{} {
+		buffer := make([]byte, maxRead)
+		return &buffer
+	},
+}
+
+func SendStream(ctx context.Context, r io.Reader, stream streaming.Stream) {
+	window := make(chan int32)
+	go func() {
+		defer close(window)
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
+
+			any, err := stream.Recv()
+			if err != nil {
+				if !errors.Is(err, io.EOF) && !errors.Is(err, context.Canceled) {
+					log.G(ctx).WithError(err).Error("send stream ended without EOF")
+				}
+				return
+			}
+			i, err := typeurl.UnmarshalAny(any)
+			if err != nil {
+				log.G(ctx).WithError(err).Error("failed to unmarshal stream object")
+				continue
+			}
+			switch v := i.(type) {
+			case *transferapi.WindowUpdate:
+				select {
+				case <-ctx.Done():
+					return
+				case window <- v.Update:
+				}
+			default:
+				log.G(ctx).Errorf("unexpected stream object of type %T", i)
+			}
+		}
+	}()
+	go func() {
+		defer stream.Close()
+
+		buf := bufPool.Get().(*[]byte)
+		defer bufPool.Put(buf)
+
+		var remaining int32
+
+		for {
+			if remaining > 0 {
+				// Don't wait for window update since there are remaining
+				select {
+				case <-ctx.Done():
+					// TODO: Send error message on stream before close to allow remote side to return error
+					return
+				case update := <-window:
+					remaining += update
+				default:
+				}
+			} else {
+				// Block until window updated
+				select {
+				case <-ctx.Done():
+					// TODO: Send error message on stream before close to allow remote side to return error
+					return
+				case update := <-window:
+					remaining = update
+				}
+			}
+			var max int32 = maxRead
+			if max > remaining {
+				max = remaining
+			}
+			b := (*buf)[:max]
+			n, err := r.Read(b)
+			if err != nil {
+				if err != io.EOF {
+					log.G(ctx).WithError(err).Errorf("failed to read stream source")
+					// TODO: Send error message on stream before close to allow remote side to return error
+				}
+				return
+			}
+			remaining = remaining - int32(n)
+
+			data := &transferapi.Data{
+				Data: b[:n],
+			}
+			any, err := typeurl.MarshalAny(data)
+			if err != nil {
+				log.G(ctx).WithError(err).Errorf("failed to marshal data for send")
+				// TODO: Send error message on stream before close to allow remote side to return error
+				return
+			}
+			if err := stream.Send(any); err != nil {
+				log.G(ctx).WithError(err).Errorf("send failed")
+				return
+			}
+		}
+	}()
+}
+
+func ReceiveStream(ctx context.Context, stream streaming.Stream) io.Reader {
+	r, w := io.Pipe()
+	go func() {
+		defer stream.Close()
+		var window int32
+		for {
+			var werr error
+			if window < windowSize {
+				update := &transferapi.WindowUpdate{
+					Update: windowSize,
+				}
+				any, err := typeurl.MarshalAny(update)
+				if err != nil {
+					w.CloseWithError(fmt.Errorf("failed to marshal window update: %w", err))
+					return
+				}
+				// check window update error after recv, stream may be complete
+				if werr = stream.Send(any); werr == nil {
+					window += windowSize
+				} else if werr == io.EOF {
+					// TODO: Why does send return EOF here
+					werr = nil
+				}
+			}
+			any, err := stream.Recv()
+			if err != nil {
+				if err == io.EOF {
+					err = nil
+				} else {
+					err = fmt.Errorf("received failed: %w", err)
+				}
+				w.CloseWithError(err)
+				return
+			} else if werr != nil {
+				// Try receive before erroring out
+				w.CloseWithError(fmt.Errorf("failed to send window update: %w", werr))
+				return
+			}
+			i, err := typeurl.UnmarshalAny(any)
+			if err != nil {
+				w.CloseWithError(fmt.Errorf("failed to unmarshal received object: %w", err))
+				return
+			}
+			switch v := i.(type) {
+			case *transferapi.Data:
+				n, err := w.Write(v.Data)
+				if err != nil {
+					w.CloseWithError(fmt.Errorf("failed to unmarshal received object: %w", err))
+					// Close will error out sender
+					return
+				}
+				window = window - int32(n)
+			// TODO: Handle error case
+			default:
+				log.G(ctx).Warnf("Ignoring unknown stream object of type %T", i)
+				continue
+			}
+		}
+
+	}()
+
+	return r
+}
+
+func GenerateID(prefix string) string {
+	t := time.Now()
+	var b [3]byte
+	rand.Read(b[:])
+	return fmt.Sprintf("%s-%d-%s", prefix, t.Nanosecond(), base64.URLEncoding.EncodeToString(b[:]))
+}