address PR comments

arbnode/dataposter/data_poster.go

@@ -221,7 +221,7 @@ func rpcClient(ctx context.Context, opts *ExternalSignerCfg) (*rpc.Client, error
 		// that it expects to be signed with. So signer is already authenticated
 		// on application level and does not need to rely on TLS for authentication.
 		InsecureSkipVerify: opts.InsecureSkipVerify, // #nosec G402
-	}
+	} // #nosec G402
 
 	if opts.ClientCert != "" && opts.ClientPrivateKey != "" {
 		log.Info("Client certificate for external signer is enabled")

arbos/programs/programs.go

@@ -248,9 +248,10 @@ func (p Programs) CallProgram(
 
 func getWasm(statedb vm.StateDB, program common.Address) ([]byte, error) {
 	prefixedWasm := statedb.GetCode(program)
-	return getWasmInternal(prefixedWasm)
+	return getWasmFromContractCode(prefixedWasm)
 }
-func getWasmInternal(prefixedWasm []byte) ([]byte, error) {
+
+func getWasmFromContractCode(prefixedWasm []byte) ([]byte, error) {
 	if prefixedWasm == nil {
 		return nil, ProgramNotWasmError()
 	}
@@ -321,9 +322,9 @@ func (p Programs) SaveActiveProgramToWasmStore(statedb *state.StateDB, codeHash
 		return nil
 	}
 
-	wasm, err := getWasmInternal(code)
+	wasm, err := getWasmFromContractCode(code)
 	if err != nil {
-		log.Error("Failed to reactivate program while rebuilding wasm store: getWasmInternal", "expected moduleHash", moduleHash, "err", err)
+		log.Error("Failed to reactivate program while rebuilding wasm store: getWasmFromContractCode", "expected moduleHash", moduleHash, "err", err)
 		return fmt.Errorf("failed to reactivate program while rebuilding wasm store: %w", err)
 	}
 

cmd/nitro/init.go

@@ -38,7 +38,6 @@ import (
 	"github.com/offchainlabs/nitro/execution/gethexec"
 	"github.com/offchainlabs/nitro/statetransfer"
 	"github.com/offchainlabs/nitro/util/arbmath"
-	"github.com/offchainlabs/nitro/util/wasmstorerebuilder"
 )
 
 func downloadInit(ctx context.Context, initConfig *conf.InitConfig) (string, error) {
@@ -210,28 +209,28 @@ func openInitializeChainDb(ctx context.Context, stack *node.Node, config *NodeCo
 				if latestBlock.Number.Uint64() <= chainConfig.ArbitrumChainParams.GenesisBlockNum {
 					// If there is only genesis block or no blocks in the blockchain, set Rebuilding of wasmdb to Done
 					log.Info("setting rebuilding of wasmdb to done")
-					if err = wasmstorerebuilder.SetRebuildingParam(wasmDb, wasmstorerebuilder.RebuildingPositionKey, wasmstorerebuilder.RebuildingDone); err != nil {
+					if err = gethexec.WriteToKeyValueStore(wasmDb, gethexec.RebuildingPositionKey, gethexec.RebuildingDone); err != nil {
 						return nil, nil, fmt.Errorf("unable to set rebuilding status of wasmdb to done: %w", err)
 					}
 				} else {
-					key, err := wasmstorerebuilder.GetRebuildingParam[common.Hash](wasmDb, wasmstorerebuilder.RebuildingPositionKey)
+					position, err := gethexec.ReadFromKeyValueStore[common.Hash](wasmDb, gethexec.RebuildingPositionKey)
 					if err != nil {
 						log.Info("unable to get codehash position in rebuilding of wasmdb, its possible it isnt initialized yet, so initializing it and starting rebuilding", "err", err)
-						if err := wasmstorerebuilder.SetRebuildingParam(wasmDb, wasmstorerebuilder.RebuildingPositionKey, common.Hash{}); err != nil {
+						if err := gethexec.WriteToKeyValueStore(wasmDb, gethexec.RebuildingPositionKey, common.Hash{}); err != nil {
 							return nil, nil, fmt.Errorf("unable to set rebuilding status of wasmdb to beginning: %w", err)
 						}
 					}
-					startBlockTime, err := wasmstorerebuilder.GetRebuildingParam[uint64](wasmDb, wasmstorerebuilder.RebuildingStartBlockTimeKey)
-					if err != nil {
-						log.Info("unable to get rebuilding start time of wasmdb so initializing it to current block time", "err", err)
-						if err := wasmstorerebuilder.SetRebuildingParam(wasmDb, wasmstorerebuilder.RebuildingStartBlockTimeKey, latestBlock.Time); err != nil {
-							return nil, nil, fmt.Errorf("unable to set rebuilding status of wasmdb to beginning: %w", err)
+					if position != gethexec.RebuildingDone {
+						startBlockHash, err := gethexec.ReadFromKeyValueStore[common.Hash](wasmDb, gethexec.RebuildingStartBlockHashKey)
+						if err != nil {
+							log.Info("unable to get rebuilding start time of wasmdb so initializing it to current block time", "err", err)
+							if err := gethexec.WriteToKeyValueStore(wasmDb, gethexec.RebuildingStartBlockHashKey, latestBlock.Hash()); err != nil {
+								return nil, nil, fmt.Errorf("unable to set rebuilding status of wasmdb to beginning: %w", err)
+							}
+							startBlockHash = latestBlock.Hash()
 						}
-						startBlockTime = latestBlock.Time
-					}
-					if key != wasmstorerebuilder.RebuildingDone {
-						log.Info("starting or continuing rebuilding of wasm store", "codeHash", key, "startBlockTime", startBlockTime)
-						go wasmstorerebuilder.RebuildWasmStore(ctx, wasmDb, l2BlockChain, key, startBlockTime)
+						log.Info("starting or continuing rebuilding of wasm store", "codeHash", position, "startBlockHash", startBlockHash)
+						gethexec.RebuildWasmStore(ctx, wasmDb, l2BlockChain, position, startBlockHash)
 					}
 				}
 				return chainDb, l2BlockChain, nil
@@ -274,7 +273,7 @@ func openInitializeChainDb(ctx context.Context, stack *node.Node, config *NodeCo
 	chainDb := rawdb.WrapDatabaseWithWasm(chainData, wasmDb, 1)
 
 	// Rebuilding wasmdb is not required when just starting out
-	err = wasmstorerebuilder.SetRebuildingParam(wasmDb, wasmstorerebuilder.RebuildingPositionKey, wasmstorerebuilder.RebuildingDone)
+	err = gethexec.WriteToKeyValueStore(wasmDb, gethexec.RebuildingPositionKey, gethexec.RebuildingDone)
 	log.Info("setting rebuilding of wasmdb to done")
 	if err != nil {
 		return nil, nil, fmt.Errorf("unable to set rebuilding of wasmdb to done: %w", err)

util/wasmstorerebuilder/wasmstorerebuilder.go → execution/gethexec/wasmstorerebuilder.go

@@ -1,7 +1,7 @@
 // Copyright 2021-2024, Offchain Labs, Inc.
 // For license information, see https://github.com/nitro/blob/master/LICENSE
 
-package wasmstorerebuilder
+package gethexec
 
 import (
 	"bytes"
@@ -19,29 +19,29 @@ import (
 )
 
 var RebuildingPositionKey []byte = []byte("_rebuildingPosition")             // contains the codehash upto where rebuilding of wasm store was last completed
-var RebuildingStartBlockTimeKey []byte = []byte("_rebuildingStartBlockTime") // contains the block time when rebuilding of wasm store first began
+var RebuildingStartBlockHashKey []byte = []byte("_rebuildingStartBlockHash") // contains the block hash of starting block when rebuilding of wasm store first began
 var RebuildingDone common.Hash = common.BytesToHash([]byte("_done"))         // indicates that the rebuilding is done, if RebuildingPositionKey holds this value it implies rebuilding was completed
 
-func GetRebuildingParam[T any](wasmStore ethdb.KeyValueStore, key []byte) (T, error) {
+func ReadFromKeyValueStore[T any](store ethdb.KeyValueStore, key []byte) (T, error) {
 	var empty T
-	posBytes, err := wasmStore.Get(key)
+	posBytes, err := store.Get(key)
 	if err != nil {
 		return empty, err
 	}
 	var val T
 	err = rlp.DecodeBytes(posBytes, &val)
 	if err != nil {
-		return empty, fmt.Errorf("invalid value stored in rebuilding key or error decoding rebuildingBytes: %w", err)
+		return empty, fmt.Errorf("error decoding value stored for key in the KeyValueStore: %w", err)
 	}
 	return val, nil
 }
 
-func SetRebuildingParam[T any](wasmStore ethdb.KeyValueStore, key []byte, val T) error {
+func WriteToKeyValueStore[T any](store ethdb.KeyValueStore, key []byte, val T) error {
 	valBytes, err := rlp.EncodeToBytes(val)
 	if err != nil {
 		return err
 	}
-	err = wasmStore.Put(key, valBytes)
+	err = store.Put(key, valBytes)
 	if err != nil {
 		return err
 	}
@@ -54,51 +54,60 @@ func SetRebuildingParam[T any](wasmStore ethdb.KeyValueStore, key []byte, val T)
 // It stores the status of rebuilding to wasm store by updating the codehash (of the latest sucessfully checked contract) in
 // RebuildingPositionKey every 50 checks.
 //
-// It also stores a special value that is only set once when rebuilding commenced in RebuildingStartBlockTimeKey as the block
+// It also stores a special value that is only set once when rebuilding commenced in RebuildingStartBlockHashKey as the block
 // time of the latest block when rebuilding was first called, this is used to avoid recomputing of assembly and module of
 // contracts that were created after rebuilding commenced since they would anyway already be added during sync.
-func RebuildWasmStore(ctx context.Context, wasmStore ethdb.KeyValueStore, l2Blockchain *core.BlockChain, key common.Hash, rebuildingStartBlockTime uint64) {
+func RebuildWasmStore(ctx context.Context, wasmStore ethdb.KeyValueStore, l2Blockchain *core.BlockChain, position, rebuildingStartBlockHash common.Hash) {
+	var err error
+	var stateDb *state.StateDB
 	latestHeader := l2Blockchain.CurrentBlock()
-	latestState, err := l2Blockchain.StateAt(latestHeader.Root)
+	// Attempt to get state at the start block when rebuilding commenced, if not available (in case of non-archival nodes) use latest state
+	rebuildingStartHeader := l2Blockchain.GetHeaderByHash(rebuildingStartBlockHash)
+	stateDb, err = l2Blockchain.StateAt(rebuildingStartHeader.Root)
 	if err != nil {
-		log.Error("error getting state at latest block, aborting rebuilding", "err", err)
-		return
+		log.Info("error getting state at start block of rebuilding wasm store, attempting rebuilding with latest state", "err", err)
+		stateDb, err = l2Blockchain.StateAt(latestHeader.Root)
+		if err != nil {
+			log.Error("error getting state at latest block, aborting rebuilding", "err", err)
+			return
+		}
 	}
-	diskDb := latestState.Database().DiskDB()
-	arbState, err := arbosState.OpenSystemArbosState(latestState, nil, true)
+	diskDb := stateDb.Database().DiskDB()
+	arbState, err := arbosState.OpenSystemArbosState(stateDb, nil, true)
 	if err != nil {
 		log.Error("error getting arbos state, aborting rebuilding", "err", err)
 		return
 	}
 	programs := arbState.Programs()
-	iter := diskDb.NewIterator(rawdb.CodePrefix, key[:])
+	iter := diskDb.NewIterator(rawdb.CodePrefix, position[:])
 	for count := 1; iter.Next(); count++ {
-		// If outer context is cancelled we should terminate rebuilding. We probably wont be able to save codehash to wasm store (it might be closed)
-		if ctx.Err() != nil {
-			return
-		}
 		codeHashBytes := bytes.TrimPrefix(iter.Key(), rawdb.CodePrefix)
 		codeHash := common.BytesToHash(codeHashBytes)
 		code := iter.Value()
 		if state.IsStylusProgram(code) {
-			if err := programs.SaveActiveProgramToWasmStore(latestState, codeHash, code, latestHeader.Time, l2Blockchain.Config().DebugMode(), rebuildingStartBlockTime); err != nil {
+			if err := programs.SaveActiveProgramToWasmStore(stateDb, codeHash, code, latestHeader.Time, l2Blockchain.Config().DebugMode(), rebuildingStartHeader.Time); err != nil {
 				log.Error("error while rebuilding wasm store, aborting rebuilding", "err", err)
 				return
 			}
 		}
 		// After every fifty codeHash checks, update the rebuilding position
 		// This also notifies user that we are working on rebuilding
-		if count%50 == 0 {
+		if count%50 == 0 || ctx.Err() != nil {
 			log.Info("Storing rebuilding status to disk", "codeHash", codeHash)
-			if err := SetRebuildingParam(wasmStore, RebuildingPositionKey, codeHash); err != nil {
+			if err := WriteToKeyValueStore(wasmStore, RebuildingPositionKey, codeHash); err != nil {
 				log.Error("error updating position to wasm store mid way though rebuilding, aborting", "err", err)
 				return
 			}
+			// If outer context is cancelled we should terminate rebuilding
+			// We attempted to write the latest checked codeHash to wasm store
+			if ctx.Err() != nil {
+				return
+			}
 		}
 	}
 	iter.Release()
 	// Set rebuilding position to done indicating completion
-	if err := SetRebuildingParam(wasmStore, RebuildingPositionKey, RebuildingDone); err != nil {
+	if err := WriteToKeyValueStore(wasmStore, RebuildingPositionKey, RebuildingDone); err != nil {
 		log.Error("error updating rebuilding position to done", "err", err)
 		return
 	}

system_tests/program_test.go

@@ -31,12 +31,12 @@ import (
 	"github.com/offchainlabs/nitro/arbos/programs"
 	"github.com/offchainlabs/nitro/arbos/util"
 	"github.com/offchainlabs/nitro/arbutil"
+	"github.com/offchainlabs/nitro/execution/gethexec"
 	"github.com/offchainlabs/nitro/solgen/go/mocksgen"
 	pgen "github.com/offchainlabs/nitro/solgen/go/precompilesgen"
 	"github.com/offchainlabs/nitro/util/arbmath"
 	"github.com/offchainlabs/nitro/util/colors"
 	"github.com/offchainlabs/nitro/util/testhelpers"
-	"github.com/offchainlabs/nitro/util/wasmstorerebuilder"
 	"github.com/offchainlabs/nitro/validator/valnode"
 	"github.com/wasmerio/wasmer-go/wasmer"
 )
@@ -1633,19 +1633,19 @@ func TestWasmStoreRebuilding(t *testing.T) {
 	}
 
 	// Start rebuilding and wait for it to finish
-	log.Info("starting rebuilding wasm store")
-	wasmstorerebuilder.RebuildWasmStore(ctx, wasmDbAfterDelete, bc, common.Hash{}, bc.CurrentBlock().Time)
+	log.Info("starting rebuilding of wasm store")
+	gethexec.RebuildWasmStore(ctx, wasmDbAfterDelete, bc, common.Hash{}, bc.CurrentBlock().Hash())
 
 	wasmDbAfterRebuild := getLatestStateWasmStore(bc)
 
 	// Before comparing, check if rebuilding was set to done and then delete the keys that are used to track rebuilding status
-	status, err := wasmstorerebuilder.GetRebuildingParam[common.Hash](wasmDbAfterRebuild, wasmstorerebuilder.RebuildingPositionKey)
+	status, err := gethexec.ReadFromKeyValueStore[common.Hash](wasmDbAfterRebuild, gethexec.RebuildingPositionKey)
 	Require(t, err)
-	if status != wasmstorerebuilder.RebuildingDone {
+	if status != gethexec.RebuildingDone {
 		Fatal(t, "rebuilding was not set to done after successful completion")
 	}
-	Require(t, wasmDbAfterRebuild.Delete(wasmstorerebuilder.RebuildingPositionKey))
-	Require(t, wasmDbAfterRebuild.Delete(wasmstorerebuilder.RebuildingStartBlockTimeKey))
+	Require(t, wasmDbAfterRebuild.Delete(gethexec.RebuildingPositionKey))
+	Require(t, wasmDbAfterRebuild.Delete(gethexec.RebuildingStartBlockHashKey))
 
 	rebuiltStoreMap, err := createMapFromDb(wasmDbAfterRebuild)
 	Require(t, err)