fix(types): make UnwrapSDKContext work after context.WithValue (#11341)

cosmos · Mar 11, 2022 · 25feb23 · 25feb23
1 parent f4d50a9
commit 25feb23
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/types/context.go b/types/context.go
@@ -237,6 +237,10 @@ func (c Context) WithValue(key, value interface{}) Context {
 }
 
 func (c Context) Value(key interface{}) interface{} {
+	if key == SdkContextKey {
+		return c
+	}
+
 	return c.baseCtx.Value(key)
 }
 

diff --git a/types/context_test.go b/types/context_test.go
@@ -220,4 +220,9 @@ func (s *contextTestSuite) TestUnwrapSDKContext() {
 
 	ctx = context.Background()
 	s.Require().Panics(func() { types.UnwrapSDKContext(ctx) })
+
+	// test unwrapping when we've used context.WithValue
+	ctx = context.WithValue(sdkCtx, "foo", "bar")
+	sdkCtx2 = types.UnwrapSDKContext(ctx)
+	s.Require().Equal(sdkCtx, sdkCtx2)
 }