collect and emit codehash=>code mappings for state objects

cerc-io · Oct 13, 2020 · 9b1e316 · 9b1e316
1 parent 7c0492a
commit 9b1e316
Show file tree

Hide file tree

Showing 4 changed files with 147 additions and 53 deletions.
diff --git a/statediff/builder.go b/statediff/builder.go
@@ -36,6 +36,7 @@ var (
 	nullHashBytes     = common.Hex2Bytes("0000000000000000000000000000000000000000000000000000000000000000")
 	emptyNode, _      = rlp.EncodeToBytes([]byte{})
 	emptyContractRoot = crypto.Keccak256Hash(emptyNode)
+	nullCodeHash      = crypto.Keccak256Hash([]byte{}).Bytes()
 )
 
 // Builder interface exposes the method for building a state diff between two blocks
@@ -62,19 +63,21 @@ func (sdb *builder) BuildStateTrieObject(current *types.Block) (StateObject, err
 		return StateObject{}, fmt.Errorf("error creating trie for block %d: %v", current.Number(), err)
 	}
 	it := currentTrie.NodeIterator([]byte{})
-	stateNodes, err := sdb.buildStateTrie(it)
+	stateNodes, codeAndCodeHashes, err := sdb.buildStateTrie(it)
 	if err != nil {
 		return StateObject{}, fmt.Errorf("error collecting state nodes for block %d: %v", current.Number(), err)
 	}
 	return StateObject{
-		BlockNumber: current.Number(),
-		BlockHash:   current.Hash(),
-		Nodes:       stateNodes,
+		BlockNumber:       current.Number(),
+		BlockHash:         current.Hash(),
+		Nodes:             stateNodes,
+		CodeAndCodeHashes: codeAndCodeHashes,
 	}, nil
 }
 
-func (sdb *builder) buildStateTrie(it trie.NodeIterator) ([]StateNode, error) {
+func (sdb *builder) buildStateTrie(it trie.NodeIterator) ([]StateNode, []CodeAndCodeHash, error) {
 	stateNodes := make([]StateNode, 0)
+	codeAndCodeHashes := make([]CodeAndCodeHash, 0)
 	for it.Next(true) {
 		// skip value nodes
 		if it.Leaf() {
@@ -87,48 +90,62 @@ func (sdb *builder) buildStateTrie(it trie.NodeIterator) ([]StateNode, error) {
 		copy(nodePath, it.Path())
 		node, err := sdb.stateCache.TrieDB().Node(it.Hash())
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		var nodeElements []interface{}
 		if err := rlp.DecodeBytes(node, &nodeElements); err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		ty, err := CheckKeyType(nodeElements)
 		if err != nil {
-			return nil, err
+			return nil, nil, err
 		}
 		switch ty {
 		case Leaf:
 			var account state.Account
 			if err := rlp.DecodeBytes(nodeElements[1].([]byte), &account); err != nil {
-				return nil, fmt.Errorf("error decoding account for leaf node at path %x nerror: %v", nodePath, err)
+				return nil, nil, fmt.Errorf("error decoding account for leaf node at path %x nerror: %v", nodePath, err)
 			}
 			partialPath := trie.CompactToHex(nodeElements[0].([]byte))
 			valueNodePath := append(nodePath, partialPath...)
 			encodedPath := trie.HexToCompact(valueNodePath)
 			leafKey := encodedPath[1:]
-			storageNodes, err := sdb.buildStorageNodesEventual(account.Root, nil, true)
-			if err != nil {
-				return nil, fmt.Errorf("failed building eventual storage diffs for account %+v\r\nerror: %v", account, err)
+			node := StateNode{
+				NodeType:  ty,
+				Path:      nodePath,
+				LeafKey:   leafKey,
+				NodeValue: node,
 			}
-			stateNodes = append(stateNodes, StateNode{
-				NodeType:     ty,
-				Path:         nodePath,
-				LeafKey:      leafKey,
-				NodeValue:    node,
-				StorageNodes: storageNodes,
-			})
+			if !bytes.Equal(account.CodeHash, nullCodeHash) {
+				storageNodes, err := sdb.buildStorageNodesEventual(account.Root, nil, true)
+				if err != nil {
+					return nil, nil, fmt.Errorf("failed building eventual storage diffs for account %+v\r\nerror: %v", account, err)
+				}
+				node.StorageNodes = storageNodes
+				// emit codehash => code mappings for cod
+				codeHash := common.BytesToHash(account.CodeHash)
+				addrHash := common.BytesToHash(leafKey)
+				code, err := sdb.stateCache.ContractCode(addrHash, codeHash)
+				if err != nil {
+					return nil, nil, fmt.Errorf("failed to retrieve code for codehash %s for account with leafkey %s\r\n error: %v", codeHash.String(), addrHash.String(), err)
+				}
+				codeAndCodeHashes = append(codeAndCodeHashes, CodeAndCodeHash{
+					Hash: codeHash,
+					Code: code,
+				})
+			}
+			stateNodes = append(stateNodes, node)
 		case Extension, Branch:
 			stateNodes = append(stateNodes, StateNode{
 				NodeType:  ty,
 				Path:      nodePath,
 				NodeValue: node,
 			})
 		default:
-			return nil, fmt.Errorf("unexpected node type %s", ty)
+			return nil, nil, fmt.Errorf("unexpected node type %s", ty)
 		}
 	}
-	return stateNodes, it.Error()
+	return stateNodes, codeAndCodeHashes, it.Error()
 }
 
 // BuildStateDiffObject builds a statediff object from two blocks and the provided parameters
@@ -181,16 +198,17 @@ func (sdb *builder) buildStateDiffWithIntermediateStateNodes(args Args, params P
 		return StateObject{}, fmt.Errorf("error building diff for updated accounts: %v", err)
 	}
 	// build the diff nodes for created accounts
-	createdAccounts, err := sdb.buildAccountCreations(diffAccountsAtB, params.WatchedStorageSlots, params.IntermediateStorageNodes)
+	createdAccounts, codeAndCodeHashes, err := sdb.buildAccountCreations(diffAccountsAtB, params.WatchedStorageSlots, params.IntermediateStorageNodes)
 	if err != nil {
 		return StateObject{}, fmt.Errorf("error building diff for created accounts: %v", err)
 	}
 
 	// assemble all of the nodes into the statediff object, including the intermediate nodes
 	return StateObject{
-		BlockNumber: args.BlockNumber,
-		BlockHash:   args.BlockHash,
-		Nodes:       append(append(append(updatedAccounts, createdAccounts...), createdOrUpdatedIntermediateNodes...), emptiedPaths...),
+		BlockNumber:       args.BlockNumber,
+		BlockHash:         args.BlockHash,
+		Nodes:             append(append(append(updatedAccounts, createdAccounts...), createdOrUpdatedIntermediateNodes...), emptiedPaths...),
+		CodeAndCodeHashes: codeAndCodeHashes,
 	}, nil
 }
 
@@ -235,16 +253,17 @@ func (sdb *builder) buildStateDiffWithoutIntermediateStateNodes(args Args, param
 		return StateObject{}, fmt.Errorf("error building diff for updated accounts: %v", err)
 	}
 	// build the diff nodes for created accounts
-	createdAccounts, err := sdb.buildAccountCreations(diffAccountsAtB, params.WatchedStorageSlots, params.IntermediateStorageNodes)
+	createdAccounts, codeAndCodeHashes, err := sdb.buildAccountCreations(diffAccountsAtB, params.WatchedStorageSlots, params.IntermediateStorageNodes)
 	if err != nil {
 		return StateObject{}, fmt.Errorf("error building diff for created accounts: %v", err)
 	}
 
 	// assemble all of the nodes into the statediff object
 	return StateObject{
-		BlockNumber: args.BlockNumber,
-		BlockHash:   args.BlockHash,
-		Nodes:       append(append(updatedAccounts, createdAccounts...), emptiedPaths...),
+		BlockNumber:       args.BlockNumber,
+		BlockHash:         args.BlockHash,
+		Nodes:             append(append(updatedAccounts, createdAccounts...), emptiedPaths...),
+		CodeAndCodeHashes: codeAndCodeHashes,
 	}, nil
 }
 
@@ -470,24 +489,40 @@ func (sdb *builder) buildAccountUpdates(creations, deletions AccountMap, updated
 }
 
 // buildAccountCreations returns the statediff node objects for all the accounts that exist at B but not at A
-func (sdb *builder) buildAccountCreations(accounts AccountMap, watchedStorageKeys []common.Hash, intermediateStorageNodes bool) ([]StateNode, error) {
+// it also returns the code and codehash for created contract accounts
+func (sdb *builder) buildAccountCreations(accounts AccountMap, watchedStorageKeys []common.Hash, intermediateStorageNodes bool) ([]StateNode, []CodeAndCodeHash, error) {
 	accountDiffs := make([]StateNode, 0, len(accounts))
+	codeAndCodeHashes := make([]CodeAndCodeHash, 0)
 	for _, val := range accounts {
-		// For account creations, any storage node contained is a diff
-		storageDiffs, err := sdb.buildStorageNodesEventual(val.Account.Root, watchedStorageKeys, intermediateStorageNodes)
-		if err != nil {
-			return nil, fmt.Errorf("failed building eventual storage diffs for node %x\r\nerror: %v", val.Path, err)
+		diff := StateNode{
+			NodeType:  val.NodeType,
+			Path:      val.Path,
+			LeafKey:   val.LeafKey,
+			NodeValue: val.NodeValue,
+		}
+		if !bytes.Equal(val.Account.CodeHash, nullCodeHash) {
+			// For contract creations, any storage node contained is a diff
+			storageDiffs, err := sdb.buildStorageNodesEventual(val.Account.Root, watchedStorageKeys, intermediateStorageNodes)
+			if err != nil {
+				return nil, nil, fmt.Errorf("failed building eventual storage diffs for node %x\r\nerror: %v", val.Path, err)
+			}
+			diff.StorageNodes = storageDiffs
+			// emit codehash => code mappings for cod
+			codeHash := common.BytesToHash(val.Account.CodeHash)
+			addrHash := common.BytesToHash(val.LeafKey)
+			code, err := sdb.stateCache.ContractCode(addrHash, codeHash)
+			if err != nil {
+				return nil, nil, fmt.Errorf("failed to retrieve code for codehash %s for account with leafkey %s\r\n error: %v", codeHash.String(), addrHash.String(), err)
+			}
+			codeAndCodeHashes = append(codeAndCodeHashes, CodeAndCodeHash{
+				Hash: codeHash,
+				Code: code,
+			})
 		}
-		accountDiffs = append(accountDiffs, StateNode{
-			NodeType:     val.NodeType,
-			Path:         val.Path,
-			LeafKey:      val.LeafKey,
-			NodeValue:    val.NodeValue,
-			StorageNodes: storageDiffs,
-		})
+		accountDiffs = append(accountDiffs, diff)
 	}
 
-	return accountDiffs, nil
+	return accountDiffs, codeAndCodeHashes, nil
 }
 
 // buildStorageNodesEventual builds the storage diff node objects for a created account

diff --git a/statediff/builder_test.go b/statediff/builder_test.go
@@ -617,6 +617,12 @@ func TestBuilder(t *testing.T) {
 						StorageNodes: emptyStorage,
 					},
 				},
+				CodeAndCodeHashes: []statediff.CodeAndCodeHash{
+					{
+						Hash: testhelpers.CodeHash,
+						Code: testhelpers.ByteCodeAfterDeployment,
+					},
+				},
 			},
 		},
 		{
@@ -863,6 +869,12 @@ func TestBuilderWithIntermediateNodes(t *testing.T) {
 						StorageNodes: emptyStorage,
 					},
 				},
+				CodeAndCodeHashes: []statediff.CodeAndCodeHash{
+					{
+						Hash: testhelpers.CodeHash,
+						Code: testhelpers.ByteCodeAfterDeployment,
+					},
+				},
 			},
 		},
 		{
@@ -1071,6 +1083,12 @@ func TestBuilderWithWatchedAddressList(t *testing.T) {
 						StorageNodes: emptyStorage,
 					},
 				},
+				CodeAndCodeHashes: []statediff.CodeAndCodeHash{
+					{
+						Hash: testhelpers.CodeHash,
+						Code: testhelpers.ByteCodeAfterDeployment,
+					},
+				},
 			},
 		},
 		{
@@ -1235,6 +1253,12 @@ func TestBuilderWithWatchedAddressAndStorageKeyList(t *testing.T) {
 						StorageNodes: emptyStorage,
 					},
 				},
+				CodeAndCodeHashes: []statediff.CodeAndCodeHash{
+					{
+						Hash: testhelpers.CodeHash,
+						Code: testhelpers.ByteCodeAfterDeployment,
+					},
+				},
 			},
 		},
 		{
@@ -1826,6 +1850,12 @@ func TestBuilderWithMovedAccount(t *testing.T) {
 						},
 					},
 				},
+				CodeAndCodeHashes: []statediff.CodeAndCodeHash{
+					{
+						Hash: testhelpers.CodeHash,
+						Code: testhelpers.ByteCodeAfterDeployment,
+					},
+				},
 			},
 		},
 		{
@@ -1942,6 +1972,12 @@ func TestBuilderWithMovedAccountOnlyLeafs(t *testing.T) {
 						},
 					},
 				},
+				CodeAndCodeHashes: []statediff.CodeAndCodeHash{
+					{
+						Hash: testhelpers.CodeHash,
+						Code: testhelpers.ByteCodeAfterDeployment,
+					},
+				},
 			},
 		},
 		{
@@ -2118,6 +2154,12 @@ func TestBuildStateTrie(t *testing.T) {
 						StorageNodes: emptyStorage,
 					},
 				},
+				CodeAndCodeHashes: []statediff.CodeAndCodeHash{
+					{
+						Hash: testhelpers.CodeHash,
+						Code: testhelpers.ByteCodeAfterDeployment,
+					},
+				},
 			},
 		},
 		{
@@ -2193,6 +2235,12 @@ func TestBuildStateTrie(t *testing.T) {
 						StorageNodes: emptyStorage,
 					},
 				},
+				CodeAndCodeHashes: []statediff.CodeAndCodeHash{
+					{
+						Hash: testhelpers.CodeHash,
+						Code: testhelpers.ByteCodeAfterDeployment,
+					},
+				},
 			},
 		},
 	}

diff --git a/statediff/testhelpers/test_data.go b/statediff/testhelpers/test_data.go
@@ -57,14 +57,16 @@ var (
 	TestBankFunds   = big.NewInt(100000000)
 	Genesis         = core.GenesisBlockForTesting(Testdb, TestBankAddress, TestBankFunds)
 
-	Account1Key, _  = crypto.HexToECDSA("8a1f9a8f95be41cd7ccb6168179afb4504aefe388d1e14474d32c45c72ce7b7a")
-	Account2Key, _  = crypto.HexToECDSA("49a7b37aa6f6645917e7b807e9d1c00d4fa71f18343b0d4122a4d2df64dd6fee")
-	Account1Addr    = crypto.PubkeyToAddress(Account1Key.PublicKey) //0x703c4b2bD70c169f5717101CaeE543299Fc946C7
-	Account2Addr    = crypto.PubkeyToAddress(Account2Key.PublicKey) //0x0D3ab14BBaD3D99F4203bd7a11aCB94882050E7e
-	Account1LeafKey = AddressToLeafKey(Account1Addr)
-	Account2LeafKey = AddressToLeafKey(Account2Addr)
-	ContractCode    = common.Hex2Bytes("608060405234801561001057600080fd5b50336000806101000a81548173ffffffffffffffffffffffffffffffffffffffff021916908373ffffffffffffffffffffffffffffffffffffffff1602179055506040518060200160405280600160ff16815250600190600161007492919061007a565b506100e4565b82606481019282156100ae579160200282015b828111156100ad578251829060ff1690559160200191906001019061008d565b5b5090506100bb91906100bf565b5090565b6100e191905b808211156100dd5760008160009055506001016100c5565b5090565b90565b6101ca806100f36000396000f3fe608060405234801561001057600080fd5b50600436106100365760003560e01c806343d726d61461003b578063c16431b914610045575b600080fd5b61004361007d565b005b61007b6004803603604081101561005b57600080fd5b81019080803590602001909291908035906020019092919050505061015c565b005b6000809054906101000a900473ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff163373ffffffffffffffffffffffffffffffffffffffff1614610122576040517f08c379a00000000000000000000000000000000000000000000000000000000081526004018080602001828103825260228152602001806101746022913960400191505060405180910390fd5b6000809054906101000a900473ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff16ff5b806001836064811061016a57fe5b0181905550505056fe4f6e6c79206f776e65722063616e2063616c6c20746869732066756e6374696f6e2ea265627a7a72305820e3747183708fb6bff3f6f7a80fb57dcc1c19f83f9cb25457a3ed5c0424bde66864736f6c634300050a0032")
-	ContractAddr    common.Address
+	Account1Key, _          = crypto.HexToECDSA("8a1f9a8f95be41cd7ccb6168179afb4504aefe388d1e14474d32c45c72ce7b7a")
+	Account2Key, _          = crypto.HexToECDSA("49a7b37aa6f6645917e7b807e9d1c00d4fa71f18343b0d4122a4d2df64dd6fee")
+	Account1Addr            = crypto.PubkeyToAddress(Account1Key.PublicKey) //0x703c4b2bD70c169f5717101CaeE543299Fc946C7
+	Account2Addr            = crypto.PubkeyToAddress(Account2Key.PublicKey) //0x0D3ab14BBaD3D99F4203bd7a11aCB94882050E7e
+	Account1LeafKey         = AddressToLeafKey(Account1Addr)
+	Account2LeafKey         = AddressToLeafKey(Account2Addr)
+	ContractCode            = common.Hex2Bytes("608060405234801561001057600080fd5b50336000806101000a81548173ffffffffffffffffffffffffffffffffffffffff021916908373ffffffffffffffffffffffffffffffffffffffff1602179055506040518060200160405280600160ff16815250600190600161007492919061007a565b506100e4565b82606481019282156100ae579160200282015b828111156100ad578251829060ff1690559160200191906001019061008d565b5b5090506100bb91906100bf565b5090565b6100e191905b808211156100dd5760008160009055506001016100c5565b5090565b90565b6101ca806100f36000396000f3fe608060405234801561001057600080fd5b50600436106100365760003560e01c806343d726d61461003b578063c16431b914610045575b600080fd5b61004361007d565b005b61007b6004803603604081101561005b57600080fd5b81019080803590602001909291908035906020019092919050505061015c565b005b6000809054906101000a900473ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff163373ffffffffffffffffffffffffffffffffffffffff1614610122576040517f08c379a00000000000000000000000000000000000000000000000000000000081526004018080602001828103825260228152602001806101746022913960400191505060405180910390fd5b6000809054906101000a900473ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff16ff5b806001836064811061016a57fe5b0181905550505056fe4f6e6c79206f776e65722063616e2063616c6c20746869732066756e6374696f6e2ea265627a7a72305820e3747183708fb6bff3f6f7a80fb57dcc1c19f83f9cb25457a3ed5c0424bde66864736f6c634300050a0032")
+	ByteCodeAfterDeployment = common.Hex2Bytes("608060405234801561001057600080fd5b50600436106100365760003560e01c806343d726d61461003b578063c16431b914610045575b600080fd5b61004361007d565b005b61007b6004803603604081101561005b57600080fd5b81019080803590602001909291908035906020019092919050505061015c565b005b6000809054906101000a900473ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff163373ffffffffffffffffffffffffffffffffffffffff1614610122576040517f08c379a00000000000000000000000000000000000000000000000000000000081526004018080602001828103825260228152602001806101746022913960400191505060405180910390fd5b6000809054906101000a900473ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff16ff5b806001836064811061016a57fe5b0181905550505056fe4f6e6c79206f776e65722063616e2063616c6c20746869732066756e6374696f6e2ea265627a7a72305820e3747183708fb6bff3f6f7a80fb57dcc1c19f83f9cb25457a3ed5c0424bde66864736f6c634300050a0032")
+	CodeHash                = common.HexToHash("0xaaea5efba4fd7b45d7ec03918ac5d8b31aa93b48986af0e6b591f0f087c80127")
+	ContractAddr            common.Address
 
 	EmptyRootNode, _  = rlp.EncodeToBytes([]byte{})
 	EmptyContractRoot = crypto.Keccak256Hash(EmptyRootNode)

diff --git a/statediff/types.go b/statediff/types.go
@@ -41,6 +41,7 @@ type Params struct {
 	IncludeBlock             bool
 	IncludeReceipts          bool
 	IncludeTD                bool
+	IncludeCode              bool
 	WatchedAddresses         []common.Address
 	WatchedStorageSlots      []common.Hash
 }
@@ -82,9 +83,17 @@ func (sd *Payload) Encode() ([]byte, error) {
 
 // StateObject is the final output structure from the builder
 type StateObject struct {
-	BlockNumber *big.Int    `json:"blockNumber"     gencodec:"required"`
-	BlockHash   common.Hash `json:"blockHash"       gencodec:"required"`
-	Nodes       []StateNode `json:"nodes"           gencodec:"required"`
+	BlockNumber       *big.Int          `json:"blockNumber"     gencodec:"required"`
+	BlockHash         common.Hash       `json:"blockHash"       gencodec:"required"`
+	Nodes             []StateNode       `json:"nodes"           gencodec:"required"`
+	CodeAndCodeHashes []CodeAndCodeHash `json:"codeMapping"`
+}
+
+// CodeAndCodeHash struct for holding codehash => code mappings
+// we can't use an actual map because they are not rlp serializable
+type CodeAndCodeHash struct {
+	Hash common.Hash
+	Code []byte
 }
 
 // StateNode holds the data for a single state diff node