prevent mutation of dictionary keys

migrations/entitlements/migration_test.go

@@ -730,6 +730,7 @@ func convertEntireTestValue(
 			testEntitlementsMigration{inter: inter},
 		},
 		reporter,
+		true,
 	)
 
 	err = migration.Commit()

migrations/migration.go

@@ -155,6 +155,7 @@ func (m *StorageMigration) NewValueMigrationsPathMigrator(
 				value,
 				valueMigrations,
 				reporter,
+				true,
 			)
 		},
 	)
@@ -168,6 +169,7 @@ func (m *StorageMigration) MigrateNestedValue(
 	value interpreter.Value,
 	valueMigrations []ValueMigration,
 	reporter Reporter,
+	allowMutation bool,
 ) (migratedValue interpreter.Value) {
 
 	defer func() {
@@ -226,6 +228,7 @@ func (m *StorageMigration) MigrateNestedValue(
 			innerValue,
 			valueMigrations,
 			reporter,
+			allowMutation,
 		)
 		if newInnerValue != nil {
 			migratedValue = interpreter.NewSomeValueNonCopying(inter, newInnerValue)
@@ -249,12 +252,21 @@ func (m *StorageMigration) MigrateNestedValue(
 				element,
 				valueMigrations,
 				reporter,
+				allowMutation,
 			)
 
 			if newElement == nil {
 				continue
 			}
 
+			if !allowMutation {
+				panic(errors.NewUnexpectedError(
+					"mutation not allowed: attempting to migrate array element at index %d: %s",
+					index,
+					element,
+				))
+			}
+
 			existingStorable := array.RemoveWithoutTransfer(
 				inter,
 				emptyLocationRange,
@@ -297,12 +309,21 @@ func (m *StorageMigration) MigrateNestedValue(
 				existingValue,
 				valueMigrations,
 				reporter,
+				allowMutation,
 			)
 
 			if migratedValue == nil {
 				continue
 			}
 
+			if !allowMutation {
+				panic(errors.NewUnexpectedError(
+					"mutation not allowed: attempting to migrate composite value field %s: %s",
+					fieldName,
+					existingValue,
+				))
+			}
+
 			composite.SetMemberWithoutTransfer(
 				inter,
 				emptyLocationRange,
@@ -323,8 +344,23 @@ func (m *StorageMigration) MigrateNestedValue(
 		// The mutating iterator is only able to read new keys,
 		// as it recalculates the stored values' hashes.
 
-		m.migrateDictionaryKeys(storageKey, storageMapKey, dictionary, valueMigrations, reporter)
-		m.migrateDictionaryValues(storageKey, storageMapKey, dictionary, valueMigrations, reporter)
+		m.migrateDictionaryKeys(
+			storageKey,
+			storageMapKey,
+			dictionary,
+			valueMigrations,
+			reporter,
+			allowMutation,
+		)
+
+		m.migrateDictionaryValues(
+			storageKey,
+			storageMapKey,
+			dictionary,
+			valueMigrations,
+			reporter,
+			allowMutation,
+		)
 
 	case *interpreter.PublishedValue:
 		publishedValue := typedValue
@@ -334,6 +370,7 @@ func (m *StorageMigration) MigrateNestedValue(
 			publishedValue.Value,
 			valueMigrations,
 			reporter,
+			allowMutation,
 		)
 		if newInnerValue != nil {
 			newInnerCapability := newInnerValue.(interpreter.CapabilityValue)
@@ -418,6 +455,7 @@ func (m *StorageMigration) migrateDictionaryKeys(
 	dictionary *interpreter.DictionaryValue,
 	valueMigrations []ValueMigration,
 	reporter Reporter,
+	allowMutation bool,
 ) {
 	inter := m.interpreter
 
@@ -442,12 +480,21 @@ func (m *StorageMigration) migrateDictionaryKeys(
 			existingKey,
 			valueMigrations,
 			reporter,
+			// NOTE: Mutation of keys is not allowed.
+			false,
 		)
 
 		if newKey == nil {
 			continue
 		}
 
+		if !allowMutation {
+			panic(errors.NewUnexpectedError(
+				"mutation not allowed: attempting to migrate dictionary key: %s",
+				existingKey,
+			))
+		}
+
 		// We only reach here because key needs to be migrated.
 
 		// Remove the old key-value pair
@@ -507,6 +554,7 @@ func (m *StorageMigration) migrateDictionaryKeys(
 				existingValue,
 				valueMigrations,
 				reporter,
+				allowMutation,
 			)
 
 			var valueToSet interpreter.Value
@@ -583,6 +631,7 @@ func (m *StorageMigration) migrateDictionaryValues(
 	dictionary *interpreter.DictionaryValue,
 	valueMigrations []ValueMigration,
 	reporter Reporter,
+	allowMutation bool,
 ) {
 
 	inter := m.interpreter
@@ -621,12 +670,20 @@ func (m *StorageMigration) migrateDictionaryValues(
 			existingValue,
 			valueMigrations,
 			reporter,
+			allowMutation,
 		)
 
 		if newValue == nil {
 			continue
 		}
 
+		if !allowMutation {
+			panic(errors.NewUnexpectedError(
+				"mutation not allowed: attempting to migrate dictionary value: %s",
+				existingValue,
+			))
+		}
+
 		// Set new value with existing key in the dictionary.
 		existingValueStorable := dictionary.InsertWithoutTransfer(
 			inter,

migrations/migration_test.go

@@ -2971,3 +2971,135 @@ func TestDictionaryWithEnumKey(t *testing.T) {
 	})()
 
 }
+
+func TestDictionaryKeyMutationMigration(t *testing.T) {
+
+	t.Parallel()
+
+	testAddress := common.MustBytesToAddress([]byte{0x1})
+	storagePathDomain := common.PathDomainStorage.Identifier()
+	storageMapKey := interpreter.StringStorageMapKey("test")
+
+	ledger := NewTestLedger(nil, nil)
+
+	newStorageAndInterpreter := func(t *testing.T) (*runtime.Storage, *interpreter.Interpreter) {
+		storage := runtime.NewStorage(ledger, nil)
+		inter, err := interpreter.NewInterpreter(
+			nil,
+			utils.TestLocation,
+			&interpreter.Config{
+				Storage:                     storage,
+				AtreeValueValidationEnabled: true,
+				// NOTE: disabled, as storage is not expected to be always valid _during_ migration
+				AtreeStorageValidationEnabled: false,
+			},
+		)
+		require.NoError(t, err)
+
+		return storage, inter
+	}
+
+	fooQualifiedIdentifier := "Test.Foo"
+	fooType := &interpreter.CompositeStaticType{
+		Location:            utils.TestLocation,
+		QualifiedIdentifier: fooQualifiedIdentifier,
+		TypeID:              utils.TestLocation.TypeID(nil, fooQualifiedIdentifier),
+	}
+
+	// Prepare
+	(func() {
+		storage, inter := newStorageAndInterpreter(t)
+
+		storageMap := storage.GetStorageMap(
+			testAddress,
+			storagePathDomain,
+			true,
+		)
+
+		// {Test.Foo: Int8}
+		dictionaryValue := interpreter.NewDictionaryValueWithAddress(
+			inter,
+			emptyLocationRange,
+			interpreter.NewDictionaryStaticType(
+				nil,
+				fooType,
+				interpreter.PrimitiveStaticTypeInt8,
+			),
+			testAddress,
+		)
+
+		// Write the dictionary value to storage before inserting values into dictionary,
+		// as the insertion of values into the dictionary triggers a storage health check,
+		// which fails if the dictionary value is not yet stored (unreferenced slabs)
+
+		storageMap.WriteValue(
+			inter,
+			storageMapKey,
+			dictionaryValue,
+		)
+
+		dictionaryKey := interpreter.NewCompositeValue(
+			inter,
+			emptyLocationRange,
+			utils.TestLocation,
+			"Foo",
+			common.CompositeKindEnum,
+			[]interpreter.CompositeField{
+				{
+					Name:  sema.EnumRawValueFieldName,
+					Value: interpreter.Int8Value(10),
+				},
+			},
+			testAddress,
+		)
+
+		dictionaryValue.InsertWithoutTransfer(
+			inter,
+			emptyLocationRange,
+			dictionaryKey,
+			interpreter.Int8Value(100),
+		)
+
+		err := storage.Commit(inter, false)
+		require.NoError(t, err)
+
+		err = storage.CheckHealth()
+		require.NoError(t, err)
+	})()
+
+	// Migrate
+	(func() {
+
+		storage, inter := newStorageAndInterpreter(t)
+
+		migration, err := NewStorageMigration(inter, storage, "test", testAddress)
+		require.NoError(t, err)
+
+		reporter := newTestReporter()
+
+		migration.Migrate(
+			migration.NewValueMigrationsPathMigrator(
+				reporter,
+				testInt8Migration{},
+			),
+		)
+
+		err = migration.Commit()
+		require.NoError(t, err)
+
+		// Assert
+
+		require.Len(t, reporter.errors, 1)
+		assert.ErrorContains(
+			t,
+			reporter.errors[0],
+			"mutation not allowed: attempting to migrate composite value field rawValue",
+		)
+
+		assert.Len(t, reporter.migrated, 1)
+
+		err = storage.CheckHealth()
+		require.NoError(t, err)
+	})()
+
+}