Skip to content

Commit

Permalink
feat: support for generating proofs (#4)
Browse files Browse the repository at this point in the history
  • Loading branch information
agaffney authored Jun 25, 2024
1 parent f4ce882 commit 6ea345d
Show file tree
Hide file tree
Showing 9 changed files with 400 additions and 42 deletions.
40 changes: 17 additions & 23 deletions branch.go
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ func (b *Branch) updateHash() {
tmpVal = append(tmpVal, byte(nibble))
}
// Calculate merkle root for children and append
childrenHash := b.merkleRootChildren()
childrenHash := merkleRoot(b.children[:])
tmpVal = append(tmpVal, childrenHash.Bytes()...)
// Calculate hash
b.hash = HashValue(tmpVal)
Expand Down Expand Up @@ -260,30 +260,24 @@ func (b *Branch) delete(path []Nibble) error {
return nil
}

func (b *Branch) merkleRootChildren() Hash {
// Gather child node hashes
tmpHashes := make([]Hash, 0, len(b.children))
for _, child := range b.children {
tmpHash := NullHash
if child != nil {
tmpHash = child.Hash()
}
tmpHashes = append(tmpHashes, tmpHash)
func (b *Branch) generateProof(path []Nibble) (*Proof, error) {
// Determine path minus the current node prefix
pathMinusPrefix := path[len(b.prefix):]
// Determine which child slot the next nibble in the path fits in
childIdx := int(pathMinusPrefix[0])
// Determine sub-path for key. We strip off the first nibble, since it's implied by
// the child slot that it's in
subPath := pathMinusPrefix[1:]
if b.children[childIdx] == nil {
return nil, ErrKeyNotExist
}
// Concat and hash child hashes in pairs, repeating until only a single hash remains
for len(tmpHashes) > 1 {
newTmpHashes := make([]Hash, 0, len(tmpHashes)/2)
for i := 0; i < len(tmpHashes); i = i + 2 {
tmpVal := append(
tmpHashes[i].Bytes(),
tmpHashes[i+1].Bytes()...,
)
tmpHash := HashValue(tmpVal)
newTmpHashes = append(newTmpHashes, tmpHash)
}
tmpHashes = newTmpHashes
existingChild := b.children[childIdx]
proof, err := existingChild.generateProof(subPath)
if err != nil {
return nil, err
}
return tmpHashes[0]
proof.Rewind(childIdx, len(b.prefix), b.children[:])
return proof, nil
}

func (b *Branch) getChildren() []Node {
Expand Down
12 changes: 10 additions & 2 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,14 @@ module github.com/blinklabs-io/merkle-patricia-forestry

go 1.21.10

require golang.org/x/crypto v0.24.0
require (
github.com/blinklabs-io/gouroboros v0.89.1
golang.org/x/crypto v0.24.0
)

require golang.org/x/sys v0.21.0 // indirect
require (
github.com/fxamacker/cbor/v2 v2.7.0 // indirect
github.com/jinzhu/copier v0.4.0 // indirect
github.com/x448/float16 v0.8.4 // indirect
golang.org/x/sys v0.21.0 // indirect
)
8 changes: 8 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
@@ -1,3 +1,11 @@
github.com/blinklabs-io/gouroboros v0.89.1 h1:pcD9hc2EkiPkq915aMDBAbgQZTX4I73gUzZf2UUcggs=
github.com/blinklabs-io/gouroboros v0.89.1/go.mod h1:l6G9mwAa/p0CBGCZBjK1W67815gWrRlmcGl6fccbt4U=
github.com/fxamacker/cbor/v2 v2.7.0 h1:iM5WgngdRBanHcxugY4JySA0nk1wZorNOpTgCMedv5E=
github.com/fxamacker/cbor/v2 v2.7.0/go.mod h1:pxXPTn3joSm21Gbwsv0w9OSA2y1HFR9qXEeXQVeNoDQ=
github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8=
github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg=
github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
golang.org/x/crypto v0.24.0 h1:mnl8DM0o513X8fdIkmyFE/5hTYxbwYOjDS/+rK6qpRI=
golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws=
Expand Down
16 changes: 16 additions & 0 deletions leaf.go
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,22 @@ func (l *Leaf) Set(value []byte) {
l.updateHash()
}

func (l *Leaf) generateProof(path []Nibble) (*Proof, error) {
if string(path) != string(l.suffix) {
return nil, ErrKeyNotExist
}
leafPath := keyToPath(l.key)
var proofVal []byte
if string(path) == string(l.suffix) {
proofVal = append(proofVal, l.value...)
}
proof := newProof(
leafPath,
proofVal,
)
return proof, nil
}

func (l *Leaf) updateHash() {
tmpVal := []byte{}
head := hashHead(l.suffix)
Expand Down
7 changes: 7 additions & 0 deletions nibble.go
Original file line number Diff line number Diff line change
Expand Up @@ -61,3 +61,10 @@ func nibblesToHexString(data []Nibble) string {
}
return ret
}

// keyToPath converts an arbitrary key to the sequence of Nibbles representing the path to the value
func keyToPath(key []byte) []Nibble {
keyHash := HashValue(key)
keyHashNibbles := bytesToNibbles(keyHash.Bytes())
return keyHashNibbles
}
27 changes: 27 additions & 0 deletions node.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,4 +18,31 @@ type Node interface {
isNode()
Hash() Hash
String() string
generateProof([]Nibble) (*Proof, error)
}

func merkleRoot(nodes []Node) Hash {
// Gather child node hashes
tmpHashes := make([]Hash, 0, len(nodes))
for _, child := range nodes {
tmpHash := NullHash
if child != nil {
tmpHash = child.Hash()
}
tmpHashes = append(tmpHashes, tmpHash)
}
// Concat and hash child hashes in pairs, repeating until only a single hash remains
for len(tmpHashes) > 1 {
newTmpHashes := make([]Hash, 0, len(tmpHashes)/2)
for i := 0; i < len(tmpHashes); i = i + 2 {
tmpVal := append(
tmpHashes[i].Bytes(),
tmpHashes[i+1].Bytes()...,
)
tmpHash := HashValue(tmpVal)
newTmpHashes = append(newTmpHashes, tmpHash)
}
tmpHashes = newTmpHashes
}
return tmpHashes[0]
}
219 changes: 219 additions & 0 deletions proof.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,219 @@
// Copyright 2024 Blink Labs Software
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package mpf

import (
"errors"
"fmt"
"slices"

"github.com/blinklabs-io/gouroboros/cbor"
)

type ProofStepType int

const (
ProofStepTypeLeaf ProofStepType = 1
ProofStepTypeFork ProofStepType = 2
ProofStepTypeBranch ProofStepType = 3
)

func (p ProofStepType) String() string {
switch p {
case ProofStepTypeLeaf:
return "leaf"
case ProofStepTypeFork:
return "fork"
case ProofStepTypeBranch:
return "branch"
default:
return "unknown"
}
}

type Proof struct {
path []Nibble
value []byte
steps []ProofStep
}

func newProof(path []Nibble, value []byte) *Proof {
p := &Proof{}
if path != nil {
p.path = append(p.path, path...)
}
if value != nil {
p.value = append(p.value, value...)
}
return p
}

func (p *Proof) Rewind(targetIdx int, prefixLen int, neighbors []Node) {
var nonEmptyNeighbors []Node
var nonEmptyNeighborIdx int
for idx, neighbor := range neighbors {
if neighbor == nil {
continue
}
if idx == targetIdx {
continue
}
nonEmptyNeighbors = append(nonEmptyNeighbors, neighbor)
nonEmptyNeighborIdx = idx
}
if len(nonEmptyNeighbors) == 1 {
neighbor := nonEmptyNeighbors[0]
switch n := neighbor.(type) {
case *Leaf:
step := ProofStep{
stepType: ProofStepTypeLeaf,
prefixLength: prefixLen,
neighbor: ProofStepNeighbor{
key: keyToPath(n.key),
value: HashValue(n.value),
},
}
p.steps = slices.Insert(p.steps, 0, step)
case *Branch:
step := ProofStep{
stepType: ProofStepTypeFork,
prefixLength: prefixLen,
neighbor: ProofStepNeighbor{
prefix: n.prefix,
nibble: Nibble(nonEmptyNeighborIdx),
root: merkleRoot(n.children[:]),
},
}
p.steps = slices.Insert(p.steps, 0, step)
default:
panic(
fmt.Sprintf(
"unknown Node type %T...this should never happen",
neighbor,
),
)
}
} else {
step := ProofStep{
stepType: ProofStepTypeBranch,
prefixLength: prefixLen,
neighbors: merkleProof(neighbors, targetIdx),
}
p.steps = slices.Insert(p.steps, 0, step)
}
}

func (p *Proof) MarshalCBOR() ([]byte, error) {
tmpSteps := make([]any, 0, len(p.steps))
for _, step := range p.steps {
tmpSteps = append(tmpSteps, step)
}
tmpData := cbor.IndefLengthList(tmpSteps)
return cbor.Encode(&tmpData)
}

type ProofStep struct {
stepType ProofStepType
prefixLength int
neighbors []Hash
neighbor ProofStepNeighbor
}

func (s *ProofStep) MarshalCBOR() ([]byte, error) {
switch s.stepType {
case ProofStepTypeBranch:
var tmpNeighbors []byte
for _, neighbor := range s.neighbors {
tmpNeighbors = append(tmpNeighbors, neighbor.Bytes()...)
}
tmpData := cbor.NewConstructor(
0,
cbor.IndefLengthList{
s.prefixLength,
cbor.IndefLengthByteString{
tmpNeighbors[0:64],
tmpNeighbors[64:],
},
},
)
return cbor.Encode(&tmpData)

case ProofStepTypeFork:
tmpData := cbor.NewConstructor(
1,
cbor.IndefLengthList{
s.prefixLength,
cbor.NewConstructor(
0,
cbor.IndefLengthList{
int(s.neighbor.nibble),
s.neighbor.prefix,
s.neighbor.root,
},
),
},
)
return cbor.Encode(&tmpData)

case ProofStepTypeLeaf:
tmpData := cbor.NewConstructor(
2,
cbor.IndefLengthList{
s.prefixLength,
nibblesToBytes(s.neighbor.key),
s.neighbor.value,
},
)
return cbor.Encode(&tmpData)

default:
return nil, errors.New("unknown proof step type")
}
}

type ProofStepNeighbor struct {
key []Nibble
value Hash
prefix []Nibble
nibble Nibble
root Hash
}

func merkleProof(nodes []Node, myIdx int) []Hash {
var ret []Hash
var pivot = 8
var n = 8
for n >= 1 {
if myIdx < pivot {
ret = append(
ret,
merkleRoot(
nodes[pivot:pivot+n],
),
)
pivot -= (n >> 1)
} else {
ret = append(
ret,
merkleRoot(
nodes[pivot-n:pivot],
),
)
pivot += (n >> 1)
}
n = n >> 1
}
return ret
}
Loading

0 comments on commit 6ea345d

Please sign in to comment.