test: add key assignment to model and driver

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.CreateSigners(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

@@ -35,12 +35,24 @@ All the logic in EndBlock/BeginBlock happens here, like updating the validator s
 ### 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`.
+
+Key assignment allows an additional action:
+
+* `KeyAssignment(chain: Chain, validator: Node, consumerAddr: ConsumerAddr)`: 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.
+
+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 +106,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: Replaces ValidatorSetHasExistedInv. Validator sets are checked for equality under key assignment when checking whether they have existed.
+- [X] SameVscPacketsKeyAssignmentInv: Replaces SameVscPacketsInv. 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 +132,6 @@ The available sanity checks are:
 - CanStopConsumer
 - CanTimeoutConsumer
 - CanSendVscPackets
-- CanSendVscMaturedPackets
+- CanSendVscMaturedPackets
+- CanAssignConsumerKey (only with `--step stepKeyAssignment`)
+- CanHaveConsumerAddresses (only with `--step stepKeyAssignment`)