From a8c13ab0db8972f0209745f02e2191a94464e5b4 Mon Sep 17 00:00:00 2001 From: Eduard S Date: Thu, 23 Sep 2021 17:44:14 +0200 Subject: [PATCH] Return err before silently discarding bytes --- tree.go | 69 ++++++++++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 59 insertions(+), 10 deletions(-) diff --git a/tree.go b/tree.go index a594b30..3e51ad6 100644 --- a/tree.go +++ b/tree.go @@ -39,6 +39,9 @@ const ( // nChars is used to crop the Graphviz nodes labels nChars = 4 + + maxKeyLen = 255 + maxSiblingsFieldLen = 65535 ) var ( @@ -243,16 +246,21 @@ func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]int, err func (t *Tree) loadVT(rTx db.ReadTx) (vt, error) { vt := newVT(t.maxLevels, t.hashFunction) vt.params.dbg = t.dbg - err := t.IterateWithTx(rTx, nil, func(k, v []byte) { + var callbackErr error + err := t.IterateWithStopWithTx(rTx, nil, func(_ int, k, v []byte) bool { if v[0] != PrefixValueLeaf { - return + return false } leafK, leafV := ReadLeafValue(v) if err := vt.add(0, leafK, leafV); err != nil { - // TODO instead of panic, return this error - panic(err) + callbackErr = err + return true } + return false }) + if callbackErr != nil { + return vt, callbackErr + } return vt, err } @@ -304,6 +312,9 @@ func (t *Tree) AddWithTx(wTx db.WriteTx, k, v []byte) error { func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) { keyPath := make([]byte, t.hashFunction.Len()) + if len(k) > len(keyPath) { + return nil, fmt.Errorf("len(k) > len(keyPath)") + } copy(keyPath[:], k) path := getPath(t.maxLevels, keyPath) @@ -376,6 +387,10 @@ func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte, } oldLeafKeyFull := make([]byte, t.hashFunction.Len()) + if len(oldLeafKey) > len(oldLeafKeyFull) { + return nil, nil, nil, + fmt.Errorf("len(oldLeafKey) > len(oldLeafKeyFull)") + } copy(oldLeafKeyFull[:], oldLeafKey) // if currKey is already used, go down until paths diverge @@ -479,6 +494,9 @@ func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) { } var leafValue []byte leafValue = append(leafValue, byte(PrefixValueLeaf)) + if len(k) > maxKeyLen { + return nil, nil, fmt.Errorf("len(k) > %v", maxKeyLen) + } leafValue = append(leafValue, byte(len(k))) leafValue = append(leafValue, k...) leafValue = append(leafValue, v...) @@ -514,6 +532,9 @@ func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) { func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) { b := make([]byte, PrefixValueLen+hashFunc.Len()*2) b[0] = PrefixValueIntermediate + if len(l) > maxKeyLen { + return nil, nil, fmt.Errorf("len(k) > %v", maxKeyLen) + } b[1] = byte(len(l)) copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l) copy(b[PrefixValueLen+hashFunc.Len():], r) @@ -575,6 +596,9 @@ func (t *Tree) UpdateWithTx(wTx db.WriteTx, k, v []byte) error { var err error keyPath := make([]byte, t.hashFunction.Len()) + if len(k) > len(keyPath) { + return fmt.Errorf("len(k) > len(keyPath)") + } copy(keyPath[:], k) path := getPath(t.maxLevels, keyPath) @@ -632,6 +656,9 @@ func (t *Tree) GenProof(k []byte) ([]byte, []byte, []byte, bool, error) { // the db.ReadTx that is used. func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) { keyPath := make([]byte, t.hashFunction.Len()) + if len(k) > len(keyPath) { + return nil, nil, nil, false, fmt.Errorf("len(k) > len(keyPath)") + } copy(keyPath[:], k) root, err := t.RootWithTx(rTx) @@ -647,7 +674,10 @@ func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, return nil, nil, nil, false, err } - s := PackSiblings(t.hashFunction, siblings) + s, err := PackSiblings(t.hashFunction, siblings) + if err != nil { + return nil, nil, nil, false, err + } leafK, leafV := ReadLeafValue(value) if !bytes.Equal(k, leafK) { @@ -665,7 +695,7 @@ func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, // array. And S is the size of the output of the hash function used for the // Tree. The 2 2-byte that define the full length and bitmap length, are // encoded in little-endian. -func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte { +func PackSiblings(hashFunc HashFunction, siblings [][]byte) ([]byte, error) { var b []byte var bitmap []bool emptySibling := make([]byte, hashFunc.Len()) @@ -683,11 +713,17 @@ func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte { fullLen := 4 + l + len(b) //nolint:gomnd res := make([]byte, fullLen) + if fullLen > maxSiblingsFieldLen { + return nil, fmt.Errorf("fullLen > %v", maxSiblingsFieldLen) + } binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length - binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length + if l > maxSiblingsFieldLen { + return nil, fmt.Errorf("l > %v", maxSiblingsFieldLen) + } + binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length copy(res[4:4+l], bitmapBytes) copy(res[4+l:], b) - return res + return res, nil } // UnpackSiblings unpacks the siblings from a byte array. @@ -1038,18 +1074,31 @@ func (t *Tree) Dump(fromRoot []byte) ([]byte, error) { // WARNING current encoding only supports key & values of 255 bytes each // (due using only 1 byte for the length headers). var b []byte - err := t.Iterate(fromRoot, func(k, v []byte) { + var callbackErr error + err := t.IterateWithStop(fromRoot, func(_ int, k, v []byte) bool { if v[0] != PrefixValueLeaf { - return + return false } leafK, leafV := ReadLeafValue(v) kv := make([]byte, 2+len(leafK)+len(leafV)) + if len(leafK) > maxKeyLen { + callbackErr = fmt.Errorf("len(leafK) > %v", maxKeyLen) + return true + } kv[0] = byte(len(leafK)) + if len(leafV) > maxKeyLen { + callbackErr = fmt.Errorf("len(leafV) > %v", maxKeyLen) + return true + } kv[1] = byte(len(leafV)) copy(kv[2:2+len(leafK)], leafK) copy(kv[2+len(leafK):], leafV) b = append(b, kv...) + return false }) + if callbackErr != nil { + return nil, callbackErr + } return b, err }