fix: refuse removals that would drop record extension fields

A capable writer may no longer remove or replace a known object, map entry or array when its stored value carries members the writer's schema does not know. The check that guarded changed arrays now runs wherever the merge drops or overwrites a stored value instead of merging into it: a known member deleted or set to null, a map entry removed or stored under a key the typed decode rewrote, a value of another JSON kind, and a changed array. Encode returns ErrRecordExtensionConflict naming the changed container and the first unknown member. Removal stays allowed when the stored value has no extensions, and a change merged member by member still keeps unknown members.
2026-10-15 09:01:31 +00:00 · 2026-10-11 08:59:14 +10:00
parent c6fdc239ec
commit 0d7b4266d3
3 changed files with 210 additions and 30 deletions
@@ -45,7 +45,7 @@ var (
 	ErrRecordFormatOutOfRange   = errors.New("instance record format is below 1")
 	ErrRecordFormatAhead        = errors.New("instance record format is newer than this binary can write")
 	ErrRecordSectionsNeedFormat = errors.New("assignment and observation need record format 2")
-	ErrRecordExtensionConflict  = errors.New("instance record change would drop an unknown member inside an array")
+	ErrRecordExtensionConflict  = errors.New("instance record change would drop an unknown member")
 )

 // AssignmentReason says why an assignment changed.
@@ -192,8 +192,8 @@ func DecodeTransitionRecord(data []byte) (*TransitionRecord, error) {
 }

 // Encode writes only the typed paths that differ from the decoded snapshot
-// over the input, keeping unknown members. A changed array is replaced whole,
-// and refused if its read elements carry unknown members. It adds no marker.
+// over the input, keeping unknown members. It refuses to remove or replace a
+// value whose stored form carries unknown members, and adds no format marker.
 func (r *TransitionRecord) Encode() ([]byte, error) {
 	if r.readOnly {
 		return nil, fmt.Errorf("%w: format %d", ErrRecordFormatAhead, r.format)
@@ -249,38 +249,38 @@ func mergeChanged(s *fieldSchema, raw, snapshot, current any, path string) (any,
 	if reflect.DeepEqual(snapshot, current) {
 		return raw, nil
 	}
-	if rawArr, ok := raw.([]any); ok && s != nil {
-		if at, found := unknownMember(s, rawArr, path); found {
-			return nil, fmt.Errorf("%w: changed array %s would drop unknown member %s", ErrRecordExtensionConflict, path, at)
-		}
-	}
 	rawObj, rawIsObj := raw.(map[string]any)
 	snapObj, snapIsObj := snapshot.(map[string]any)
 	curObj, curIsObj := current.(map[string]any)
 	if s == nil || (s.fields == nil && !s.isMap) || !rawIsObj || !snapIsObj || !curIsObj {
+		if err := refuseDrop(s, raw, path); err != nil {
+			return nil, err
+		}
 		return current, nil
 	}
 	out := make(map[string]any, len(rawObj))
-	for k, v := range rawObj {
+	for _, k := range slices.Sorted(maps.Keys(rawObj)) {
 		_, inSnap := snapObj[k]
 		_, inCur := curObj[k]
 		switch {
 		case !s.isMap && (!inSnap || inCur):
-			out[k] = v // unknown, or known and still present
+			out[k] = rawObj[k] // unknown, or known and still present
 		case s.isMap && inSnap && inCur:
-			out[k] = v // a map key the typed decode did not normalise
+			out[k] = rawObj[k] // a map key the typed decode did not normalise
+		default:
+			// A removed member or entry, or a map key the typed decode rewrote.
+			if err := refuseDrop(childSchema(s, k), rawObj[k], path+"."+k); err != nil {
+				return nil, err
+			}
 		}
 	}
-	for k, curV := range curObj {
+	for _, k := range slices.Sorted(maps.Keys(curObj)) {
+		curV := curObj[k]
 		snapV, inSnap := snapObj[k]
 		rawV, inRaw := out[k]
-		child := s.elem
-		if !s.isMap {
-			child = s.fields[k]
-		}
 		switch {
 		case inSnap && inRaw:
-			merged, err := mergeChanged(child, rawV, snapV, curV, path+"."+k)
+			merged, err := mergeChanged(childSchema(s, k), rawV, snapV, curV, path+"."+k)
 			if err != nil {
 				return nil, err
 			}
@@ -288,12 +288,44 @@ func mergeChanged(s *fieldSchema, raw, snapshot, current any, path string) (any,
 		case inSnap && !s.isMap && reflect.DeepEqual(snapV, curV):
 			// A zero member the input left out stays out.
 		default:
+			if inRaw {
+				if err := refuseDrop(childSchema(s, k), rawV, path+"."+k); err != nil {
+					return nil, err
+				}
+			}
 			out[k] = curV
 		}
 	}
 	return out, nil
 }

+func childSchema(s *fieldSchema, key string) *fieldSchema {
+	if s.isMap {
+		return s.elem
+	}
+	return s.fields[key]
+}
+
+// refuseDrop fails when the raw value about to be removed or replaced holds a
+// member its schema does not know, so a writer cannot erase what it cannot read.
+func refuseDrop(s *fieldSchema, raw any, path string) error {
+	at, found := unknownMember(s, raw, path)
+	if !found {
+		return nil
+	}
+	kind := "value"
+	switch raw.(type) {
+	case []any:
+		kind = "array"
+	case map[string]any:
+		kind = "object"
+		if s.isMap {
+			kind = "map"
+		}
+	}
+	return fmt.Errorf("%w: changed %s %s would drop unknown member %s", ErrRecordExtensionConflict, kind, path, at)
+}
+
 // unknownMember finds the first member of v, at any depth, that schema s does
 // not know. An opaque value (nil schema) is not inspected.
 func unknownMember(s *fieldSchema, v any, path string) (string, bool) {
@@ -0,0 +1,95 @@
+package state
+
+import (
+	"bytes"
+	"encoding/json"
+	"reflect"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// The record has no typed map with object values and no non-omitempty object
+// pointer, so these shapes exercise the merge rules the record cannot reach.
+type probeInner struct {
+	Name string `json:"name"`
+}
+
+type probe struct {
+	Entries  map[string]probeInner `json:"entries,omitempty"`
+	Numbered map[int]probeInner    `json:"numbered,omitempty"`
+	Ptr      *probeInner           `json:"ptr"`
+}
+
+func probeTree(t *testing.T, b []byte) any {
+	t.Helper()
+	dec := json.NewDecoder(bytes.NewReader(b))
+	dec.UseNumber()
+	var v any
+	require.NoError(t, dec.Decode(&v))
+	return v
+}
+
+func mergeProbe(t *testing.T, doc string, mutate func(*probe)) (any, error) {
+	t.Helper()
+	var p probe
+	require.NoError(t, json.Unmarshal([]byte(doc), &p))
+	snap, err := json.Marshal(&p)
+	require.NoError(t, err)
+	mutate(&p)
+	cur, err := json.Marshal(&p)
+	require.NoError(t, err)
+	s := schemaOf(reflect.TypeFor[probe](), map[reflect.Type]*fieldSchema{})
+	return mergeChanged(s, probeTree(t, []byte(doc)), probeTree(t, snap), probeTree(t, cur), "$")
+}
+
+func TestMergeChanged_ProbeShapes(t *testing.T) {
+	cases := []struct {
+		name, doc string
+		mutate    func(*probe)
+		conflict  string // container and member, or empty for success
+		want      string
+	}{
+		{"map entry with an unknown member removed",
+			`{"entries":{"a":{"name":"x","future":1},"b":{"name":"y"}},"ptr":null}`,
+			func(p *probe) { delete(p.Entries, "a") },
+			"changed object $.entries.a would drop unknown member $.entries.a.future", ""},
+		{"map entry without unknown members removed",
+			`{"entries":{"a":{"name":"x"},"b":{"name":"y","future":1}},"ptr":null}`,
+			func(p *probe) { delete(p.Entries, "a") },
+			"", `{"entries":{"b":{"name":"y","future":1}},"ptr":null}`},
+		{"map entry changed in place keeps its unknown member",
+			`{"entries":{"a":{"name":"x","future":1}},"ptr":null}`,
+			func(p *probe) { p.Entries["a"] = probeInner{Name: "z"} },
+			"", `{"entries":{"a":{"name":"z","future":1}},"ptr":null}`},
+		{"map entry under a rewritten key replaced",
+			`{"numbered":{"07":{"name":"x","future":1}},"ptr":null}`,
+			func(p *probe) { p.Numbered[8] = probeInner{Name: "y"} },
+			"changed object $.numbered.07 would drop unknown member $.numbered.07.future", ""},
+		{"object with an unknown member replaced by null",
+			`{"ptr":{"name":"x","future":{"deep":true}}}`,
+			func(p *probe) { p.Ptr = nil },
+			"changed object $.ptr would drop unknown member $.ptr.future", ""},
+		{"object without unknown members replaced by null",
+			`{"ptr":{"name":"x"}}`,
+			func(p *probe) { p.Ptr = nil },
+			"", `{"ptr":null}`},
+		{"keys of a map are data, not extensions",
+			`{"entries":{"anything":{"name":"x"}},"ptr":null}`,
+			func(p *probe) { p.Entries = nil },
+			"", `{"ptr":null}`},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := mergeProbe(t, tc.doc, tc.mutate)
+			if tc.conflict != "" {
+				require.ErrorIs(t, err, ErrRecordExtensionConflict)
+				assert.Contains(t, err.Error(), tc.conflict)
+				return
+			}
+			require.NoError(t, err)
+			assert.Equal(t, probeTree(t, []byte(tc.want)), got)
+		})
+	}
+}
@@ -176,27 +176,83 @@ func TestTransitionRecord_MutationChangesOnlyItsPath(t *testing.T) {
 	}
 }

+// plainDoc carries no unknown member inside any known value, so removing a
+// known value drops nothing a newer writer added.
+const plainDoc = `{"format":2,"metadata":{"name":"i-1","owner_refs":[{"kind":"cluster","name":"c-1"}]},` +
+	`"spec":{"run_instances_input":{"MaxCount":1,"Placement":{"AvailabilityZone":"a"}}},` +
+	`"assignment":{"node":"node-1","generation":3},"future":{"keep":[{"x":1}]}}`
+
 func TestTransitionRecord_ZeroedOmitemptyFieldIsRemoved(t *testing.T) {
 	cases := []struct {
-		name   string
-		mutate func(*state.TransitionRecord)
-		gone   string
+		name, doc string
+		mutate    func(*state.TransitionRecord)
+		gone      string
 	}{
-		{"string", func(r *state.TransitionRecord) { r.Record.Status.AZ = "" }, "status.az"},
-		{"map", func(r *state.TransitionRecord) { r.Record.Metadata.Tags = nil }, "metadata.tags"},
-		{"slice", func(r *state.TransitionRecord) { r.Record.Metadata.OwnerRefs = nil }, "metadata.owner_refs"},
-		{"omitzero time", func(r *state.TransitionRecord) { r.Record.Status.TerminatedAt = time.Time{} }, "status.terminated_at"},
-		{"pointer", func(r *state.TransitionRecord) { r.Record.Spec.RunInstancesInput = nil }, "spec.run_instances_input"},
+		{"string", goldenV1, func(r *state.TransitionRecord) { r.Record.Status.AZ = "" }, "status.az"},
+		{"map of strings", goldenV1, func(r *state.TransitionRecord) { r.Record.Metadata.Tags = nil }, "metadata.tags"},
+		{"omitzero time", goldenV1, func(r *state.TransitionRecord) { r.Record.Status.TerminatedAt = time.Time{} }, "status.terminated_at"},
+		{"slice", plainDoc, func(r *state.TransitionRecord) { r.Record.Metadata.OwnerRefs = nil }, "metadata.owner_refs"},
+		{"pointer", plainDoc, func(r *state.TransitionRecord) { r.Record.Spec.RunInstancesInput = nil }, "spec.run_instances_input"},
+		{"section", plainDoc, func(r *state.TransitionRecord) { r.Assignment = nil }, "assignment"},
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			r := decode(t, goldenV1)
+			r := decode(t, tc.doc)
 			tc.mutate(r)
-			assert.Equal(t, edit(t, goldenV1, map[string]any{tc.gone: nil}), canonical(t, encode(t, r)))
+			assert.Equal(t, edit(t, tc.doc, map[string]any{tc.gone: nil}), canonical(t, encode(t, r)))
 		})
 	}
 }

+// Removing or replacing a known value whose stored form carries an unknown
+// member would erase state this writer cannot read, so Encode refuses it.
+func TestTransitionRecord_RemovalOfExtendedValueIsRefused(t *testing.T) {
+	const extended = `{"metadata":{"name":"i-1"},` +
+		`"spec":{"run_instances_input":{"MaxCount":1,"Placement":{"AvailabilityZone":"a","FuturePlacement":1}}}}`
+	cases := []struct {
+		name, doc, container, member string
+		mutate                       func(*state.TransitionRecord)
+	}{
+		{"object with a nested unknown member", extended,
+			"changed object $.spec.run_instances_input ", "$.spec.run_instances_input.Placement.FuturePlacement",
+			func(r *state.TransitionRecord) { r.Record.Spec.RunInstancesInput = nil }},
+		{"AWS SDK object replaced by null", extended,
+			"changed object $.spec.run_instances_input.Placement ", "$.spec.run_instances_input.Placement.FuturePlacement",
+			func(r *state.TransitionRecord) { r.Record.Spec.RunInstancesInput.Placement = nil }},
+		{"array with an unknown element member", goldenV1,
+			"changed array $.metadata.owner_refs ", "$.metadata.owner_refs[0].future_ref",
+			func(r *state.TransitionRecord) { r.Record.Metadata.OwnerRefs = nil }},
+		{"format-2 section", goldenV2,
+			"changed object $.assignment ", "$.assignment.future_assignment",
+			func(r *state.TransitionRecord) { r.Assignment = nil }},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			r := decode(t, tc.doc)
+			tc.mutate(r)
+			out, err := r.Encode()
+			require.ErrorIs(t, err, state.ErrRecordExtensionConflict)
+			assert.Contains(t, err.Error(), tc.container)
+			assert.Contains(t, err.Error(), "would drop unknown member "+tc.member)
+			assert.Nil(t, out)
+		})
+	}
+
+	t.Run("AWS SDK object without unknown members replaced by null", func(t *testing.T) {
+		r := decode(t, plainDoc)
+		r.Record.Spec.RunInstancesInput.Placement = nil
+		placement, _ := member(t, encode(t, r), "spec", "run_instances_input", "Placement")
+		assert.JSONEq(t, `null`, string(placement))
+	})
+
+	t.Run("sibling change inside an extended object keeps the extension", func(t *testing.T) {
+		r := decode(t, extended)
+		r.Record.Spec.RunInstancesInput.Placement.AvailabilityZone = aws.String("b")
+		want := edit(t, extended, map[string]any{"spec.run_instances_input.Placement.AvailabilityZone": "b"})
+		assert.Equal(t, want, canonical(t, encode(t, r)))
+	})
+}
+
 // terminated_at and hostfwd_port_map are read in forms that the typed encoder
 // rewrites, so these cases tell retention of the read value from re-encoding.
 func TestTransitionRecord_UnchangedKnownFieldKeepsItsRawForm(t *testing.T) {
@@ -335,9 +391,6 @@ func TestTransitionRecord_V2SectionMutationKeepsNestedUnknownMembers(t *testing.
 		{"observation state", func(r *state.TransitionRecord) {
 			r.Observation.State = state.ObservedUnconfirmed
 		}, map[string]any{"observation.state": "unconfirmed"}},
-		{"assignment removed", func(r *state.TransitionRecord) {
-			r.Assignment = nil
-		}, map[string]any{"assignment": nil}},
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {