Always apply schema level secrets

infer/configuration.go

@@ -107,7 +107,7 @@ func (c *config[T]) checkConfig(ctx context.Context, req p.CheckRequest) (p.Chec
 			return p.CheckResponse{}, err
 		}
 		return p.CheckResponse{
-			Inputs:   inputs,
+			Inputs:   applySecrets[T](inputs),
 			Failures: failures,
 		}, nil
 	}

infer/resource.go

@@ -865,41 +865,52 @@ func (*derivedResourceController[R, I, O]) getInstance() *R {
 }
 
 func (rc *derivedResourceController[R, I, O]) Check(ctx context.Context, req p.CheckRequest) (p.CheckResponse, error) {
-	encoder, i, failures, err := decodeCheckingMapErrors[I](req.News)
-	if err != nil {
-		return p.CheckResponse{}, err
-	}
-	if len(failures) > 0 {
-		return p.CheckResponse{
-			// If we failed to decode, we apply secrets pro-actively to ensure
-			// that they don't leak into previews.
-			Inputs:   applySecrets[I](req.News),
-			Failures: failures,
-		}, nil
-	}
-
 	var r R
 	if r, ok := ((interface{})(r)).(CustomCheck[I]); ok {
 		// The user implemented check manually, so call that.
 		//
-		// We do not apply defaults or secrets if the user has implemented Check
-		// themselves. Defaults and secrets are applied by [DefaultCheck].
-
-		defCheckEnc, i, failures, err := callCustomCheck(ctx, r, req.Urn.Name(), req.Olds, req.News)
+		// We do not apply defaults if the user has implemented Check
+		// themselves. Defaults are applied by [DefaultCheck].
+		encoder, i, failures, err := callCustomCheck(ctx, r, req.Urn.Name(), req.Olds, req.News)
 		if err != nil {
 			return p.CheckResponse{}, err
 		}
-		if defCheckEnc != nil {
-			encoder = *defCheckEnc
+
+		// callCustomCheck will have an encoder if and only if the custom check
+		// calls [DefaultCheck].
+		//
+		// If it doesn't have an encoder, but no error was returned, we do our
+		// best to recover secrets, unknowns, etc by calling
+		// decodeCheckingMapErrors to re-derive an encoder to use.
+		//
+		// There isn't any guaranteed relationship between the shape of `req.News`
+		// and `I`, so we are not guaranteed that `decodeCheckingMapErrors` won't
+		// produce errors.
+		if encoder == nil {
+			backupEncoder, _, _, _ := decodeCheckingMapErrors[I](req.News)
+			encoder = &backupEncoder
 		}
 
 		inputs, err := encoder.Encode(i)
 		return p.CheckResponse{
-			Inputs:   inputs,
+			Inputs:   applySecrets[I](inputs),
 			Failures: failures,
 		}, err
 	}
 
+	encoder, i, failures, err := decodeCheckingMapErrors[I](req.News)
+	if err != nil {
+		return p.CheckResponse{}, err
+	}
+	if len(failures) > 0 {
+		return p.CheckResponse{
+			// If we failed to decode, we apply secrets pro-actively to ensure
+			// that they don't leak into previews.
+			Inputs:   applySecrets[I](req.News),
+			Failures: failures,
+		}, nil
+	}
+
 	if i, err = defaultCheck(i); err != nil {
 		return p.CheckResponse{}, fmt.Errorf("unable to apply defaults: %w", err)
 	}
@@ -933,7 +944,6 @@ func callCustomCheck[T any](
 //
 // It also adds defaults to inputs as necessary, as defined by [Annotator.SetDefault].
 func DefaultCheck[I any](ctx context.Context, inputs resource.PropertyMap) (I, []p.CheckFailure, error) {
-	inputs = applySecrets[I](inputs)
 	enc, i, failures, err := decodeCheckingMapErrors[I](inputs)
 
 	if v, ok := ctx.Value(defaultCheckEncoderKey{}).(*defaultCheckEncoderValue); ok {

infer/tests/check_test.go

@@ -255,5 +255,25 @@ func TestCheckDefaultsRecursive(t *testing.T) {
 			}},
 		}},
 	}, resp.Inputs)
+}
+
+// TestCheckAlwaysAppliesSecrets checks that if a inferred provider resource has (1) a
+// field marked as secret with a field annotation (`provider:"secret"`), and (2)
+// implements [infer.CustomCheck] without calling [infer.DefaultDiff], the field is still
+// marked as secret.
+func TestCheckAlwaysAppliesSecrets(t *testing.T) {
+	t.Parallel()
+
+	prov := provider()
+	resp, err := prov.Check(p.CheckRequest{
+		Urn: urn("CustomCheckNoDefaults", "check-env"),
+		News: resource.PropertyMap{
+			"input": resource.NewProperty("value"),
+		},
+	})
+	require.NoError(t, err)
 
+	assert.Equal(t, resource.PropertyMap{
+		"input": resource.MakeSecret(resource.NewProperty("value")),
+	}, resp.Inputs)
 }

infer/tests/provider.go

@@ -386,6 +386,31 @@ func (w *ReadConfigCustom) Create(
 	return "read", ReadConfigCustomOutput{Config: string(bytes)}, err
 }
 
+var (
+	_ infer.CustomResource[CustomCheckNoDefaultsArgs, CustomCheckNoDefaultsOutput] = &CustomCheckNoDefaults{}
+	_ infer.CustomCheck[CustomCheckNoDefaultsArgs]                                 = &CustomCheckNoDefaults{}
+)
+
+type (
+	CustomCheckNoDefaults     struct{}
+	CustomCheckNoDefaultsArgs struct {
+		Input string `pulumi:"input" provider:"secret"`
+	}
+	CustomCheckNoDefaultsOutput struct{ CustomCheckNoDefaultsArgs }
+)
+
+func (w *CustomCheckNoDefaults) Check(_ context.Context,
+	_ string, _ resource.PropertyMap, m resource.PropertyMap,
+) (CustomCheckNoDefaultsArgs, []p.CheckFailure, error) {
+	return CustomCheckNoDefaultsArgs{Input: m["input"].StringValue()}, nil, nil
+}
+
+func (w *CustomCheckNoDefaults) Create(
+	ctx context.Context, name string, inputs CustomCheckNoDefaultsArgs, preview bool,
+) (string, CustomCheckNoDefaultsOutput, error) {
+	return "id", CustomCheckNoDefaultsOutput{inputs}, nil
+}
+
 func providerOpts(config infer.InferredConfig) infer.Options {
 	return infer.Options{
 		Config: config,
@@ -399,6 +424,7 @@ func providerOpts(config infer.InferredConfig) infer.Options {
 			infer.Resource[*Recursive, RecursiveArgs, RecursiveOutput](),
 			infer.Resource[*ReadConfig, ReadConfigArgs, ReadConfigOutput](),
 			infer.Resource[*ReadConfigCustom, ReadConfigCustomArgs, ReadConfigCustomOutput](),
+			infer.Resource[*CustomCheckNoDefaults, CustomCheckNoDefaultsArgs, CustomCheckNoDefaultsOutput](),
 		},
 		Functions: []infer.InferredFunction{
 			infer.Function[*GetJoin, JoinArgs, JoinResult](),

middleware/schema/schema.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"reflect"
 	"strings"
+	"sync"
 
 	"github.com/hashicorp/go-multierror"
 	"github.com/pulumi/pulumi/pkg/v3/codegen/schema"
@@ -92,6 +93,8 @@ type state struct {
 	lowerSchema    *cache
 	combinedSchema *cache
 	innerGetSchema func(ctx context.Context, req p.GetSchemaRequest) (p.GetSchemaResponse, error)
+
+	m sync.Mutex
 }
 
 // Options sets the schema options used by [Wrap].
@@ -179,6 +182,9 @@ func Wrap(provider p.Provider, opts Options) p.Provider {
 }
 
 func (s *state) GetSchema(ctx context.Context, req p.GetSchemaRequest) (p.GetSchemaResponse, error) {
+	s.m.Lock()
+	defer s.m.Unlock()
+
 	if s.schema.isEmpty() {
 		spec, err := s.generateSchema(ctx)
 		if err != nil {

tests/config_test.go

@@ -191,38 +191,27 @@ func TestInferCustomCheckConfig(t *testing.T) {
 		Config: infer.Config[*config](),
 	}))
 
-	t.Run("with-default-check", func(t *testing.T) {
-		resp, err := s.CheckConfig(p.CheckRequest{
-			Urn: resource.CreateURN("p", "pulumi:providers:test", "", "test", "dev"),
-			News: resource.PropertyMap{
-				"field":         resource.NewProperty("value"),
+	// Test that our manual implementation of check works the same as the default
+	// version, and that secrets are applied regardless of if check is used.
+	for _, applyDefaults := range []bool{true, false} {
+		applyDefaults := applyDefaults
+		t.Run(fmt.Sprintf("%t", applyDefaults), func(t *testing.T) {
+			t.Parallel()
+			resp, err := s.CheckConfig(p.CheckRequest{
+				Urn: resource.CreateURN("p", "pulumi:providers:test", "", "test", "dev"),
+				News: resource.PropertyMap{
+					"field":         resource.NewProperty("value"),
+					"not":           resource.NewProperty("not-secret"),
+					"applyDefaults": resource.NewProperty(applyDefaults),
+				},
+			})
+			require.NoError(t, err)
+			require.Empty(t, resp.Failures)
+			assert.Equal(t, resource.PropertyMap{
+				"field":         resource.MakeSecret(resource.NewProperty("value")),
 				"not":           resource.NewProperty("not-secret"),
-				"applyDefaults": resource.NewProperty(true),
-			},
-		})
-		require.NoError(t, err)
-		require.Empty(t, resp.Failures)
-		assert.Equal(t, resource.PropertyMap{
-			"field":         resource.MakeSecret(resource.NewProperty("value")),
-			"not":           resource.NewProperty("not-secret"),
-			"applyDefaults": resource.NewProperty(true),
-		}, resp.Inputs)
-	})
-
-	t.Run("without-default-check", func(t *testing.T) {
-		resp, err := s.CheckConfig(p.CheckRequest{
-			News: resource.PropertyMap{
-				"field":         resource.NewProperty("value"),
-				"not":           resource.NewProperty("not-secret"),
-				"applyDefaults": resource.NewProperty(false),
-			},
+				"applyDefaults": resource.NewProperty(applyDefaults),
+			}, resp.Inputs)
 		})
-		require.NoError(t, err)
-		require.Empty(t, resp.Failures)
-		assert.Equal(t, resource.PropertyMap{
-			"field":         resource.NewProperty("value"),
-			"not":           resource.NewProperty("not-secret"),
-			"applyDefaults": resource.NewProperty(false),
-		}, resp.Inputs)
-	})
+	}
 }