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 }