feat: snapshot Solo environment

packages/cryptolib/keypair.go

@@ -58,7 +58,9 @@ func (k *KeyPair) Address() *iotago.Ed25519Address {
 
 func (k *KeyPair) Read(r io.Reader) error {
 	rr := rwutil.NewReader(r)
+	k.publicKey = new(PublicKey)
 	rr.Read(k.publicKey)
+	k.privateKey = new(PrivateKey)
 	rr.Read(k.privateKey)
 	return rr.Err
 }

packages/cryptolib/private_key.go

@@ -83,9 +83,7 @@ func (pkT *PrivateKey) AddressKeys(addr iotago.Address) iotago.AddressKeys {
 
 func (pkT *PrivateKey) Read(r io.Reader) error {
 	rr := rwutil.NewReader(r)
-	if len(pkT.key) != PrivateKeySize {
-		panic("unexpected private key size for read")
-	}
+	pkT.key = make([]byte, PrivateKeySize)
 	rr.ReadN(pkT.key)
 	return rr.Err
 }

packages/cryptolib/public_key.go

@@ -93,9 +93,7 @@ func (pkT *PublicKey) String() string {
 
 func (pkT *PublicKey) Read(r io.Reader) error {
 	rr := rwutil.NewReader(r)
-	if len(pkT.key) != PublicKeySize {
-		panic("unexpected public key size for read")
-	}
+	pkT.key = make([]byte, PublicKeySize)
 	rr.ReadN(pkT.key)
 	return rr.Err
 }

packages/solo/snapshot.go

@@ -0,0 +1,105 @@
+package solo
+
+import (
+	"encoding/json"
+	"os"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/iotaledger/hive.go/kvstore"
+	"github.com/iotaledger/wasp/packages/cryptolib"
+	"github.com/iotaledger/wasp/packages/isc"
+	"github.com/iotaledger/wasp/packages/testutil/utxodb"
+	"github.com/iotaledger/wasp/packages/util/rwutil"
+)
+
+type soloSnapshot struct {
+	UtxoDB *utxodb.UtxoDBState
+	Chains []soloChainSnapshot
+}
+
+type soloChainSnapshot struct {
+	Name                   string
+	StateControllerKeyPair []byte
+	ChainID                []byte
+	OriginatorPrivateKey   []byte
+	ValidatorFeeTarget     []byte
+	DB                     [][]byte
+}
+
+// SaveSnapshot generates a snapshot of the Solo environment
+func (env *Solo) SaveSnapshot(fname string) {
+	env.glbMutex.Lock()
+	defer env.glbMutex.Unlock()
+
+	snapshot := soloSnapshot{
+		UtxoDB: env.utxoDB.State(),
+	}
+
+	for _, ch := range env.chains {
+		chainSnapshot := soloChainSnapshot{
+			Name:                   ch.Name,
+			StateControllerKeyPair: rwutil.WriteToBytes(ch.StateControllerKeyPair),
+			ChainID:                ch.ChainID.Bytes(),
+			OriginatorPrivateKey:   rwutil.WriteToBytes(ch.OriginatorPrivateKey),
+			ValidatorFeeTarget:     ch.ValidatorFeeTarget.Bytes(),
+		}
+
+		err := ch.db.Iterate(kvstore.EmptyPrefix, func(k, v []byte) bool {
+			chainSnapshot.DB = append(chainSnapshot.DB, k, v)
+			return true
+		})
+		require.NoError(env.T, err)
+
+		snapshot.Chains = append(snapshot.Chains, chainSnapshot)
+	}
+
+	b, err := json.Marshal(snapshot)
+	require.NoError(env.T, err)
+	err = os.WriteFile(fname, b, 0o600)
+	require.NoError(env.T, err)
+}
+
+// LoadSnapshot restores the Solo environment from the given snapshot
+func (env *Solo) LoadSnapshot(fname string) {
+	env.glbMutex.Lock()
+	defer env.glbMutex.Unlock()
+
+	b, err := os.ReadFile(fname)
+	require.NoError(env.T, err)
+	var snapshot soloSnapshot
+	err = json.Unmarshal(b, &snapshot)
+	require.NoError(env.T, err)
+
+	env.utxoDB.SetState(snapshot.UtxoDB)
+	for _, chainSnapshot := range snapshot.Chains {
+		sckp, err := rwutil.ReadFromBytes(chainSnapshot.StateControllerKeyPair, new(cryptolib.KeyPair))
+		require.NoError(env.T, err)
+
+		chainID, err := isc.ChainIDFromBytes(chainSnapshot.ChainID)
+		require.NoError(env.T, err)
+
+		okp, err := rwutil.ReadFromBytes(chainSnapshot.OriginatorPrivateKey, new(cryptolib.KeyPair))
+		require.NoError(env.T, err)
+
+		val, err := isc.AgentIDFromBytes(chainSnapshot.ValidatorFeeTarget)
+		require.NoError(env.T, err)
+
+		db, err := env.chainStateDatabaseManager.ChainStateKVStore(chainID)
+		require.NoError(env.T, err)
+		for i := 0; i < len(chainSnapshot.DB); i += 2 {
+			err = db.Set(chainSnapshot.DB[i], chainSnapshot.DB[i+1])
+			require.NoError(env.T, err)
+		}
+
+		chainData := chainData{
+			Name:                   chainSnapshot.Name,
+			StateControllerKeyPair: sckp,
+			ChainID:                chainID,
+			OriginatorPrivateKey:   okp,
+			ValidatorFeeTarget:     val,
+			db:                     db,
+		}
+		env.addChain(chainData)
+	}
+}

packages/solo/solo.go

@@ -15,6 +15,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap/zapcore"
 
+	"github.com/iotaledger/hive.go/kvstore"
 	hivedb "github.com/iotaledger/hive.go/kvstore/database"
 	"github.com/iotaledger/hive.go/logger"
 	iotago "github.com/iotaledger/iota.go/v3"
@@ -67,33 +68,40 @@ type Solo struct {
 	ctx                             context.Context
 }
 
-// Chain represents state of individual chain.
-// There may be several parallel instances of the chain in the 'solo' test
-type Chain struct {
-	// Env is a pointer to the global structure of the 'solo' test
-	Env *Solo
-
+// data to be persisted in the snapshot
+type chainData struct {
 	// Name is the name of the chain
 	Name string
 
 	// StateControllerKeyPair signature scheme of the chain address, the one used to control funds owned by the chain.
 	// In Solo it is Ed25519 signature scheme (in full Wasp environment is is a BLS address)
 	StateControllerKeyPair *cryptolib.KeyPair
-	StateControllerAddress iotago.Address
 
 	// ChainID is the ID of the chain (in this version alias of the ChainAddress)
 	ChainID isc.ChainID
 
 	// OriginatorPrivateKey the key pair used to create the chain (origin transaction).
 	// It is a default key pair in many of Solo calls which require private key.
 	OriginatorPrivateKey *cryptolib.KeyPair
-	OriginatorAddress    iotago.Address
-	// OriginatorAgentID is the OriginatorAddress represented in the form of AgentID
-	OriginatorAgentID isc.AgentID
 
 	// ValidatorFeeTarget is the agent ID to which all fees are accrued. By default, it is equal to OriginatorAgentID
 	ValidatorFeeTarget isc.AgentID
 
+	db kvstore.KVStore
+}
+
+// Chain represents state of individual chain.
+// There may be several parallel instances of the chain in the 'solo' test
+type Chain struct {
+	chainData
+
+	StateControllerAddress iotago.Address
+	OriginatorAddress      iotago.Address
+	OriginatorAgentID      isc.AgentID
+
+	// Env is a pointer to the global structure of the 'solo' test
+	Env *Solo
+
 	// Store is where the chain data (blocks, state) is stored
 	store indexedstore.IndexedStore
 	// Log is the named logger of the chain
@@ -199,6 +207,17 @@ func (env *Solo) Publisher() *publisher.Publisher {
 	return env.publisher
 }
 
+func (env *Solo) GetChainByName(name string) *Chain {
+	env.glbMutex.Lock()
+	defer env.glbMutex.Unlock()
+	for _, ch := range env.chains {
+		if ch.Name == name {
+			return ch
+		}
+	}
+	panic("chain not found")
+}
+
 // WithNativeContract registers a native contract so that it may be deployed
 func (env *Solo) WithNativeContract(c *coreutil.ContractProcessor) *Solo {
 	env.processorConfig.RegisterNativeContract(c)
@@ -216,26 +235,12 @@ func (env *Solo) NewChain(depositFundsForOriginator ...bool) *Chain {
 	return ret
 }
 
-// NewChainExt returns also origin and init transactions. Used for core testing
-//
-// If 'chainOriginator' is nil, new one is generated and utxodb.FundsFromFaucetAmount (many) base tokens are loaded from the UTXODB faucet.
-// ValidatorFeeTarget will be set to OriginatorAgentID, and can be changed after initialization.
-// To deploy a chain instance the following steps are performed:
-//   - chain signature scheme (private key), chain address and chain ID are created
-//   - empty virtual state is initialized
-//   - origin transaction is created by the originator and added to the UTXODB
-//   - 'init' request transaction to the 'root' contract is created and added to UTXODB
-//   - backlog processing threads (goroutines) are started
-//   - VM processor cache is initialized
-//   - 'init' request is run by the VM. The 'root' contracts deploys the rest of the core contracts:
-//
-// Upon return, the chain is fully functional to process requests
-func (env *Solo) NewChainExt(
+func (env *Solo) deployChain(
 	chainOriginator *cryptolib.KeyPair,
 	initBaseTokens uint64,
 	name string,
 	originParams ...dict.Dict,
-) (*Chain, *iotago.Transaction) {
+) (chainData, *iotago.Transaction) {
 	env.logger.Debugf("deploying new chain '%s'", name)
 
 	if chainOriginator == nil {
@@ -288,46 +293,74 @@ func (env *Solo) NewChainExt(
 	env.logger.Infof("     chain '%s'. state controller address: %s", chainID.String(), stateControllerAddr.Bech32(parameters.L1().Protocol.Bech32HRP))
 	env.logger.Infof("     chain '%s'. originator address: %s", chainID.String(), originatorAddr.Bech32(parameters.L1().Protocol.Bech32HRP))
 
-	chainlog := env.logger.Named(name)
-
-	kvStore, err := env.chainStateDatabaseManager.ChainStateKVStore(chainID)
+	db, err := env.chainStateDatabaseManager.ChainStateKVStore(chainID)
 	require.NoError(env.T, err)
 	originAOMinSD := parameters.L1().Protocol.RentStructure.MinRent(originAO)
-	store := indexedstore.New(state.NewStore(kvStore))
+	store := indexedstore.New(state.NewStore(db))
 	origin.InitChain(store, initParams, originAO.Amount-originAOMinSD)
 
 	{
 		block, err2 := store.LatestBlock()
 		require.NoError(env.T, err2)
-		env.logger.Infof("     chain '%s'. origin trie root: %s", chainID.String(), block.TrieRoot())
+		env.logger.Infof("     chain '%s'. origin trie root: %s", chainID, block.TrieRoot())
 	}
 
-	ret := &Chain{
-		Env:                    env,
+	return chainData{
 		Name:                   name,
 		ChainID:                chainID,
 		StateControllerKeyPair: stateControllerKey,
-		StateControllerAddress: stateControllerAddr,
 		OriginatorPrivateKey:   chainOriginator,
-		OriginatorAddress:      originatorAddr,
-		OriginatorAgentID:      originatorAgentID,
 		ValidatorFeeTarget:     originatorAgentID,
-		store:                  store,
-		proc:                   processors.MustNew(env.processorConfig),
-		log:                    chainlog,
-		metrics:                metrics.NewChainMetricsProvider().GetChainMetrics(chainID),
-	}
+		db:                     db,
+	}, originTx
+}
 
-	ret.mempool = newMempool(env.utxoDB.GlobalTime)
+// NewChainExt returns also origin and init transactions. Used for core testing
+//
+// If 'chainOriginator' is nil, new one is generated and utxodb.FundsFromFaucetAmount (many) base tokens are loaded from the UTXODB faucet.
+// ValidatorFeeTarget will be set to OriginatorAgentID, and can be changed after initialization.
+// To deploy a chain instance the following steps are performed:
+//   - chain signature scheme (private key), chain address and chain ID are created
+//   - empty virtual state is initialized
+//   - origin transaction is created by the originator and added to the UTXODB
+//   - 'init' request transaction to the 'root' contract is created and added to UTXODB
+//   - backlog processing threads (goroutines) are started
+//   - VM processor cache is initialized
+//   - 'init' request is run by the VM. The 'root' contracts deploys the rest of the core contracts:
+//
+// Upon return, the chain is fully functional to process requests
+func (env *Solo) NewChainExt(
+	chainOriginator *cryptolib.KeyPair,
+	initBaseTokens uint64,
+	name string,
+	originParams ...dict.Dict,
+) (*Chain, *iotago.Transaction) {
+	chData, originTx := env.deployChain(chainOriginator, initBaseTokens, name, originParams...)
 
 	env.glbMutex.Lock()
-	env.chains[chainID] = ret
-	env.glbMutex.Unlock()
+	defer env.glbMutex.Unlock()
+	ch := env.addChain(chData)
 
-	go ret.batchLoop()
+	ch.log.Infof("chain '%s' deployed. Chain ID: %s", ch.Name, ch.ChainID.String())
+	return ch, originTx
+}
 
-	ret.log.Infof("chain '%s' deployed. Chain ID: %s", ret.Name, ret.ChainID.String())
-	return ret, originTx
+func (env *Solo) addChain(chData chainData) *Chain {
+	ch := &Chain{
+		chainData:              chData,
+		StateControllerAddress: chData.StateControllerKeyPair.GetPublicKey().AsEd25519Address(),
+		OriginatorAddress:      chData.OriginatorPrivateKey.GetPublicKey().AsEd25519Address(),
+		OriginatorAgentID:      isc.NewAgentID(chData.OriginatorPrivateKey.GetPublicKey().AsEd25519Address()),
+		Env:                    env,
+		store:                  indexedstore.New(state.NewStore(chData.db)),
+		proc:                   processors.MustNew(env.processorConfig),
+		log:                    env.logger.Named(chData.Name),
+		metrics:                metrics.NewChainMetricsProvider().GetChainMetrics(chData.ChainID),
+		mempool:                newMempool(env.utxoDB.GlobalTime),
+	}
+	env.chains[chData.ChainID] = ch
+	go ch.batchLoop()
+	return ch
 }
 
 // AddToLedger adds (synchronously confirms) transaction to the UTXODB ledger. Return error if it is