From fccd0b9b70255099691deca5dc1d577efcfc889b Mon Sep 17 00:00:00 2001
From: Ian Lance Taylor <iant@golang.org>
Date: Mon, 11 Sep 2023 16:43:20 -0700
Subject: [PATCH] context: support non-standard Context in Cause

If Cause is called on a non-standard Context, call ctx.Err.

Fixes #62582

Change-Id: Iac4ed93203eb5529f8839eb479b6ee2ee5ff6cbc
Reviewed-on: https://go-review.googlesource.com/c/go/+/527277
Reviewed-by: Bryan Mills <bcmills@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Auto-Submit: Ian Lance Taylor <iant@google.com>
Reviewed-by: Ian Lance Taylor <iant@google.com>
---
 src/context/context.go |   7 ++-
 src/context/x_test.go  | 120 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 126 insertions(+), 1 deletion(-)

diff --git a/src/context/context.go b/src/context/context.go
index ee66b43c85..80e1787576 100644
--- a/src/context/context.go
+++ b/src/context/context.go
@@ -286,7 +286,12 @@ func Cause(c Context) error {
 		defer cc.mu.Unlock()
 		return cc.cause
 	}
-	return nil
+	// There is no cancelCtxKey value, so we know that c is
+	// not a descendant of some Context created by WithCancelCause.
+	// Therefore, there is no specific cause to return.
+	// If this is not one of the standard Context types,
+	// it might still have an error even though it won't have a cause.
+	return c.Err()
 }
 
 // AfterFunc arranges to call f in its own goroutine after ctx is done
diff --git a/src/context/x_test.go b/src/context/x_test.go
index 57fe60b4ee..e006e53470 100644
--- a/src/context/x_test.go
+++ b/src/context/x_test.go
@@ -867,6 +867,126 @@ func TestCustomContextPropagation(t *testing.T) {
 	}
 }
 
+// customCauseContext is a custom Context used to test context.Cause.
+type customCauseContext struct {
+	mu   sync.Mutex
+	done chan struct{}
+	err  error
+
+	cancelChild CancelFunc
+}
+
+func (ccc *customCauseContext) Deadline() (deadline time.Time, ok bool) {
+	return
+}
+
+func (ccc *customCauseContext) Done() <-chan struct{} {
+	ccc.mu.Lock()
+	defer ccc.mu.Unlock()
+	return ccc.done
+}
+
+func (ccc *customCauseContext) Err() error {
+	ccc.mu.Lock()
+	defer ccc.mu.Unlock()
+	return ccc.err
+}
+
+func (ccc *customCauseContext) Value(key any) any {
+	return nil
+}
+
+func (ccc *customCauseContext) cancel() {
+	ccc.mu.Lock()
+	ccc.err = Canceled
+	close(ccc.done)
+	cancelChild := ccc.cancelChild
+	ccc.mu.Unlock()
+
+	if cancelChild != nil {
+		cancelChild()
+	}
+}
+
+func (ccc *customCauseContext) setCancelChild(cancelChild CancelFunc) {
+	ccc.cancelChild = cancelChild
+}
+
+func TestCustomContextCause(t *testing.T) {
+	// Test if we cancel a custom context, Err and Cause return Canceled.
+	ccc := &customCauseContext{
+		done: make(chan struct{}),
+	}
+	ccc.cancel()
+	if got := ccc.Err(); got != Canceled {
+		t.Errorf("ccc.Err() = %v, want %v", got, Canceled)
+	}
+	if got := Cause(ccc); got != Canceled {
+		t.Errorf("Cause(ccc) = %v, want %v", got, Canceled)
+	}
+
+	// Test that if we pass a custom context to WithCancelCause,
+	// and then cancel that child context with a cause,
+	// that the cause of the child canceled context is correct
+	// but that the parent custom context is not canceled.
+	ccc = &customCauseContext{
+		done: make(chan struct{}),
+	}
+	ctx, causeFunc := WithCancelCause(ccc)
+	cause := errors.New("TestCustomContextCause")
+	causeFunc(cause)
+	if got := ctx.Err(); got != Canceled {
+		t.Errorf("after CancelCauseFunc ctx.Err() = %v, want %v", got, Canceled)
+	}
+	if got := Cause(ctx); got != cause {
+		t.Errorf("after CancelCauseFunc Cause(ctx) = %v, want %v", got, cause)
+	}
+	if got := ccc.Err(); got != nil {
+		t.Errorf("after CancelCauseFunc ccc.Err() = %v, want %v", got, nil)
+	}
+	if got := Cause(ccc); got != nil {
+		t.Errorf("after CancelCauseFunc Cause(ccc) = %v, want %v", got, nil)
+	}
+
+	// Test that if we now cancel the parent custom context,
+	// the cause of the child canceled context is still correct,
+	// and the parent custom context is canceled without a cause.
+	ccc.cancel()
+	if got := ctx.Err(); got != Canceled {
+		t.Errorf("after CancelCauseFunc ctx.Err() = %v, want %v", got, Canceled)
+	}
+	if got := Cause(ctx); got != cause {
+		t.Errorf("after CancelCauseFunc Cause(ctx) = %v, want %v", got, cause)
+	}
+	if got := ccc.Err(); got != Canceled {
+		t.Errorf("after CancelCauseFunc ccc.Err() = %v, want %v", got, Canceled)
+	}
+	if got := Cause(ccc); got != Canceled {
+		t.Errorf("after CancelCauseFunc Cause(ccc) = %v, want %v", got, Canceled)
+	}
+
+	// Test that if we associate a custom context with a child,
+	// then canceling the custom context cancels the child.
+	ccc = &customCauseContext{
+		done: make(chan struct{}),
+	}
+	ctx, cancelFunc := WithCancel(ccc)
+	ccc.setCancelChild(cancelFunc)
+	ccc.cancel()
+	if got := ctx.Err(); got != Canceled {
+		t.Errorf("after CancelCauseFunc ctx.Err() = %v, want %v", got, Canceled)
+	}
+	if got := Cause(ctx); got != Canceled {
+		t.Errorf("after CancelCauseFunc Cause(ctx) = %v, want %v", got, Canceled)
+	}
+	if got := ccc.Err(); got != Canceled {
+		t.Errorf("after CancelCauseFunc ccc.Err() = %v, want %v", got, Canceled)
+	}
+	if got := Cause(ccc); got != Canceled {
+		t.Errorf("after CancelCauseFunc Cause(ccc) = %v, want %v", got, Canceled)
+	}
+}
+
 func TestAfterFuncCalledAfterCancel(t *testing.T) {
 	ctx, cancel := WithCancel(Background())
 	donec := make(chan struct{})
-- 
2.50.0