From 158643ca500f1ad23b4fd6b19747ee6257736c7a Mon Sep 17 00:00:00 2001 From: gpsanant Date: Wed, 28 Feb 2024 22:12:37 -0800 Subject: [PATCH 1/3] add GenerateProofWithIndex --- merkletree.go | 12 ++++++++++++ proof_test.go | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/merkletree.go b/merkletree.go index 7ee6d78..88e9f81 100644 --- a/merkletree.go +++ b/merkletree.go @@ -108,6 +108,18 @@ func (t *MerkleTree) GenerateProof(data []byte, height int) (*Proof, error) { return nil, err } + return t.GenerateProofWithIndex(index, height) +} + +// GenerateProofWithIndex generates the proof for the data at the given index. It is faster than GenerateProof() if the index is already known. +// Height is the height of the pollard to verify the proof. If using the Merkle root to verify this should be 0. +// If the index is out of range this will return an error. +// If the data is present in the tree this will return the hashes for each level in the tree and the index of the value in the tree. +func (t *MerkleTree) GenerateProofWithIndex(index uint64, height int) (*Proof, error) { + if index >= uint64(len(t.Data)) { + return nil, errors.New("index out of range") + } + proofLen := int(math.Ceil(math.Log2(float64(len(t.Data))))) - height hashes := make([][]byte, proofLen) diff --git a/proof_test.go b/proof_test.go index 8719247..2d72f29 100644 --- a/proof_test.go +++ b/proof_test.go @@ -21,6 +21,25 @@ import ( "github.com/stretchr/testify/assert" ) +func TestProofWithIndex(t *testing.T) { + for i, test := range tests { + if test.createErr == nil { + tree, err := NewTree( + WithData(test.data), + WithHashType(test.hashType), + ) + assert.Nil(t, err, fmt.Sprintf("failed to create tree at test %d", i)) + for j, data := range test.data { + proof, err := tree.GenerateProofWithIndex(uint64(j), 0) + assert.Nil(t, err, fmt.Sprintf("failed to create proof at test %d data %d", i, j)) + proven, err := VerifyProofUsing(data, false, proof, [][]byte{tree.Root()}, test.hashType) + assert.Nil(t, err, fmt.Sprintf("error verifying proof at test %d", i)) + assert.True(t, proven, fmt.Sprintf("failed to verify proof at test %d data %d", i, j)) + } + } + } +} + func TestProof(t *testing.T) { for i, test := range tests { if test.createErr == nil { @@ -100,6 +119,20 @@ func TestMissingProof(t *testing.T) { } } +func TestProveInvalidIndex(t *testing.T) { + for i, test := range tests { + if test.createErr == nil { + tree, err := NewTree( + WithData(test.data), + WithHashType(test.hashType), + ) + assert.Nil(t, err, fmt.Sprintf("failed to create tree at test %d", i)) + _, err = tree.GenerateProofWithIndex(uint64(len(test.data)+i), 0) + assert.Equal(t, err.Error(), "index out of range") + } + } +} + func TestBadProof(t *testing.T) { for i, test := range tests { if test.createErr == nil && len(test.data) > 1 { From afa030a7c76d15282f7c87235344fe7d5721dd67 Mon Sep 17 00:00:00 2001 From: gpsanant Date: Wed, 28 Feb 2024 22:22:48 -0800 Subject: [PATCH 2/3] add GenerateMultiProofWithIndices --- merkletree.go | 22 ++++++++++++++++++---- multiproof_test.go | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 4 deletions(-) diff --git a/merkletree.go b/merkletree.go index 88e9f81..aaa4afa 100644 --- a/merkletree.go +++ b/merkletree.go @@ -135,17 +135,31 @@ func (t *MerkleTree) GenerateProofWithIndex(index uint64, height int) (*Proof, e // GenerateMultiProof generates the proof for multiple pieces of data. func (t *MerkleTree) GenerateMultiProof(data [][]byte) (*MultiProof, error) { - hashes := make([][][]byte, len(data)) indices := make([]uint64, len(data)) // Step 1: generate individual proofs. for dataIndex := range data { - tmpProof, err := t.GenerateProof(data[dataIndex], 0) + index, err := t.indexOf(data[dataIndex]) if err != nil { return nil, err } - hashes[dataIndex] = tmpProof.Hashes - indices[dataIndex] = tmpProof.Index + indices[dataIndex] = index + } + + return t.GenerateMultiProofWithIndices(indices) +} + +// GenerateMultiProof generates the proof for multiple pieces of data. +func (t *MerkleTree) GenerateMultiProofWithIndices(indices []uint64) (*MultiProof, error) { + hashes := make([][][]byte, len(indices)) + + // Step 1: generate individual proofs. + for i, leafIndex := range indices { + tmpProof, err := t.GenerateProofWithIndex(leafIndex, 0) + if err != nil { + return nil, err + } + hashes[i] = tmpProof.Hashes } // Step 2: combine the hashes across all proofs and highlight all calculated indices. diff --git a/multiproof_test.go b/multiproof_test.go index 4b295b3..5b32e07 100644 --- a/multiproof_test.go +++ b/multiproof_test.go @@ -24,6 +24,39 @@ import ( "github.com/stretchr/testify/require" ) +func TestMultiProofWithIndices(t *testing.T) { + for i, test := range tests { + if test.createErr == nil { + tree, err := NewTree( + WithData(test.data), + WithHashType(test.hashType), + WithSalt(test.salt), + WithSorted(test.sorted), + ) + assert.Nil(t, err, fmt.Sprintf("failed to create tree at test %d", i)) + + // Test proof for all combinations of data. + var proof *MultiProof + combinations := 1<>k)&1 == 1 { + indices = append(indices, uint64(k)) + items = append(items, test.data[k]) + } + } + proof, err = tree.GenerateMultiProofWithIndices(indices) + assert.Nil(t, err, fmt.Sprintf("failed to create multiproof at test %d data %d", i, j)) + proven, err := proof.Verify(items, tree.Root()) + assert.Nil(t, err, fmt.Sprintf("error verifying multiproof at test %d data %d", i, j)) + assert.True(t, proven, fmt.Sprintf("failed to verify multiproof at test %d data %d", i, j)) + } + } + } +} + func TestMultiProof(t *testing.T) { for i, test := range tests { if test.createErr == nil { From f0686ad2dffb402cb7dad1ecacd38b5ec986553d Mon Sep 17 00:00:00 2001 From: gpsanant Date: Sat, 2 Mar 2024 14:21:46 -0800 Subject: [PATCH 3/3] fix lint --- merkletree.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/merkletree.go b/merkletree.go index aaa4afa..0e0c6c2 100644 --- a/merkletree.go +++ b/merkletree.go @@ -111,7 +111,8 @@ func (t *MerkleTree) GenerateProof(data []byte, height int) (*Proof, error) { return t.GenerateProofWithIndex(index, height) } -// GenerateProofWithIndex generates the proof for the data at the given index. It is faster than GenerateProof() if the index is already known. +// GenerateProofWithIndex generates the proof for the data at the given index. +// It is faster than GenerateProof() if the index is already known. // Height is the height of the pollard to verify the proof. If using the Merkle root to verify this should be 0. // If the index is out of range this will return an error. // If the data is present in the tree this will return the hashes for each level in the tree and the index of the value in the tree.