github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/core/mpt/billet.go (about)

     1  package mpt
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"errors"
     7  	"fmt"
     8  
     9  	"github.com/nspcc-dev/neo-go/pkg/core/storage"
    10  	"github.com/nspcc-dev/neo-go/pkg/io"
    11  	"github.com/nspcc-dev/neo-go/pkg/util"
    12  )
    13  
    14  var (
    15  	// ErrRestoreFailed is returned when replacing HashNode by its "unhashed"
    16  	// candidate fails.
    17  	ErrRestoreFailed = errors.New("failed to restore MPT node")
    18  	errStop          = errors.New("stop condition is met")
    19  )
    20  
    21  // Billet is a part of an MPT trie with missing hash nodes that need to be restored.
    22  // Billet is based on the following assumptions:
    23  //  1. Refcount can only be incremented (we don't change the MPT structure during restore,
    24  //     thus don't need to decrease refcount).
    25  //  2. Each time a part of a Billet is completely restored, it is collapsed into
    26  //     HashNode.
    27  //  3. Any pair (node, path) must be restored only once. It's a duty of an MPT pool to manage
    28  //     MPT paths in order to provide this assumption.
    29  type Billet struct {
    30  	TempStoragePrefix storage.KeyPrefix
    31  	Store             *storage.MemCachedStore
    32  
    33  	root Node
    34  	mode TrieMode
    35  }
    36  
    37  // NewBillet returns a new billet for MPT trie restoring. It accepts a MemCachedStore
    38  // to decouple storage errors from logic errors so that all storage errors are
    39  // processed during `store.Persist()` at the caller. Another benefit is
    40  // that every `Put` can be considered an atomic operation.
    41  func NewBillet(rootHash util.Uint256, mode TrieMode, prefix storage.KeyPrefix, store *storage.MemCachedStore) *Billet {
    42  	return &Billet{
    43  		TempStoragePrefix: prefix,
    44  		Store:             store,
    45  		root:              NewHashNode(rootHash),
    46  		mode:              mode,
    47  	}
    48  }
    49  
    50  // RestoreHashNode replaces HashNode located at the provided path by the specified Node
    51  // and stores it. It also maintains the MPT as small as possible by collapsing those parts
    52  // of the MPT that have been completely restored.
    53  func (b *Billet) RestoreHashNode(path []byte, node Node) error {
    54  	if _, ok := node.(*HashNode); ok {
    55  		return fmt.Errorf("%w: unable to restore node into HashNode", ErrRestoreFailed)
    56  	}
    57  	if _, ok := node.(EmptyNode); ok {
    58  		return fmt.Errorf("%w: unable to restore node into EmptyNode", ErrRestoreFailed)
    59  	}
    60  	r, err := b.putIntoNode(b.root, path, node)
    61  	if err != nil {
    62  		return err
    63  	}
    64  	b.root = r
    65  
    66  	// If it's a leaf, then put into temporary contract storage.
    67  	if leaf, ok := node.(*LeafNode); ok {
    68  		if b.TempStoragePrefix == 0 {
    69  			panic("invalid storage prefix")
    70  		}
    71  		k := append([]byte{byte(b.TempStoragePrefix)}, fromNibbles(path)...)
    72  		b.Store.Put(k, leaf.value)
    73  	}
    74  	return nil
    75  }
    76  
    77  // putIntoNode puts val with the provided path inside curr and returns an updated node.
    78  // Reference counters are updated for both curr and returned value.
    79  func (b *Billet) putIntoNode(curr Node, path []byte, val Node) (Node, error) {
    80  	switch n := curr.(type) {
    81  	case *LeafNode:
    82  		return b.putIntoLeaf(n, path, val)
    83  	case *BranchNode:
    84  		return b.putIntoBranch(n, path, val)
    85  	case *ExtensionNode:
    86  		return b.putIntoExtension(n, path, val)
    87  	case *HashNode:
    88  		return b.putIntoHash(n, path, val)
    89  	case EmptyNode:
    90  		return nil, fmt.Errorf("%w: can't modify EmptyNode during restore", ErrRestoreFailed)
    91  	default:
    92  		panic("invalid MPT node type")
    93  	}
    94  }
    95  
    96  func (b *Billet) putIntoLeaf(curr *LeafNode, path []byte, val Node) (Node, error) {
    97  	if len(path) != 0 {
    98  		return nil, fmt.Errorf("%w: can't modify LeafNode during restore", ErrRestoreFailed)
    99  	}
   100  	if curr.Hash() != val.Hash() {
   101  		return nil, fmt.Errorf("%w: bad Leaf node hash: expected %s, got %s", ErrRestoreFailed, curr.Hash().StringBE(), val.Hash().StringBE())
   102  	}
   103  	// Once Leaf node is restored, it will be collapsed into HashNode forever, so
   104  	// there shouldn't be such situation when we try to restore a Leaf node.
   105  	panic("bug: can't restore LeafNode")
   106  }
   107  
   108  func (b *Billet) putIntoBranch(curr *BranchNode, path []byte, val Node) (Node, error) {
   109  	if len(path) == 0 && curr.Hash().Equals(val.Hash()) {
   110  		// This node has already been restored, so it's an MPT pool duty to avoid
   111  		// duplicating restore requests.
   112  		panic("bug: can't perform restoring of BranchNode twice")
   113  	}
   114  	i, path := splitPath(path)
   115  	r, err := b.putIntoNode(curr.Children[i], path, val)
   116  	if err != nil {
   117  		return nil, err
   118  	}
   119  	curr.Children[i] = r
   120  	return b.tryCollapseBranch(curr), nil
   121  }
   122  
   123  func (b *Billet) putIntoExtension(curr *ExtensionNode, path []byte, val Node) (Node, error) {
   124  	if len(path) == 0 {
   125  		if curr.Hash() != val.Hash() {
   126  			return nil, fmt.Errorf("%w: bad Extension node hash: expected %s, got %s", ErrRestoreFailed, curr.Hash().StringBE(), val.Hash().StringBE())
   127  		}
   128  		// This node has already been restored, so it's an MPT pool duty to avoid
   129  		// duplicating restore requests.
   130  		panic("bug: can't perform restoring of ExtensionNode twice")
   131  	}
   132  	if !bytes.HasPrefix(path, curr.key) {
   133  		return nil, fmt.Errorf("%w: can't modify ExtensionNode during restore", ErrRestoreFailed)
   134  	}
   135  
   136  	r, err := b.putIntoNode(curr.next, path[len(curr.key):], val)
   137  	if err != nil {
   138  		return nil, err
   139  	}
   140  	curr.next = r
   141  	return b.tryCollapseExtension(curr), nil
   142  }
   143  
   144  func (b *Billet) putIntoHash(curr *HashNode, path []byte, val Node) (Node, error) {
   145  	// Once a part of the MPT Billet is completely restored, it will be collapsed forever, so
   146  	// it's an MPT pool duty to avoid duplicating restore requests.
   147  	if len(path) != 0 {
   148  		return nil, fmt.Errorf("%w: node has already been collapsed", ErrRestoreFailed)
   149  	}
   150  
   151  	// `curr` hash node can be either of
   152  	// 1) saved in the storage (i.g. if we've already restored a node with the same hash from the
   153  	//    other part of the MPT), so just add it to the local in-memory MPT.
   154  	// 2) missing from the storage. It's OK because we're syncing MPT state, and the purpose
   155  	//    is to store missing hash nodes.
   156  	// both cases are OK, but we still need to validate `val` against `curr`.
   157  	if val.Hash() != curr.Hash() {
   158  		return nil, fmt.Errorf("%w: can't restore HashNode: expected and actual hashes mismatch (%s vs %s)", ErrRestoreFailed, curr.Hash().StringBE(), val.Hash().StringBE())
   159  	}
   160  
   161  	if curr.Collapsed {
   162  		// This node has already been restored and collapsed, so it's an MPT pool duty to avoid
   163  		// duplicating restore requests.
   164  		panic("bug: can't perform restoring of collapsed node")
   165  	}
   166  
   167  	// We also need to increment refcount in both cases. That's the only place where refcount
   168  	// is changed during restore process. Also flush right now, because sync process can be
   169  	// interrupted at any time.
   170  	b.incrementRefAndStore(val.Hash(), val.Bytes())
   171  
   172  	if val.Type() == LeafT {
   173  		return b.tryCollapseLeaf(val.(*LeafNode)), nil
   174  	}
   175  	return val, nil
   176  }
   177  
   178  func (b *Billet) incrementRefAndStore(h util.Uint256, bs []byte) {
   179  	key := makeStorageKey(h)
   180  	if b.mode.RC() {
   181  		var (
   182  			err  error
   183  			data []byte
   184  			cnt  int32
   185  		)
   186  		// An item may already be in store.
   187  		data, err = b.Store.Get(key)
   188  		if err == nil {
   189  			cnt = int32(binary.LittleEndian.Uint32(data[len(data)-4:]))
   190  		}
   191  		cnt++
   192  		if len(data) == 0 {
   193  			data = append(bs, 1, 0, 0, 0, 0)
   194  		}
   195  		binary.LittleEndian.PutUint32(data[len(data)-4:], uint32(cnt))
   196  		b.Store.Put(key, data)
   197  	} else {
   198  		b.Store.Put(key, bs)
   199  	}
   200  }
   201  
   202  // Traverse traverses MPT nodes (pre-order) starting from the billet root down
   203  // to its children calling `process` for each serialised node until true is
   204  // returned from `process` function. It also replaces all HashNodes to their
   205  // "unhashed" counterparts until the stop condition is satisfied.
   206  func (b *Billet) Traverse(process func(pathToNode []byte, node Node, nodeBytes []byte) bool, ignoreStorageErr bool) error {
   207  	r, err := b.traverse(b.root, []byte{}, []byte{}, process, ignoreStorageErr, false)
   208  	if err != nil && !errors.Is(err, errStop) {
   209  		return err
   210  	}
   211  	b.root = r
   212  	return nil
   213  }
   214  
   215  func (b *Billet) traverse(curr Node, path, from []byte, process func(pathToNode []byte, node Node, nodeBytes []byte) bool, ignoreStorageErr bool, backwards bool) (Node, error) {
   216  	if _, ok := curr.(EmptyNode); ok {
   217  		// We're not interested in EmptyNodes, and they do not affect the
   218  		// traversal process, thus remain them untouched.
   219  		return curr, nil
   220  	}
   221  	if hn, ok := curr.(*HashNode); ok {
   222  		r, err := b.GetFromStore(hn.Hash())
   223  		if err != nil {
   224  			if ignoreStorageErr && errors.Is(err, storage.ErrKeyNotFound) {
   225  				return hn, nil
   226  			}
   227  			return nil, err
   228  		}
   229  		return b.traverse(r, path, from, process, ignoreStorageErr, backwards)
   230  	}
   231  	if len(from) == 0 {
   232  		bytes := bytes.Clone(curr.Bytes())
   233  		if process(fromNibbles(path), curr, bytes) {
   234  			return curr, errStop
   235  		}
   236  	}
   237  	switch n := curr.(type) {
   238  	case *LeafNode:
   239  		return b.tryCollapseLeaf(n), nil
   240  	case *BranchNode:
   241  		var (
   242  			startIndex byte
   243  			endIndex   byte = childrenCount
   244  			cmp             = func(i int) bool {
   245  				return i < int(endIndex)
   246  			}
   247  			step = 1
   248  		)
   249  		if backwards {
   250  			startIndex, endIndex = lastChild, startIndex
   251  			cmp = func(i int) bool {
   252  				return i >= int(endIndex)
   253  			}
   254  			step = -1
   255  		}
   256  		if len(from) != 0 {
   257  			endIndex = lastChild
   258  			if backwards {
   259  				endIndex = 0
   260  			}
   261  			startIndex, from = splitPath(from)
   262  		}
   263  		for i := int(startIndex); cmp(i); i += step {
   264  			var newPath []byte
   265  			if i == lastChild {
   266  				newPath = path
   267  			} else {
   268  				newPath = append(path, byte(i))
   269  			}
   270  			if byte(i) != startIndex {
   271  				from = []byte{}
   272  			}
   273  			r, err := b.traverse(n.Children[i], newPath, from, process, ignoreStorageErr, backwards)
   274  			if err != nil {
   275  				if !errors.Is(err, errStop) {
   276  					return nil, err
   277  				}
   278  				n.Children[i] = r
   279  				return b.tryCollapseBranch(n), err
   280  			}
   281  			n.Children[i] = r
   282  		}
   283  		return b.tryCollapseBranch(n), nil
   284  	case *ExtensionNode:
   285  		if len(from) != 0 && bytes.HasPrefix(from, n.key) {
   286  			from = from[len(n.key):]
   287  		} else if len(from) == 0 || bytes.Compare(n.key, from) > 0 {
   288  			from = []byte{}
   289  		} else {
   290  			return b.tryCollapseExtension(n), nil
   291  		}
   292  		r, err := b.traverse(n.next, append(path, n.key...), from, process, ignoreStorageErr, backwards)
   293  		if err != nil && !errors.Is(err, errStop) {
   294  			return nil, err
   295  		}
   296  		n.next = r
   297  		return b.tryCollapseExtension(n), err
   298  	default:
   299  		return nil, ErrNotFound
   300  	}
   301  }
   302  
   303  func (b *Billet) tryCollapseLeaf(curr *LeafNode) Node {
   304  	// Leaf can always be collapsed.
   305  	res := NewHashNode(curr.Hash())
   306  	res.Collapsed = true
   307  	return res
   308  }
   309  
   310  func (b *Billet) tryCollapseExtension(curr *ExtensionNode) Node {
   311  	if !(curr.next.Type() == HashT && curr.next.(*HashNode).Collapsed) {
   312  		return curr
   313  	}
   314  	res := NewHashNode(curr.Hash())
   315  	res.Collapsed = true
   316  	return res
   317  }
   318  
   319  func (b *Billet) tryCollapseBranch(curr *BranchNode) Node {
   320  	canCollapse := true
   321  	for i := 0; i < childrenCount; i++ {
   322  		if curr.Children[i].Type() == EmptyT {
   323  			continue
   324  		}
   325  		if curr.Children[i].Type() == HashT && curr.Children[i].(*HashNode).Collapsed {
   326  			continue
   327  		}
   328  		canCollapse = false
   329  		break
   330  	}
   331  	if !canCollapse {
   332  		return curr
   333  	}
   334  	res := NewHashNode(curr.Hash())
   335  	res.Collapsed = true
   336  	return res
   337  }
   338  
   339  // GetFromStore returns MPT node from the storage.
   340  func (b *Billet) GetFromStore(h util.Uint256) (Node, error) {
   341  	data, err := b.Store.Get(makeStorageKey(h))
   342  	if err != nil {
   343  		return nil, err
   344  	}
   345  
   346  	var n NodeObject
   347  	r := io.NewBinReaderFromBuf(data)
   348  	n.DecodeBinary(r)
   349  	if r.Err != nil {
   350  		return nil, r.Err
   351  	}
   352  
   353  	if b.mode.RC() {
   354  		data = data[:len(data)-5]
   355  	}
   356  	n.Node.(flushedNode).setCache(data, h)
   357  	return n.Node, nil
   358  }