test: add key assignment to model and driver (#1573)

* Start implementing key assignment in Quint

* Update model

* Update model

* Update model

* Add key assignment to functional layer

* Start writing test for key assignment

* Continue adding key assignment to model

* Fix some behaviours in key assignment logic

* Pull key assignment into its own module

* Split key assignment into extra module

* Start updating invariants for key assignment

* Update key assignment spec

* Merge key assignment into main model

* Update spec, fix bug

* Fix map building for key assignment

* Fix key assignment invariants

* Update driver to include key assignment

* Fix key assignment in bounded steps

* Finish key assignment integration in driver

* Update clients every block

* Add key assignment to invariant check

* Use key assignment in its step

* Revert changes to extraSpells

* Add check to avoid empty set of running consumers

* Update tests/mbt/model/ccv.qnt

Co-authored-by: insumity <[email protected]>

* Update tests/mbt/model/ccv.qnt

Co-authored-by: insumity <[email protected]>

* Incorporate review comments

* Refactor keyAssignmentReplacement

* Start refactoring a few things

* Refactor, take into account comments

* Refactor utility functions to type module

* Minor refactor

* Refactor utils into own file

* Add key assignment invariant checks to make target

* Add empty-list-handling to inv

* Start refactoring keys of sentVscPackets map

* Fix keys for packet map

* Fix path to model

* Add more comments, remove unnecessary disjuncts

---------

Co-authored-by: insumity <[email protected]>

Makefile

@@ -131,7 +131,8 @@ test-trace:
 verify-models:
 	quint test tests/mbt/model/ccv_test.qnt;\
 	quint test tests/mbt/model/ccv_model.qnt;\
-	quint run --invariant "all{ValidatorUpdatesArePropagatedInv,ValidatorSetHasExistedInv,SameVscPacketsInv,MatureOnTimeInv,EventuallyMatureOnProviderInv}" tests/mbt/model/ccv_model.qnt --max-steps 200 --max-samples 200
+	quint run --invariant "all{ValidatorUpdatesArePropagatedInv,ValidatorSetHasExistedInv,SameVscPacketsInv,MatureOnTimeInv,EventuallyMatureOnProviderInv}" tests/mbt/model/ccv_model.qnt --max-steps 200 --max-samples 200;\
+	quint run --invariant "all{ValidatorUpdatesArePropagatedKeyAssignmentInv,ValidatorSetHasExistedKeyAssignmentInv,SameVscPacketsKeyAssignmentInv,MatureOnTimeInv,EventuallyMatureOnProviderInv,KeyAssignmentRulesInv}" tests/mbt/model/ccv_model.qnt --step stepKeyAssignment --max-steps 200 --max-samples 200
 
 
 

tests/mbt/driver/core.go

@@ -20,6 +20,7 @@ import (
 	abcitypes "github.com/cometbft/cometbft/abci/types"
 	cmttypes "github.com/cometbft/cometbft/types"
 
+	"github.com/cometbft/cometbft/proto/tendermint/crypto"
 	appConsumer "github.com/cosmos/interchain-security/v4/app/consumer"
 	appProvider "github.com/cosmos/interchain-security/v4/app/provider"
 	simibc "github.com/cosmos/interchain-security/v4/testutil/simibc"
@@ -123,9 +124,13 @@ func (s *Driver) consumerPower(i int64, chain ChainId) (int64, error) {
 	return v.Power, nil
 }
 
+func (s *Driver) stakingValidator(i int64) (stakingtypes.Validator, bool) {
+	return s.providerStakingKeeper().GetValidator(s.ctx(PROVIDER), s.validator(i))
+}
+
 // providerPower returns the power(=number of bonded tokens) of the i-th validator on the provider.
 func (s *Driver) providerPower(i int64) (int64, error) {
-	v, found := s.providerStakingKeeper().GetValidator(s.ctx(PROVIDER), s.validator(i))
+	v, found := s.stakingValidator(i)
 	if !found {
 		return 0, fmt.Errorf("validator with id %v not found on provider", i)
 	} else {
@@ -370,6 +375,14 @@ func (s *Driver) setTime(chain ChainId, newTime time.Time) {
 	testChain.App.BeginBlock(abcitypes.RequestBeginBlock{Header: testChain.CurrentHeader})
 }
 
+func (s *Driver) AssignKey(chain ChainId, valIndex int64, value crypto.PublicKey) error {
+	stakingVal, found := s.stakingValidator(valIndex)
+	if !found {
+		return fmt.Errorf("validator with id %v not found on provider", valIndex)
+	}
+	return s.providerKeeper().AssignConsumerKey(s.providerCtx(), string(chain), stakingVal, value)
+}
+
 // DeliverPacketToConsumer delivers a packet from the provider to the given consumer recipient.
 // It updates the client before delivering the packet.
 // Since the channel is ordered, the packet that is delivered is the first packet in the outbox.

tests/mbt/driver/generate_more_traces.sh

@@ -9,4 +9,5 @@ go run ./... -modelPath=../model/ccv_boundeddrift.qnt -step stepBoundedDrift -in
 echo "Generating synced traces with maturations"
 go run ./... -modelPath=../model/ccv_sync.qnt -init initSync -step stepSync -invariant CanReceiveMaturations -traceFolder traces/sync_mat -numTraces 20 -numSteps 300 -numSamples 20
 echo "Generating long synced traces without invariants"
-go run ./... -modelPath=../model/ccv_sync.qnt -init initSync -step stepSync -traceFolder traces/sync_noinv -numTraces 20 -numSteps 500 -numSamples 1
+go run ./... -modelPath=../model/ccv_sync.qnt -init initSync -step stepSync -traceFolder traces/sync_noinv -numTraces 20 -numSteps 500 -numSamples 1
+go run ./... -modelPath=../model/ccv_boundeddrift.qnt --step stepBoundedDriftKeyAssignment --traceFolder traces/bound_key -numTraces 20 -numSteps 100 -numSamples 20

tests/mbt/driver/generate_traces.sh

@@ -9,4 +9,5 @@ go run ./... -modelPath=../model/ccv_boundeddrift.qnt -step stepBoundedDrift -in
 echo "Generating synced traces with maturations"
 go run ./... -modelPath=../model/ccv_sync.qnt -init initSync -step stepSync -invariant CanReceiveMaturations -traceFolder traces/sync_mat -numTraces 1 -numSteps 300 -numSamples 20
 echo "Generating long synced traces without invariants"
-go run ./... -modelPath=../model/ccv_sync.qnt -init initSync -step stepSync -traceFolder traces/sync_noinv -numTraces 1 -numSteps 500 -numSamples 1
+go run ./... -modelPath=../model/ccv_sync.qnt -init initSync -step stepSync -traceFolder traces/sync_noinv -numTraces 1 -numSteps 500 -numSamples 1
+go run ./... -modelPath=../model/ccv_boundeddrift.qnt --step stepBoundedDriftKeyAssignment --traceFolder traces/bound_key -numTraces 1 -numSteps 100 -numSamples 20

tests/mbt/driver/mbt_test.go

@@ -15,10 +15,13 @@ import (
 	"github.com/kylelemons/godebug/pretty"
 	"github.com/stretchr/testify/require"
 
-	sdktypes "github.com/cosmos/cosmos-sdk/types"
-
 	cmttypes "github.com/cometbft/cometbft/types"
 
+	tmencoding "github.com/cometbft/cometbft/crypto/encoding"
+	"github.com/cosmos/interchain-security/v4/testutil/integration"
+
+	sdktypes "github.com/cosmos/cosmos-sdk/types"
+
 	providertypes "github.com/cosmos/interchain-security/v4/x/ccv/provider/types"
 )
 
@@ -69,6 +72,7 @@ func TestMBT(t *testing.T) {
 	t.Logf("Number of sent packets: %v", stats.numSentPackets)
 	t.Logf("Number of blocks: %v", stats.numBlocks)
 	t.Logf("Number of transactions: %v", stats.numTxs)
+	t.Logf("Number of key assignments: %v", stats.numKeyAssignments)
 	t.Logf("Average summed block time delta passed per trace: %v", stats.totalBlockTimePassedPerTrace/time.Duration(numTraces))
 }
 
@@ -117,6 +121,21 @@ func RunItfTrace(t *testing.T, path string) {
 
 	t.Log("Chains are: ", chains)
 
+	// generate keys that can be assigned on consumers, according to the ConsumerAddresses in the trace
+	consumerAddressesExpr := params["ConsumerAddresses"].Value.(itf.ListExprType)
+
+	_, _, consumerPrivVals, err := integration.CreateValidators(len(consumerAddressesExpr))
+	require.NoError(t, err, "Error creating consumer signers")
+
+	consumerAddrNamesToPrivVals := make(map[string]cmttypes.PrivValidator, len(consumerAddressesExpr))
+	realAddrsToModelConsAddrs := make(map[string]string, len(consumerAddressesExpr))
+	i := 0
+	for address, privVal := range consumerPrivVals {
+		consumerAddrNamesToPrivVals[consumerAddressesExpr[i].Value.(string)] = privVal
+		realAddrsToModelConsAddrs[address] = consumerAddressesExpr[i].Value.(string)
+		i++
+	}
+
 	// create params struct
 	vscTimeout := time.Duration(params["VscTimeout"].Value.(int64)) * time.Second
 
@@ -145,6 +164,15 @@ func RunItfTrace(t *testing.T, path string) {
 	valSet, addressMap, signers, err := CreateValSet(initialValSet)
 	require.NoError(t, err, "Error creating validator set")
 
+	// get the set of signers for consumers: the validator signers, plus signers for the assignable addresses
+	consumerSigners := make(map[string]cmttypes.PrivValidator, 0)
+	for consAddr, consPrivVal := range consumerPrivVals {
+		consumerSigners[consAddr] = consPrivVal
+	}
+	for consAddr, signer := range signers {
+		consumerSigners[consAddr] = signer
+	}
+
 	// get a slice of validators in the right order
 	nodes := make([]*cmttypes.Validator, len(valNames))
 	for i, valName := range valNames {
@@ -211,6 +239,10 @@ func RunItfTrace(t *testing.T, path string) {
 			// and then increment the rest of the time
 			runningConsumersBefore := driver.runningConsumers()
 			driver.endAndBeginBlock("provider", 1*time.Nanosecond)
+			for _, consumer := range driver.runningConsumers() {
+				UpdateProviderClientOnConsumer(t, driver, consumer.ChainId)
+			}
+
 			driver.endAndBeginBlock("provider", time.Duration(timeAdvancement)*time.Second-1*time.Nanosecond)
 			runningConsumersAfter := driver.runningConsumers()
 
@@ -243,7 +275,7 @@ func RunItfTrace(t *testing.T, path string) {
 					consumer.Value.(string),
 					modelParams,
 					driver.providerChain().Vals,
-					signers,
+					consumerSigners,
 					nodes,
 					valNames,
 					driver.providerChain(),
@@ -268,11 +300,8 @@ func RunItfTrace(t *testing.T, path string) {
 				if len(consumersToStart) > 0 && consumer.ChainId == consumersToStart[len(consumersToStart)-1].Value.(string) {
 					continue
 				}
-				consumerChainId := consumer.ChainId
 
-				driver.path(ChainId(consumerChainId)).AddClientHeader(PROVIDER, driver.providerHeader())
-				err := driver.path(ChainId(consumerChainId)).UpdateClient(consumerChainId, false)
-				require.True(t, err == nil, "Error updating client from %v on provider: %v", consumerChainId, err)
+				UpdateProviderClientOnConsumer(t, driver, consumer.ChainId)
 			}
 
 		case "EndAndBeginBlockForConsumer":
@@ -286,13 +315,12 @@ func RunItfTrace(t *testing.T, path string) {
 			_ = headerBefore
 
 			driver.endAndBeginBlock(ChainId(consumerChain), 1*time.Nanosecond)
+			UpdateConsumerClientOnProvider(t, driver, consumerChain)
+
 			driver.endAndBeginBlock(ChainId(consumerChain), time.Duration(timeAdvancement)*time.Second-1*time.Nanosecond)
 
 			// update the client on the provider
-			consumerHeader := driver.chain(ChainId(consumerChain)).LastHeader
-			driver.path(ChainId(consumerChain)).AddClientHeader(consumerChain, consumerHeader)
-			err := driver.path(ChainId(consumerChain)).UpdateClient(PROVIDER, false)
-			require.True(t, err == nil, "Error updating client from %v on provider: %v", consumerChain, err)
+			UpdateConsumerClientOnProvider(t, driver, consumerChain)
 
 		case "DeliverVscPacket":
 			consumerChain := lastAction["consumerChain"].Value.(string)
@@ -328,8 +356,26 @@ func RunItfTrace(t *testing.T, path string) {
 				expectError = false
 				driver.DeliverPacketFromConsumer(ChainId(consumerChain), expectError)
 			}
-		default:
+		case "KeyAssignment":
+			consumerChain := lastAction["consumerChain"].Value.(string)
+			node := lastAction["validator"].Value.(string)
+			consumerAddr := lastAction["consumerAddr"].Value.(string)
+
+			t.Log("KeyAssignment", consumerChain, node, consumerAddr)
+			stats.numKeyAssignments++
 
+			valIndex := getIndexOfString(node, valNames)
+			assignedPrivVal := consumerAddrNamesToPrivVals[consumerAddr]
+			assignedKey, err := assignedPrivVal.GetPubKey()
+			require.NoError(t, err, "Error getting pubkey")
+
+			protoPubKey, err := tmencoding.PubKeyToProto(assignedKey)
+			require.NoError(t, err, "Error converting pubkey to proto")
+
+			error := driver.AssignKey(ChainId(consumerChain), int64(valIndex), protoPubKey)
+			require.NoError(t, error, "Error assigning key")
+
+		default:
 			log.Fatalf("Error loading trace file %s, step %v: do not know action type %s",
 				path, index, actionKind)
 		}
@@ -364,7 +410,7 @@ func RunItfTrace(t *testing.T, path string) {
 		require.Equal(t, modelRunningConsumers, actualRunningConsumers, "Running consumers do not match")
 
 		// check validator sets - provider current validator set should be the one from the staking keeper
-		CompareValidatorSets(t, driver, currentModelState, actualRunningConsumers)
+		CompareValidatorSets(t, driver, currentModelState, actualRunningConsumers, realAddrsToModelConsAddrs)
 
 		// check times - sanity check that the block times match the ones from the model
 		CompareTimes(driver, actualRunningConsumers, currentModelState, timeOffset)
@@ -383,7 +429,27 @@ func RunItfTrace(t *testing.T, path string) {
 	t.Log("🟢 Trace is ok!")
 }
 
-func CompareValidatorSets(t *testing.T, driver *Driver, currentModelState map[string]itf.Expr, consumers []string) {
+func UpdateProviderClientOnConsumer(t *testing.T, driver *Driver, consumerChainId string) {
+	driver.path(ChainId(consumerChainId)).AddClientHeader(PROVIDER, driver.providerHeader())
+	err := driver.path(ChainId(consumerChainId)).UpdateClient(consumerChainId, false)
+	require.True(t, err == nil, "Error updating client from %v on provider: %v", consumerChainId, err)
+}
+
+func UpdateConsumerClientOnProvider(t *testing.T, driver *Driver, consumerChain string) {
+	consumerHeader := driver.chain(ChainId(consumerChain)).LastHeader
+	driver.path(ChainId(consumerChain)).AddClientHeader(consumerChain, consumerHeader)
+	err := driver.path(ChainId(consumerChain)).UpdateClient(PROVIDER, false)
+	require.True(t, err == nil, "Error updating client from %v on provider: %v", consumerChain, err)
+}
+
+func CompareValidatorSets(
+	t *testing.T,
+	driver *Driver,
+	currentModelState map[string]itf.Expr,
+	consumers []string,
+	// a map from real addresses to the names of those consumer addresses in the model
+	keyAddrsToModelConsAddrName map[string]string,
+) {
 	t.Helper()
 	modelValSet := ValidatorSet(currentModelState, "provider")
 
@@ -407,23 +473,28 @@ func CompareValidatorSets(t *testing.T, driver *Driver, currentModelState map[st
 			pubkey, err := val.ConsPubKey()
 			require.NoError(t, err, "Error getting pubkey")
 
-			consAddr := providertypes.NewConsumerConsAddress(sdktypes.ConsAddress(pubkey.Address().Bytes()))
+			consAddrModelName, ok := keyAddrsToModelConsAddrName[pubkey.Address().String()]
+			if ok { // the node has a key assigned, use the name of the consumer address in the model
+				consumerCurValSet[consAddrModelName] = val.Power
+			} else { // the node doesn't have a key assigned yet, get the validator moniker
+				consAddr := providertypes.NewConsumerConsAddress(sdktypes.ConsAddress(pubkey.Address().Bytes()))
 
-			// the consumer vals right now are CrossChainValidators, for which we don't know their mnemonic
-			// so we need to find the mnemonic of the consumer val now to enter it by name in the map
+				// the consumer vals right now are CrossChainValidators, for which we don't know their mnemonic
+				// so we need to find the mnemonic of the consumer val now to enter it by name in the map
 
-			// get the address on the provider that corresponds to the consumer address
-			providerConsAddr, found := driver.providerKeeper().GetValidatorByConsumerAddr(driver.providerCtx(), consumer, consAddr)
-			if !found {
-				providerConsAddr = providertypes.NewProviderConsAddress(consAddr.Address)
-			}
+				// get the address on the provider that corresponds to the consumer address
+				providerConsAddr, found := driver.providerKeeper().GetValidatorByConsumerAddr(driver.providerCtx(), consumer, consAddr)
+				if !found {
+					providerConsAddr = providertypes.NewProviderConsAddress(consAddr.Address)
+				}
 
-			// get the validator for that address on the provider
-			providerVal, found := driver.providerStakingKeeper().GetValidatorByConsAddr(driver.providerCtx(), providerConsAddr.Address)
-			require.True(t, found, "Error getting provider validator")
+				// get the validator for that address on the provider
+				providerVal, found := driver.providerStakingKeeper().GetValidatorByConsAddr(driver.providerCtx(), providerConsAddr.Address)
+				require.True(t, found, "Error getting provider validator")
 
-			// use the moniker of that validator
-			consumerCurValSet[providerVal.GetMoniker()] = val.Power
+				// use the moniker of that validator
+				consumerCurValSet[providerVal.GetMoniker()] = val.Power
+			}
 		}
 		require.NoError(t, CompareValSet(modelValSet, consumerCurValSet), "Validator sets do not match for consumer %v", consumer)
 	}

tests/mbt/driver/stats.go

@@ -16,4 +16,6 @@ type Stats struct {
 	numTxs         int
 
 	totalBlockTimePassedPerTrace time.Duration
+
+	numKeyAssignments int
 }

tests/mbt/model/README.md

@@ -31,16 +31,25 @@ All the logic in EndBlock/BeginBlock happens here, like updating the validator s
 * `EndAndBeginBlockForConsumer(chain: Chain, timeAdvancement: Time)`: On the consumer `chain`, ends the current block, and begins a new one. Again, all the logic in EndBlock/BeginBlock happens here, like validator set change maturations.
 * `DeliverVscPacket(receiver: Chain)`: Delivers a pending VSCPacket from the provider to the consumer `receiver`.
 * `DeliverVscMaturedPacket(receiver: Chain)`: Delivers a pending VSCMaturedPacket from the consumer `receiver` to the provider.
+* `KeyAssignment(chain: Chain, validator: Node, consumerAddr: ConsumerAddr)` (only when running with `--step stepKeyAssignment`): Assigns the `consumerAddr` to the `validator` on the `chain`. Note that we use "key" and "consumerAddr" pretty much interchangeably, as the model makes no differentiation between private keys, public keys, addresses, etc, as it doesn't model the cryptography.
 
 ### State machines
 
 There are 3 different "state machine layers" that can be put on top of the core logic.
+Some layers include extra logic, need other invariants, ...
 
 #### ccv_model.qnt
 This is the most general state machine layer. It allows the most behaviour,
 in particular it allows abitrary clock drift between chains, it allows starting and
 stopping consumer chains during runtime, etc.
 This layer is most useful for model checking, because it encompasses the most behaviour.
+As an optional module, it can also include KeyAssignment.
+
+##### KeyAssignment
+
+To run with key assignment, specify the step flag: `--step stepKeyAssignment`.
+
+KeyAssignment also needs some different invariants, see below.
 
 #### ccv_boundeddrift.qnt
 This state machine layer is more restricted to generate more interesting traces:
@@ -94,6 +103,13 @@ with a timestamp >= t + UnbondingPeriod on that consumer.
 - [X] EventuallyMatureOnProviderInv: If we send a VscPacket, this is eventually responded to by all consumers
 that were running at the time the packet was sent (and are still running).
 
+Invariants only relevant when running with key assignment (`--step stepKeyAssignment`):
+- [X] ValidatorSetHasExistedKeyAssignmentInv: Should replace ValidatorSetHasExistedInv when running with `--step stepKeyAssignment`. Validator sets are checked for equality under key assignment when checking whether they have existed.
+- [X] SameVscPacketsKeyAssignmentInv: Should replace SameVscPacketsInv when running with `--step stepKeyAssignment`. VscPackets are checked for equality under key assignment when ensuring consumers receive the same ones.
+- [X] KeyAssignmentRulesInv: Ensures the rules of key assignment are never violated. The two rules relevant for the model are: 1) validator A cannot assign consumer key K to consumer chain X if there is already a validator B (B!=A)
+using K on the provider, and 2) validator A cannot assign consumer key K to consumer chain X if there is already a validator B using K on X
+
+
 Invariants can also be model-checked by Apalache, using this command:
 ```
 quint verify --invariant ValidatorUpdatesArePropagatedInv,ValidatorSetHasExistedInv,SameVscPacketsInv,MatureOnTimeInv,EventuallyMatureOnProviderInv \
@@ -113,4 +129,6 @@ The available sanity checks are:
 - CanStopConsumer
 - CanTimeoutConsumer
 - CanSendVscPackets
-- CanSendVscMaturedPackets
+- CanSendVscMaturedPackets
+- CanAssignConsumerKey (only with `--step stepKeyAssignment`)
+- CanHaveConsumerAddresses (only with `--step stepKeyAssignment`)