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

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)
 }
 

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)
 }