From e0d3b97c8e0e8091473537d8fc14601f06964000 Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Fri, 8 Jan 2021 09:56:03 -0500
Subject: [PATCH] preserve resource private data on error

If the provider returns an empty response during apply, we restore the
prior state for the resource, but the private data was not being
restored.
---
 terraform/context_apply_test.go              | 42 ++++++++++++++++++++
 terraform/node_resource_abstract_instance.go | 25 +++++-------
 2 files changed, 51 insertions(+), 16 deletions(-)

diff --git a/terraform/context_apply_test.go b/terraform/context_apply_test.go
index 3b5901d96..5d15472f4 100644
--- a/terraform/context_apply_test.go
+++ b/terraform/context_apply_test.go
@@ -12468,3 +12468,45 @@ func TestContext2Apply_dataSensitive(t *testing.T) {
 		t.Errorf("wrong marks\n got: %#v\nwant: %#v", gotMarks, wantMarks)
 	}
 }
+
+func TestContext2Apply_errorRestorePrivateData(t *testing.T) {
+	// empty config to remove our resource
+	m := testModuleInline(t, map[string]string{
+		"main.tf": "",
+	})
+
+	p := simpleMockProvider()
+	p.ApplyResourceChangeResponse = &providers.ApplyResourceChangeResponse{
+		// we error during apply, which will trigger core to preserve the last
+		// known state, including private data
+		Diagnostics: tfdiags.Diagnostics(nil).Append(errors.New("oops")),
+	}
+
+	addr := mustResourceInstanceAddr("test_object.a")
+
+	state := states.BuildState(func(s *states.SyncState) {
+		s.SetResourceInstanceCurrent(addr, &states.ResourceInstanceObjectSrc{
+			Status:    states.ObjectReady,
+			AttrsJSON: []byte(`{"id":"foo"}`),
+			Private:   []byte("private"),
+		}, mustProviderConfig(`provider["registry.terraform.io/hashicorp/test"]`))
+	})
+
+	ctx := testContext2(t, &ContextOpts{
+		Config: m,
+		State:  state,
+		Providers: map[addrs.Provider]providers.Factory{
+			addrs.NewDefaultProvider("test"): testProviderFuncFixed(p),
+		},
+	})
+
+	_, diags := ctx.Plan()
+	if diags.HasErrors() {
+		t.Fatal(diags.Err())
+	}
+
+	state, _ = ctx.Apply()
+	if string(state.ResourceInstance(addr).Current.Private) != "private" {
+		t.Fatal("missing private data in state")
+	}
+}
diff --git a/terraform/node_resource_abstract_instance.go b/terraform/node_resource_abstract_instance.go
index 3278759e5..b3d8e5fec 100644
--- a/terraform/node_resource_abstract_instance.go
+++ b/terraform/node_resource_abstract_instance.go
@@ -1821,16 +1821,16 @@ func (n *NodeAbstractResourceInstance) apply(
 	if state == nil {
 		state = &states.ResourceInstanceObject{}
 	}
-	var newState *states.ResourceInstanceObject
+
 	provider, providerSchema, err := getProvider(ctx, n.ResolvedProvider)
 	if err != nil {
-		return newState, diags.Append(err), applyError
+		return nil, diags.Append(err), applyError
 	}
 	schema, _ := providerSchema.SchemaForResourceType(n.Addr.Resource.Resource.Mode, n.Addr.Resource.Resource.Type)
 	if schema == nil {
 		// Should be caught during validation, so we don't bother with a pretty error here
 		diags = diags.Append(fmt.Errorf("provider does not support resource type %q", n.Addr.Resource.Resource.Type))
-		return newState, diags, applyError
+		return nil, diags, applyError
 	}
 
 	log.Printf("[INFO] Starting apply for %s", n.Addr)
@@ -1843,7 +1843,7 @@ func (n *NodeAbstractResourceInstance) apply(
 		configVal, _, configDiags = ctx.EvaluateBlock(applyConfig.Config, schema, nil, keyData)
 		diags = diags.Append(configDiags)
 		if configDiags.HasErrors() {
-			return newState, diags, applyError
+			return nil, diags, applyError
 		}
 	}
 
@@ -1852,13 +1852,13 @@ func (n *NodeAbstractResourceInstance) apply(
 			"configuration for %s still contains unknown values during apply (this is a bug in Terraform; please report it!)",
 			n.Addr,
 		))
-		return newState, diags, applyError
+		return nil, diags, applyError
 	}
 
 	metaConfigVal, metaDiags := n.providerMetas(ctx)
 	diags = diags.Append(metaDiags)
 	if diags.HasErrors() {
-		return newState, diags, applyError
+		return nil, diags, applyError
 	}
 
 	log.Printf("[DEBUG] %s: applying the planned %s change", n.Addr, change.Action)
@@ -1870,6 +1870,7 @@ func (n *NodeAbstractResourceInstance) apply(
 	unmarkedBefore, beforePaths := change.Before.UnmarkDeepWithPaths()
 	unmarkedAfter, afterPaths := change.After.UnmarkDeepWithPaths()
 
+	var newState *states.ResourceInstanceObject
 	// If we have an Update action, our before and after values are equal,
 	// and only differ on their sensitivity, the newVal is the after val
 	// and we should not communicate with the provider. We do need to update
@@ -2071,8 +2072,6 @@ func (n *NodeAbstractResourceInstance) apply(
 		}
 	}
 
-	newStatus := states.ObjectReady
-
 	// Sometimes providers return a null value when an operation fails for some
 	// reason, but we'd rather keep the prior state so that the error can be
 	// corrected on a subsequent run. We must only do this for null new value
@@ -2084,18 +2083,12 @@ func (n *NodeAbstractResourceInstance) apply(
 		// deleted then our next refresh will detect that and fix it up.
 		// If change.Action is Create then change.Before will also be null,
 		// which is fine.
-		newVal = change.Before
-
-		// If we're recovering the previous state, we also want to restore the
-		// the tainted status of the object.
-		if state.Status == states.ObjectTainted {
-			newStatus = states.ObjectTainted
-		}
+		newState = state.DeepCopy()
 	}
 
 	if !newVal.IsNull() { // null value indicates that the object is deleted, so we won't set a new state in that case
 		newState = &states.ResourceInstanceObject{
-			Status:              newStatus,
+			Status:              states.ObjectReady,
 			Value:               newVal,
 			Private:             resp.Private,
 			CreateBeforeDestroy: createBeforeDestroy,