From ad8b87ba234e0a4993fe84983ff4e45a3e47a58f Mon Sep 17 00:00:00 2001
From: Hamza El-Saawy <hamzaelsaawy@microsoft.com>
Date: Mon, 9 May 2022 17:15:00 -0400
Subject: [PATCH] Add `Wait` to `binaryProcessor`

Add exported `Wait(ctx context.Context) error` interface that waits on
the underlying command (or context cancellation) and returns the error.

This fixes a race condition between `.wait()` and `.Err error`:
https://github.com/containerd/containerd/issues/6914

Signed-off-by: Hamza El-Saawy <hamzaelsaawy@microsoft.com>
---
 diff/stream_unix.go    | 16 ++++++++++++++++
 diff/stream_windows.go | 16 ++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/diff/stream_unix.go b/diff/stream_unix.go
index 4e458f448..6e0a44e01 100644
--- a/diff/stream_unix.go
+++ b/diff/stream_unix.go
@@ -89,6 +89,7 @@ func NewBinaryProcessor(ctx context.Context, imt, rmt string, stream StreamProce
 		r:      r,
 		mt:     rmt,
 		stderr: stderr,
+		done:   make(chan struct{}),
 	}
 	go p.wait()
 
@@ -111,6 +112,11 @@ type binaryProcessor struct {
 
 	mu  sync.Mutex
 	err error
+
+	// There is a race condition between waiting on c.cmd.Wait() and setting c.err within
+	// c.wait(), and reading that value from c.Err().
+	// Use done to wait for the returned error to be captured and set.
+	done chan struct{}
 }
 
 func (c *binaryProcessor) Err() error {
@@ -127,6 +133,16 @@ func (c *binaryProcessor) wait() {
 			c.mu.Unlock()
 		}
 	}
+	close(c.done)
+}
+
+func (c *binaryProcessor) Wait(ctx context.Context) error {
+	select {
+	case <-c.done:
+		return c.Err()
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 func (c *binaryProcessor) File() *os.File {
diff --git a/diff/stream_windows.go b/diff/stream_windows.go
index d03923f1f..a5ae19728 100644
--- a/diff/stream_windows.go
+++ b/diff/stream_windows.go
@@ -98,6 +98,7 @@ func NewBinaryProcessor(ctx context.Context, imt, rmt string, stream StreamProce
 		r:      r,
 		mt:     rmt,
 		stderr: stderr,
+		done:   make(chan struct{}),
 	}
 	go p.wait()
 
@@ -117,6 +118,11 @@ type binaryProcessor struct {
 
 	mu  sync.Mutex
 	err error
+
+	// There is a race condition between waiting on c.cmd.Wait() and setting c.err within
+	// c.wait(), and reading that value from c.Err().
+	// Use done to wait for the returned error to be captured and set.
+	done chan struct{}
 }
 
 func (c *binaryProcessor) Err() error {
@@ -133,6 +139,16 @@ func (c *binaryProcessor) wait() {
 			c.mu.Unlock()
 		}
 	}
+	close(c.done)
+}
+
+func (c *binaryProcessor) Wait(ctx context.Context) error {
+	select {
+	case <-c.done:
+		return c.Err()
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 func (c *binaryProcessor) File() *os.File {