tofu: Include schema-based marks in final state after apply

During the "refresh" and "plan" steps we build the marks for a managed
resource object's value as a combination of the marks from the input
(prior state or configuration, respectively) and the marks implied by the
provider schema.

However, the apply step was previously relying only on the marks from the
planned new state, without considering marks from the provider schema. That
meant that a sensitive attribute contained within a container that is
unknown during planning could not be marked as sensitive once the container
became known, because the corresponding value did not exist at all in the
planned new state and therefore could not carry a sensitive mark.

To fix this problem, this changes the apply step to match the strategy
already used in the refresh and plan steps: using combinePathValueMarks
to blend the dynamic marks with the static information from the schema,
so that the final value saved in the new state snapshot will have a full
set of sensitive markings for the next plan/apply round to rely on. Without
this the next plan/apply round would produce a spurious diff due to the
sensitivity of the nested attributes appearing to have changed.

This introduces a new test TestContext2Apply_sensitiveInsideUnknown which
covers the case where the sensitivity information comes from schema. The
preexisting test TestContext2Apply_additionalSensitiveFromState already
covered the case of dynamically-tracked sensitivity information, and
remains passing without modification after this change.

Signed-off-by: Martin Atkins <mart@degeneration.co.uk>

internal/tofu/context_apply2_test.go

@@ -599,6 +599,118 @@ output "out" {
 	}
 }
 
+// TestContext2Apply_sensitiveInsideUnknown verifies that OpenTofu uses
+// sensitive value information from provider schema when deciding what to
+// mark as sensitive in the final state after apply, even if the values in
+// question were not yet available during the planning phase.
+//
+// For additional context, refer to https://github.com/opentofu/opentofu/issues/3367 .
+func TestContext2Apply_sensitiveInsideUnknown(t *testing.T) {
+	m := testModuleInline(t, map[string]string{
+		"main.tf": `
+			terraform {
+				required_providers {
+					test = {
+						source = "example.com/foo/test"
+					}
+				}
+			}
+
+			resource "test_sensitive" "test" {
+			}
+		`,
+	})
+	p := &MockProvider{
+		GetProviderSchemaResponse: &providers.GetProviderSchemaResponse{
+			ResourceTypes: map[string]providers.Schema{
+				"test_sensitive": {
+					Block: &configschema.Block{
+						Attributes: map[string]*configschema.Attribute{
+							// We need at least one level of container
+							// indirection here, because we're verifying
+							// that the nested attribute has its sensitivity
+							// recorded in the final state even though
+							// the entire container will be unknown during
+							// planning and therefore the nested value cannot
+							// be marked in the "planned new state".
+							"container": {
+								Computed: true,
+								NestedType: &configschema.Object{
+									Nesting: configschema.NestingSingle,
+									Attributes: map[string]*configschema.Attribute{
+										"sensitive": {
+											Type:      cty.String,
+											Sensitive: true,
+											Computed:  true,
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		PlanResourceChangeResponse: &providers.PlanResourceChangeResponse{
+			PlannedState: cty.ObjectVal(map[string]cty.Value{
+				// The whole container is initially unknown, so the sensitivity
+				// of the nested "sensitive" attribute cannot be tracked as
+				// part of this value, forcing the apply phase to rely on
+				// the provider schema for that information.
+				"container": cty.UnknownVal(cty.Object(map[string]cty.Type{"sensitive": cty.String})),
+			}),
+		},
+		ApplyResourceChangeResponse: &providers.ApplyResourceChangeResponse{
+			NewState: cty.ObjectVal(map[string]cty.Value{
+				"container": cty.ObjectVal(map[string]cty.Value{
+					// This nested string value should be automatically marked
+					// as sensitive in the final state due to the provider
+					// schema, even though it wasn't present in the "planned
+					// state" in PlanResourceChangeResponse above.
+					"sensitive": cty.StringVal("hello"),
+				}),
+			}),
+		},
+	}
+
+	tofuCtx := testContext2(t, &ContextOpts{
+		Providers: map[addrs.Provider]providers.Factory{
+			addrs.MustParseProviderSourceString("example.com/foo/test"): testProviderFuncFixed(p),
+		},
+	})
+
+	plan, diags := tofuCtx.Plan(t.Context(), m, states.NewState(), SimplePlanOpts(plans.NormalMode, nil))
+	assertNoErrors(t, diags)
+	state, diags := tofuCtx.Apply(context.Background(), plan, m, &ApplyOpts{})
+	assertNoErrors(t, diags)
+
+	riState := state.ResourceInstance(addrs.Resource{
+		Mode: addrs.ManagedResourceMode,
+		Type: "test_sensitive",
+		Name: "test",
+	}.Absolute(addrs.RootModuleInstance).Instance(addrs.NoKey))
+	if riState == nil {
+		t.Fatal("resource instance state is missing")
+	}
+	if riState.Current == nil {
+		t.Fatal("resource instance has no current object")
+	}
+	foundSensitive := false
+	for _, pvm := range riState.Current.AttrSensitivePaths {
+		t.Logf("marks at %s: %#v", tfdiags.FormatCtyPath(pvm.Path), pvm.Marks)
+		if _, ok := pvm.Marks[marks.Sensitive]; !ok {
+			continue
+		}
+		if !pvm.Path.Equals(cty.GetAttrPath("container").GetAttr("sensitive")) {
+			continue
+		}
+		foundSensitive = true
+	}
+	if !foundSensitive {
+		t.Errorf("no sensitive mark for .container.sensitive in %s", spew.Sdump(riState.Current))
+	}
+}
+
 func TestContext2Apply_ignoreImpureFunctionChanges(t *testing.T) {
 	// The impure function call should not cause a planned change with
 	// ignore_changes

internal/tofu/node_resource_abstract_instance.go

@@ -2779,8 +2779,9 @@ func (n *NodeAbstractResourceInstance) apply(
 	newVal := resp.NewState
 
 	// If we have paths to mark, mark those on this new value
-	if len(afterPaths) > 0 {
-		newVal = newVal.MarkWithPaths(afterPaths)
+	newValMarks := combinePathValueMarks(afterPaths, schema.ValueMarks(newVal, nil))
+	if len(newValMarks) > 0 {
+		newVal = newVal.MarkWithPaths(newValMarks)
 	}
 
 	if newVal == cty.NilVal {