Fixes #1231 panic when accessing GetRawState

Diff customizers can request access to RawState, which prior to this PR was always set to null, causing a panic.

pkg/tfbridge/diff_test.go

@@ -2,6 +2,7 @@ package tfbridge
 
 import (
 	"context"
+	"fmt"
 	"testing"
 
 	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
@@ -143,6 +144,47 @@ func TestCustomizeDiff(t *testing.T) {
 		assert.Equal(t, changes, pulumirpc.DiffResponse_DIFF_NONE)
 		assert.Equal(t, expectedDiff, diff)
 	})
+
+	t.Run("CustomDiffDoesNotPanicOnGetRawState", func(t *testing.T) {
+		for _, diffStrat := range []shimv2.DiffStrategy{shimv2.PlanState, shimv2.ClassicDiff} {
+			diffStrat := diffStrat
+			t.Run(fmt.Sprintf("%v", diffStrat), func(t *testing.T) {
+				customDiffRes := &v2Schema.Resource{
+					Schema: tfs,
+					CustomizeDiff: func(_ context.Context, diff *v2Schema.ResourceDiff, _ interface{}) error {
+						rawStateType := diff.GetRawState().Type()
+						if !rawStateType.HasAttribute("outp") {
+							return fmt.Errorf("Expected rawState type to have attribute: outp")
+						}
+						return nil
+					},
+				}
+
+				v2Provider := &v2Schema.Provider{
+					ResourcesMap: map[string]*v2Schema.Resource{
+						"resource": customDiffRes,
+					},
+				}
+
+				provider := shimv2.NewProvider(v2Provider, shimv2.WithDiffStrategy(diffStrat))
+
+				// Convert the inputs and state to TF config and resource attributes.
+				r := Resource{
+					TF:     shimv2.NewResource(customDiffRes),
+					Schema: &ResourceInfo{Fields: info},
+				}
+				tfState, err := MakeTerraformState(r, "id", stateMap)
+				assert.NoError(t, err)
+
+				config, _, err := MakeTerraformConfig(&Provider{tf: provider}, inputsMap, sch, info)
+				assert.NoError(t, err)
+
+				// Calling Diff with the given CustomizeDiff used to panic, no more asserts needed.
+				_, err = provider.Diff("resource", tfState, config)
+				assert.NoError(t, err)
+			})
+		}
+	})
 }
 
 func diffTest(t *testing.T, tfs map[string]*schema.Schema, info map[string]*SchemaInfo,

pkg/tfshim/sdk-v2/provider_diff.go

@@ -79,14 +79,22 @@ func (p v2Provider) simpleDiff(
 
 	switch diffStrat {
 	case ClassicDiff:
-		if s.RawPlan.IsNull() {
+		state := s.DeepCopy()
+		if state.RawPlan.IsNull() {
 			// SimpleDiff may read RawPlan and panic if it is nil; while in the case of ClassicDiff we do
 			// not yet do what TF CLI does (that is PlanState), it is better to approximate and assume that
 			// the RawPlan is the same as RawConfig than to have the code panic.
-			rawPlan := rawConfigVal
-			return res.SimpleDiff(ctx, instanceStateWithUpdatedRawPlan(s, rawPlan), c, meta)
+			state.RawConfig = rawConfigVal
 		}
-		return res.SimpleDiff(ctx, s, c, meta)
+		if state.RawState.IsNull() {
+			// Same trick as for nil RawPlan.
+			priorStateVal, err := state.AttrsAsObjectValue(res.CoreConfigSchema().ImpliedType())
+			if err != nil {
+				return nil, err
+			}
+			state.RawState = priorStateVal
+		}
+		return res.SimpleDiff(ctx, state, c, meta)
 	case PlanState:
 		return simpleDiffViaPlanState(ctx, res, s, rawConfigVal, meta)
 	case TryPlanState:
@@ -128,15 +136,6 @@ func (p v2Provider) simpleDiff(
 	}
 }
 
-func instanceStateWithUpdatedRawPlan(
-	s *terraform.InstanceState,
-	rawPlan hcty.Value,
-) *terraform.InstanceState {
-	c := s.DeepCopy()
-	c.RawPlan = rawPlan
-	return c
-}
-
 func simpleDiffViaPlanState(
 	ctx context.Context,
 	res *schema.Resource,
@@ -154,7 +153,12 @@ func simpleDiffViaPlanState(
 	}
 
 	planned := terraform.NewResourceConfigShimmed(proposedNewStateVal, res.CoreConfigSchema())
-	return res.SimpleDiff(ctx, instanceStateWithUpdatedRawPlan(s, proposedNewStateVal), planned, meta)
+	state := s.DeepCopy()
+	state.RawPlan = proposedNewStateVal
+	if state.RawState.IsNull() {
+		state.RawState = priorStateVal
+	}
+	return res.SimpleDiff(ctx, state, planned, meta)
 }
 
 func showDiffChangeType(b byte) string {