Merge branch 'extension-into-branch' (only .go files) into feat/#1752-fix-witness-generation-json-export

Author: KimiWu123

geth-utils/gethutil/mpt/oracle/prefetch.go

@@ -81,7 +81,8 @@ var RemoteUrl = "https://mainnet.infura.io/v3/9aa3d95b3bc440fa88ea12eaa4456161"
 var LocalUrl = "http://localhost:8545"
 
 // For generating special tests for MPT circuit:
-var PreventHashingInSecureTrie = false
+var PreventHashingInSecureTrie = false // storage
+var AccountPreventHashingInSecureTrie = false
 
 func toFilename(key string) string {
 	return fmt.Sprintf("/tmp/eth/json_%s", key)

geth-utils/gethutil/mpt/state/database.go

@@ -52,7 +52,7 @@ func (db *Database) CopyTrie(t Trie) Trie {
 
 // OpenTrie opens the main account trie at a specific root hash.
 func (db *Database) OpenTrie(root common.Hash) (Trie, error) {
-	tr, err := trie.NewSecure(root, db.db)
+	tr, err := trie.NewSecure(root, db.db, false)
 	if err != nil {
 		return nil, err
 	}
@@ -62,7 +62,7 @@ func (db *Database) OpenTrie(root common.Hash) (Trie, error) {
 // OpenStorageTrie opens the storage trie of an account.
 func (db *Database) OpenStorageTrie(addrHash, root common.Hash) (Trie, error) {
 	//return SimpleTrie{db.BlockNumber, root, true, addrHash}, nil
-	tr, err := trie.NewSecure(root, db.db)
+	tr, err := trie.NewSecure(root, db.db, true)
 	if err != nil {
 		return nil, err
 	}

geth-utils/gethutil/mpt/state/statedb.go

@@ -298,7 +298,12 @@ func (s *StateDB) GetState(addr common.Address, hash common.Hash) common.Hash {
 
 // GetProof returns the Merkle proof for a given account.
 func (s *StateDB) GetProof(addr common.Address) ([][]byte, []byte, [][]byte, bool, bool, error) {
-	return s.GetProofByHash(crypto.Keccak256Hash(addr.Bytes()))
+	newAddr := crypto.Keccak256Hash(addr.Bytes())
+	if oracle.AccountPreventHashingInSecureTrie {
+		bytes := append(addr.Bytes(), []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
+		newAddr = common.BytesToHash(bytes)
+	}
+	return s.GetProofByHash(newAddr)
 }
 
 // GetProofByHash returns the Merkle proof for a given account.
@@ -311,7 +316,13 @@ func (s *StateDB) GetProofByHash(addrHash common.Hash) ([][]byte, []byte, [][]by
 // GetStorageProof returns the Merkle proof for given storage slot.
 func (s *StateDB) GetStorageProof(a common.Address, key common.Hash) ([][]byte, []byte, [][]byte, bool, bool, error) {
 	var proof proofList
-	trie := s.StorageTrie(a)
+	newAddr := a
+	if oracle.AccountPreventHashingInSecureTrie {
+		bytes := append(a.Bytes(), []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
+		newAddr = common.BytesToAddress(bytes)
+	}
+
+	trie := s.StorageTrie(newAddr)
 	if trie == nil {
 		return proof, nil, nil, false, false, errors.New("storage trie for requested address does not exist")
 	}
@@ -457,7 +468,7 @@ func (s *StateDB) SetStateObjectIfExists(addr common.Address) {
 			}
 
 			/*
-				When an account that does not exist is tried to be fetched by PrefetchAccount and when the some other account
+				When an account that does not exist is being fetched by PrefetchAccount and when some other account
 				exist at the overlapping address (the beginning of it), this (wrong) account is obtained by PrefetchAccount
 				and needs to be ignored.
 			*/
@@ -533,8 +544,15 @@ func (s *StateDB) updateStateObject(obj *stateObject) {
 	if err != nil {
 		panic(fmt.Errorf("can't encode object at %x: %v", addr[:], err))
 	}
-	if err = s.trie.TryUpdateAlwaysHash(addr[:], data); err != nil {
-		s.setError(fmt.Errorf("updateStateObject (%x) error: %v", addr[:], err))
+
+	if !oracle.AccountPreventHashingInSecureTrie {
+		if err = s.trie.TryUpdateAlwaysHash(addr[:], data); err != nil {
+			s.setError(fmt.Errorf("updateStateObject (%x) error: %v", addr[:], err))
+		}
+	} else {
+		if err = s.trie.TryUpdate(addr[:], data); err != nil {
+			s.setError(fmt.Errorf("updateStateObject (%x) error: %v", addr[:], err))
+		}
 	}
 
 	// If state snapshotting is active, cache the data til commit. Note, this

geth-utils/gethutil/mpt/trie/proof.go

@@ -102,33 +102,27 @@ func (t *Trie) Prove(key []byte, fromLevel uint, proofDb ethdb.KeyValueWriter) (
 			fromLevel--
 			continue
 		}
-		// var hn Node
 
 		// We need nibbles in witness for extension keys.
 		// copy n.Key before it gets changed in ProofHash
 		var nCopy []byte
 		if short, ok := n.(*ShortNode); ok {
-			if !hasTerm(short.Key) { // only for extension keys
+			if !hasTerm(short.Key) { // extension keys
 				nCopy = make([]byte, len(short.Key))
 				copy(nCopy, short.Key)
 				extNibbles = append(extNibbles, nCopy)
+			} else {
+				extNibbles = append(extNibbles, []byte{})
 			}
+		} else {
+			extNibbles = append(extNibbles, []byte{})
 		}
 
-		// n, hn = hasher.ProofHash(n)
 		n, _ = hasher.ProofHash(n)
-		// if hash, ok := hn.(HashNode); ok || i == 0 {
 		// If the node's database encoding is a hash (or is the
 		// root node), it becomes a proof element.
 		enc, _ := rlp.EncodeToBytes(n)
-		/*
-			if !ok {
-				hash = hasher.HashData(enc)
-			}
-		*/
-		// proofDb.Put(hash, enc)
 		proofDb.Put([]byte{1, 1, 1}, enc)
-		// }
 	}
 
 	isNeighbourNodeHashed := false

geth-utils/gethutil/mpt/trie/secure_trie.go

@@ -40,6 +40,7 @@ type SecureTrie struct {
 	hashKeyBuf       [common.HashLength]byte
 	secKeyCache      map[string][]byte
 	secKeyCacheOwner *SecureTrie // Pointer to self, replace the key cache on mismatch
+	isStorageTrie    bool
 }
 
 // NewSecure creates a trie with an existing root node from a backing database
@@ -53,15 +54,15 @@ type SecureTrie struct {
 // Loaded nodes are kept around until their 'cache generation' expires.
 // A new cache generation is created by each call to Commit.
 // cachelimit sets the number of past cache generations to keep.
-func NewSecure(root common.Hash, db *Database) (*SecureTrie, error) {
+func NewSecure(root common.Hash, db *Database, isStorageTrie bool) (*SecureTrie, error) {
 	if db == nil {
 		panic("trie.NewSecure called without a database")
 	}
 	trie, err := New(root, db)
 	if err != nil {
 		return nil, err
 	}
-	return &SecureTrie{trie: *trie}, nil
+	return &SecureTrie{trie: *trie, isStorageTrie: isStorageTrie}, nil
 }
 
 // Get returns the value for key stored in the trie.
@@ -202,7 +203,8 @@ func (t *SecureTrie) NodeIterator(start []byte) NodeIterator {
 // The caller must not hold onto the return value because it will become
 // invalid on the next call to hashKey or secKey.
 func (t *SecureTrie) hashKey(key []byte) []byte {
-	if !oracle.PreventHashingInSecureTrie {
+	preventHashing := (oracle.PreventHashingInSecureTrie && t.isStorageTrie) || (oracle.AccountPreventHashingInSecureTrie && !t.isStorageTrie)
+	if !preventHashing {
 		h := NewHasher(false)
 		h.sha.Reset()
 		h.sha.Write(key)
@@ -211,6 +213,9 @@ func (t *SecureTrie) hashKey(key []byte) []byte {
 		return t.hashKeyBuf[:]
 	} else {
 		// For generating special tests for MPT circuit.
+		if len(key) < 32 { // accounts
+			key = append(key, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}...)
+		}
 		return key
 	}
 }

geth-utils/gethutil/mpt/witness/branch.go

@@ -70,8 +70,7 @@ func prepareBranchWitness(rows [][]byte, branch []byte, branchStart int, branchR
 func prepareBranchNode(
 	branch1, branch2, extNode1, extNode2, extListRlpBytes []byte,
 	extValues [][]byte, key, driftedInd byte,
-	isBranchSPlaceholder, isBranchCPlaceholder, isExtension bool) Node {
-
+	isBranchSPlaceholder, isBranchCPlaceholder, isExtension, isLastLevel bool) Node {
 	extensionNode := ExtensionNode{
 		ListRlpBytes: extListRlpBytes,
 	}
@@ -118,10 +117,11 @@ func prepareBranchNode(
 	}
 
 	extensionBranch := ExtensionBranchNode{
-		IsExtension:   isExtension,
-		IsPlaceholder: [2]bool{isBranchSPlaceholder, isBranchCPlaceholder},
-		Extension:     extensionNode,
-		Branch:        branchNode,
+		IsExtension:                isExtension,
+		IsPlaceholder:              [2]bool{isBranchSPlaceholder, isBranchCPlaceholder},
+		IsLastLevelAndWrongExtCase: isLastLevel,
+		Extension:                  extensionNode,
+		Branch:                     branchNode,
 	}
 
 	values := make([][]byte, 17)
@@ -311,10 +311,14 @@ func addBranchAndPlaceholder(
 	driftedInd := getDriftedPosition(leafRow0, numberOfNibbles)
 	if len1 > len2 {
 		node = prepareBranchNode(proof1[len1-2], proof1[len1-2], extNode, extNode, extListRlpBytes, extValues,
-			key[keyIndex+numberOfNibbles], driftedInd, false, true, isExtension)
+			key[keyIndex+numberOfNibbles], driftedInd, false, true, isExtension, false)
+
+		// We now get the first nibble of the leaf that was turned into branch.
+		// This first nibble presents the position of the leaf once it moved
+		// into the new branch.
 	} else {
 		node = prepareBranchNode(proof2[len2-2], proof2[len2-2], extNode, extNode, extListRlpBytes, extValues,
-			key[keyIndex+numberOfNibbles], driftedInd, true, false, isExtension)
+			key[keyIndex+numberOfNibbles], driftedInd, true, false, isExtension, false)
 	}
 
 	return isModifiedExtNode, isExtension, numberOfNibbles, node

geth-utils/gethutil/mpt/witness/leaf.go

@@ -122,7 +122,7 @@ func getStorageRootCodeHashValue(leaf []byte, storageStart int) ([]byte, []byte)
 	return storageRootValue, codeHashValue
 }
 
-func prepareAccountLeafNode(addr common.Address, addrh []byte, leafS, leafC, neighbourNode, addressNibbles []byte, isPlaceholder, isSModExtension, isCModExtension bool) Node {
+func prepareAccountLeafNode(addr common.Address, addrh []byte, leafS, leafC, constructedLeaf, neighbourNode, addressNibbles []byte, isPlaceholder, isSModExtension, isCModExtension bool) Node {
 	// For non existing account proof there are two cases:
 	// 1. A leaf is returned that is not at the required address (wrong leaf).
 	// 2. A branch is returned as the last element of getProof and
@@ -174,12 +174,19 @@ func prepareAccountLeafNode(addr common.Address, addrh []byte, leafS, leafC, nei
 
 	// wrongValue is used only for proof that account doesn't exist
 
+	wrongLeaf := leafC
+	wrongLen := keyLenC
+	if constructedLeaf != nil {
+		wrongLeaf = constructedLeaf
+		wrongLen = int(constructedLeaf[2]) - 128
+	}
+
 	offset := 0
-	nibblesNum := (keyLenC - 1) * 2
-	wrongRlpBytes[0] = leafC[0]
-	wrongRlpBytes[1] = leafC[1]
-	wrongValue[0] = leafC[2] // length
-	if leafC[3] != 32 {      // odd number of nibbles
+	nibblesNum := (wrongLen - 1) * 2
+	wrongRlpBytes[0] = wrongLeaf[0]
+	wrongRlpBytes[1] = wrongLeaf[1]
+	wrongValue[0] = wrongLeaf[2] // length
+	if wrongLeaf[3] != 32 {      // odd number of nibbles
 		nibblesNum = nibblesNum + 1
 		wrongValue[1] = addressNibbles[64-nibblesNum] + 48
 		offset = 1
@@ -327,7 +334,7 @@ func prepareLeafAndPlaceholderNode(addr common.Address, addrh []byte, proof1, pr
 
 		// When generating a proof that account doesn't exist, the length of both proofs is the same (doesn't reach
 		// this code).
-		return prepareAccountLeafNode(addr, addrh, leafS, leafC, nil, key, false, isSModExtension, isCModExtension)
+		return prepareAccountLeafNode(addr, addrh, leafS, leafC, nil, nil, key, false, isSModExtension, isCModExtension)
 	} else {
 		var leaf []byte
 		isSPlaceholder := false
@@ -341,7 +348,7 @@ func prepareLeafAndPlaceholderNode(addr common.Address, addrh []byte, proof1, pr
 			isSPlaceholder = true
 		}
 
-		return prepareStorageLeafNode(leaf, leaf, nil, storage_key, key, false, isSPlaceholder, isCPlaceholder, isSModExtension, isCModExtension)
+		return prepareStorageLeafNode(leaf, leaf, nil, nil, storage_key, key, false, isSPlaceholder, isCPlaceholder, isSModExtension, isCModExtension)
 	}
 }
 
@@ -451,7 +458,7 @@ func prepareAccountLeafPlaceholderNode(addr common.Address, addrh, key []byte, k
 		leaf[4+i] = remainingNibbles[2*i+offset]*16 + remainingNibbles[2*i+1+offset]
 	}
 
-	node := prepareAccountLeafNode(addr, addrh, leaf, leaf, nil, key, true, false, false)
+	node := prepareAccountLeafNode(addr, addrh, leaf, leaf, nil, nil, key, true, false, false)
 
 	node.Account.ValueRlpBytes[0][0] = 184
 	node.Account.ValueRlpBytes[0][1] = 70
@@ -480,7 +487,7 @@ func prepareStorageLeafPlaceholderNode(storage_key common.Hash, key []byte, keyI
 	keyLen := getLeafKeyLen(keyIndex)
 	leaf[0] = 192 + 1 + byte(keyLen) + 1
 
-	return prepareStorageLeafNode(leaf, leaf, nil, storage_key, key, false, true, true, false, false)
+	return prepareStorageLeafNode(leaf, leaf, nil, nil, storage_key, key, false, true, true, false, false)
 }
 
 func setKeyValue(row, keyRlp []byte, keyLen, offset byte, valueIsZero, isPlaceholder bool) ([]byte, []byte, []byte) {
@@ -592,7 +599,7 @@ func prepareStorageLeafInfo(row []byte, valueIsZero, isPlaceholder bool) ([]byte
 	return key, value, keyRlp, valueRlp
 }
 
-func prepareStorageLeafNode(leafS, leafC, neighbourNode []byte, storage_key common.Hash, key []byte, nonExistingStorageProof, isSPlaceholder, isCPlaceholder, isSModExtension, isCModExtension bool) Node {
+func prepareStorageLeafNode(leafS, leafC, constructedLeaf, neighbourNode []byte, storage_key common.Hash, key []byte, nonExistingStorageProof, isSPlaceholder, isCPlaceholder, isSModExtension, isCModExtension bool) Node {
 	var rows [][]byte
 
 	keyS, valueS, listRlpBytes1, valueRlpBytes1 := prepareStorageLeafInfo(leafS, false, isSPlaceholder)
@@ -623,7 +630,11 @@ func prepareStorageLeafNode(leafS, leafC, neighbourNode []byte, storage_key comm
 	var nonExistingStorageRow []byte
 	var wrongRlpBytes []byte
 	if nonExistingStorageProof {
-		wrongRlpBytes, nonExistingStorageRow = prepareNonExistingStorageRow(leafC, key)
+		if constructedLeaf != nil {
+			wrongRlpBytes, nonExistingStorageRow = prepareNonExistingStorageRow(constructedLeaf, key)
+		} else {
+			wrongRlpBytes, nonExistingStorageRow = prepareNonExistingStorageRow(leafC, key)
+		}
 	} else {
 		nonExistingStorageRow = prepareEmptyNonExistingStorageRow()
 	}
@@ -645,6 +656,7 @@ func prepareStorageLeafNode(leafS, leafC, neighbourNode []byte, storage_key comm
 		ValueRlpBytes:   valueRlpBytes,
 		IsModExtension:  [2]bool{isSModExtension, isCModExtension},
 	}
+
 	keccakData := [][]byte{leafS, leafC, storage_key.Bytes()}
 	if neighbourNode != nil {
 		keccakData = append(keccakData, neighbourNode)
@@ -657,3 +669,38 @@ func prepareStorageLeafNode(leafS, leafC, neighbourNode []byte, storage_key comm
 
 	return node
 }
+
+func equipLeafWithWrongExtension(leafNode Node, keyMiddle, keyAfter, nibblesMiddle, nibblesAfter []byte) Node {
+	l := len(leafNode.Values)
+	leafNode.Values[l-modifiedExtensionNodeRowLen] = keyMiddle
+	startNibblePos := 2 // we don't need any nibbles for case keyLen = 1
+	if len(keyMiddle) > 1 {
+		if len(nibblesMiddle)%2 == 0 {
+			startNibblePos = 1
+		} else {
+			startNibblePos = 2
+		}
+	}
+	ind := 0
+	for j := startNibblePos; j < len(nibblesMiddle); j += 2 {
+		leafNode.Values[l-modifiedExtensionNodeRowLen+1][2+ind] = nibblesMiddle[j]
+		ind++
+	}
+
+	leafNode.Values[l-modifiedExtensionNodeRowLen+3] = keyAfter
+	startNibblePos = 2 // we don't need any nibbles for case keyLen = 1
+	if len(keyAfter) > 1 {
+		if len(nibblesAfter)%2 == 0 {
+			startNibblePos = 1
+		} else {
+			startNibblePos = 2
+		}
+	}
+	ind = 0
+	for j := startNibblePos; j < len(nibblesAfter); j += 2 {
+		leafNode.Values[l-modifiedExtensionNodeRowLen+4][2+ind] = nibblesAfter[j]
+		ind++
+	}
+
+	return leafNode
+}

geth-utils/gethutil/mpt/witness/modified_extension_node.go

@@ -56,19 +56,6 @@ func equipLeafWithModExtensionNode(
 		longExtNodeKey[j] = longNibbles[j-byte(keyIndex)]
 	}
 
-	k := trie.HexToKeybytes(longExtNodeKey)
-	ky := common.BytesToHash(k)
-	var proof [][]byte
-	var err error
-	if !isTxProof {
-		if isAccountProof {
-			proof, _, _, _, _, err = statedb.GetProof(addr)
-		} else {
-			proof, _, _, _, _, err = statedb.GetStorageProof(addr, ky)
-		}
-		check(err)
-	}
-
 	// There is no short extension node when `len(longNibbles) - numberOfNibbles = 1`, in this case there
 	// is simply a branch instead.
 	// stack trie is always a short ext node.
@@ -79,6 +66,19 @@ func equipLeafWithModExtensionNode(
 	var extValuesC [][]byte
 
 	if !shortExtNodeIsBranch {
+		k := trie.HexToKeybytes(longExtNodeKey)
+		ky := common.BytesToHash(k)
+		var proof [][]byte
+		var err error
+		if !isTxProof {
+			if isAccountProof {
+				proof, _, _, _, _, err = statedb.GetProof(addr)
+			} else {
+				proof, _, _, _, _, err = statedb.GetStorageProof(addr, ky)
+			}
+		}
+		check(err)
+
 		if len2 > len1 {
 			isItBranch := isBranch(proof[len(proof)-1])