From 192c4d3b67be1b9bef7ac3e6bea55018e6a6e7f9 Mon Sep 17 00:00:00 2001 From: harry <53987565+h5law@users.noreply.github.com> Date: Mon, 4 Sep 2023 22:08:35 +0400 Subject: [PATCH 1/4] feat: add json marshalling for smt proofs --- proofs.go | 53 +++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 45 insertions(+), 8 deletions(-) diff --git a/proofs.go b/proofs.go index 78bbb1f..1f02764 100644 --- a/proofs.go +++ b/proofs.go @@ -3,6 +3,7 @@ package smt import ( "bytes" "encoding/binary" + "encoding/json" "errors" "math" ) @@ -13,16 +14,34 @@ var ErrBadProof = errors.New("bad proof") // SparseMerkleProof is a Merkle proof for an element in a SparseMerkleTree. type SparseMerkleProof struct { // SideNodes is an array of the sibling nodes leading up to the leaf of the proof. - SideNodes [][]byte + SideNodes [][]byte `json:"side_nodes"` // NonMembershipLeafData is the data of the unrelated leaf at the position // of the key being proven, in the case of a non-membership proof. For // membership proofs, is nil. - NonMembershipLeafData []byte + NonMembershipLeafData []byte `json:"non_membership_leaf_data"` // SiblingData is the data of the sibling node to the leaf being proven, // required for updatable proofs. For unupdatable proofs, is nil. - SiblingData []byte + SiblingData []byte `json:"sibling_data"` +} + +// MarshalJSON implements the json.Marshaler interface for SparseMerkleProof +func (proof *SparseMerkleProof) MarshalJSON() ([]byte, error) { + buf := bytes.NewBuffer(nil) + enc := json.NewEncoder(buf) + err := enc.Encode(proof) + if err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +// UnmarshalJSON implements the json.Unmarshaler interface for SparseMerkleProof. +func (proof *SparseMerkleProof) UnmarshalJSON(bz []byte) error { + buf := bytes.NewBuffer(bz) + dec := json.NewDecoder(buf) + return dec.Decode(proof) } func (proof *SparseMerkleProof) sanityCheck(spec *TreeSpec) bool { @@ -57,25 +76,43 @@ func (proof *SparseMerkleProof) sanityCheck(spec *TreeSpec) bool { // SparseCompactMerkleProof is a compact Merkle proof for an element in a SparseMerkleTree. type SparseCompactMerkleProof struct { // SideNodes is an array of the sibling nodes leading up to the leaf of the proof. - SideNodes [][]byte + SideNodes [][]byte `json:"side_nodes"` // NonMembershipLeafData is the data of the unrelated leaf at the position // of the key being proven, in the case of a non-membership proof. For // membership proofs, is nil. - NonMembershipLeafData []byte + NonMembershipLeafData []byte `json:"non_membership_leaf_data"` // BitMask, in the case of a compact proof, is a bit mask of the sidenodes // of the proof where an on-bit indicates that the sidenode at the bit's // index is a placeholder. This is only set if the proof is compact. - BitMask []byte + BitMask []byte `json:"bit_mask"` // NumSideNodes, in the case of a compact proof, indicates the number of // sidenodes in the proof when decompacted. This is only set if the proof is compact. - NumSideNodes int + NumSideNodes int `json:"num_side_nodes"` // SiblingData is the data of the sibling node to the leaf being proven, // required for updatable proofs. For unupdatable proofs, is nil. - SiblingData []byte + SiblingData []byte `json:"sibling_data"` +} + +// MarshalJSON implements the json.Marshaler interface for SparseCompactMerkleProof +func (proof *SparseCompactMerkleProof) MarshalJSON() ([]byte, error) { + buf := bytes.NewBuffer(nil) + enc := json.NewEncoder(buf) + err := enc.Encode(proof) + if err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +// UnmarshalJSON implements the json.Unmarshaler interface for SparseCompactMerkleProof +func (proof *SparseCompactMerkleProof) UnmarshalJSON(bz []byte) error { + buf := bytes.NewBuffer(bz) + dec := json.NewDecoder(buf) + return dec.Decode(proof) } func (proof *SparseCompactMerkleProof) sanityCheck(spec *TreeSpec) bool { From 50741577befa578652af6fb2bbc97626057c501a Mon Sep 17 00:00:00 2001 From: harry <53987565+h5law@users.noreply.github.com> Date: Mon, 4 Sep 2023 22:23:40 +0400 Subject: [PATCH 2/4] feat: add unit tests --- proofs.go | 16 ++--- proofs_test.go | 159 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 167 insertions(+), 8 deletions(-) diff --git a/proofs.go b/proofs.go index 1f02764..c741f0a 100644 --- a/proofs.go +++ b/proofs.go @@ -26,8 +26,8 @@ type SparseMerkleProof struct { SiblingData []byte `json:"sibling_data"` } -// MarshalJSON implements the json.Marshaler interface for SparseMerkleProof -func (proof *SparseMerkleProof) MarshalJSON() ([]byte, error) { +// Marshal serialises the SparseMerkleProof to bytes +func (proof *SparseMerkleProof) Marshal() ([]byte, error) { buf := bytes.NewBuffer(nil) enc := json.NewEncoder(buf) err := enc.Encode(proof) @@ -37,8 +37,8 @@ func (proof *SparseMerkleProof) MarshalJSON() ([]byte, error) { return buf.Bytes(), nil } -// UnmarshalJSON implements the json.Unmarshaler interface for SparseMerkleProof. -func (proof *SparseMerkleProof) UnmarshalJSON(bz []byte) error { +// Unmarshal deserialises the SparseMerkleProof from bytes +func (proof *SparseMerkleProof) Unmarshal(bz []byte) error { buf := bytes.NewBuffer(bz) dec := json.NewDecoder(buf) return dec.Decode(proof) @@ -97,8 +97,8 @@ type SparseCompactMerkleProof struct { SiblingData []byte `json:"sibling_data"` } -// MarshalJSON implements the json.Marshaler interface for SparseCompactMerkleProof -func (proof *SparseCompactMerkleProof) MarshalJSON() ([]byte, error) { +// Marshal serialises the SparseCompactMerkleProof to bytes +func (proof *SparseCompactMerkleProof) Marshal() ([]byte, error) { buf := bytes.NewBuffer(nil) enc := json.NewEncoder(buf) err := enc.Encode(proof) @@ -108,8 +108,8 @@ func (proof *SparseCompactMerkleProof) MarshalJSON() ([]byte, error) { return buf.Bytes(), nil } -// UnmarshalJSON implements the json.Unmarshaler interface for SparseCompactMerkleProof -func (proof *SparseCompactMerkleProof) UnmarshalJSON(bz []byte) error { +// Unmarshal deserialises the SparseCompactMerkleProof from bytes +func (proof *SparseCompactMerkleProof) Unmarshal(bz []byte) error { buf := bytes.NewBuffer(bz) dec := json.NewDecoder(buf) return dec.Decode(proof) diff --git a/proofs_test.go b/proofs_test.go index 5451ab3..3f4ed30 100644 --- a/proofs_test.go +++ b/proofs_test.go @@ -3,9 +3,168 @@ package smt import ( "bytes" "crypto/rand" + "crypto/sha256" "testing" + + "github.com/stretchr/testify/require" ) +func TestSparseMerkleProof_Marshal(t *testing.T) { + tree := setupTree(t) + + proof, err := tree.Prove([]byte("key")) + require.NoError(t, err) + bz, err := proof.Marshal() + require.NoError(t, err) + require.NotNil(t, bz) + require.Greater(t, len(bz), 0) + + proof2, err := tree.Prove([]byte("key2")) + require.NoError(t, err) + bz2, err := proof2.Marshal() + require.NoError(t, err) + require.NotNil(t, bz2) + require.Greater(t, len(bz2), 0) + require.NotEqual(t, bz, bz2) + + proof3 := randomiseProof(proof) + bz3, err := proof3.Marshal() + require.NoError(t, err) + require.NotNil(t, bz3) + require.Greater(t, len(bz3), 0) + require.NotEqual(t, bz, bz3) +} + +func TestSparseMerkleProof_Unmarshal(t *testing.T) { + tree := setupTree(t) + + proof, err := tree.Prove([]byte("key")) + require.NoError(t, err) + bz, err := proof.Marshal() + require.NoError(t, err) + require.NotNil(t, bz) + require.Greater(t, len(bz), 0) + uproof := new(SparseMerkleProof) + require.NoError(t, uproof.Unmarshal(bz)) + require.Equal(t, proof, uproof) + + proof2, err := tree.Prove([]byte("key2")) + require.NoError(t, err) + bz2, err := proof2.Marshal() + require.NoError(t, err) + require.NotNil(t, bz2) + require.Greater(t, len(bz2), 0) + uproof2 := new(SparseMerkleProof) + require.NoError(t, uproof2.Unmarshal(bz2)) + require.Equal(t, proof2, uproof2) + + proof3 := randomiseProof(proof) + bz3, err := proof3.Marshal() + require.NoError(t, err) + require.NotNil(t, bz3) + require.Greater(t, len(bz3), 0) + uproof3 := new(SparseMerkleProof) + require.NoError(t, uproof3.Unmarshal(bz3)) + require.Equal(t, proof3, uproof3) +} + +func TestSparseCompactMerkletProof_Marshal(t *testing.T) { + tree := setupTree(t) + + proof, err := tree.Prove([]byte("key")) + require.NoError(t, err) + compactProof, err := CompactProof(proof, tree.Spec()) + require.NoError(t, err) + bz, err := compactProof.Marshal() + require.NoError(t, err) + require.NotNil(t, bz) + require.Greater(t, len(bz), 0) + + proof2, err := tree.Prove([]byte("key2")) + require.NoError(t, err) + compactProof2, err := CompactProof(proof2, tree.Spec()) + require.NoError(t, err) + bz2, err := compactProof2.Marshal() + require.NoError(t, err) + require.NotNil(t, bz2) + require.Greater(t, len(bz2), 0) + require.NotEqual(t, bz, bz2) + + proof3 := randomiseProof(proof) + compactProof3, err := CompactProof(proof3, tree.Spec()) + require.NoError(t, err) + bz3, err := compactProof3.Marshal() + require.NoError(t, err) + require.NotNil(t, bz3) + require.Greater(t, len(bz3), 0) + require.NotEqual(t, bz, bz3) +} + +func TestSparseCompactMerkleProof_Unmarshal(t *testing.T) { + tree := setupTree(t) + + proof, err := tree.Prove([]byte("key")) + require.NoError(t, err) + compactProof, err := CompactProof(proof, tree.Spec()) + require.NoError(t, err) + bz, err := compactProof.Marshal() + require.NoError(t, err) + require.NotNil(t, bz) + require.Greater(t, len(bz), 0) + uCproof := new(SparseCompactMerkleProof) + require.NoError(t, uCproof.Unmarshal(bz)) + require.Equal(t, compactProof, uCproof) + uproof, err := DecompactProof(uCproof, tree.Spec()) + require.NoError(t, err) + require.Equal(t, proof, uproof) + + proof2, err := tree.Prove([]byte("key2")) + require.NoError(t, err) + compactProof2, err := CompactProof(proof2, tree.Spec()) + require.NoError(t, err) + bz2, err := compactProof2.Marshal() + require.NoError(t, err) + require.NotNil(t, bz2) + require.Greater(t, len(bz2), 0) + uCproof2 := new(SparseCompactMerkleProof) + require.NoError(t, uCproof2.Unmarshal(bz2)) + require.Equal(t, compactProof2, uCproof2) + uproof2, err := DecompactProof(uCproof2, tree.Spec()) + require.NoError(t, err) + require.Equal(t, proof2, uproof2) + + proof3 := randomiseProof(proof) + compactProof3, err := CompactProof(proof3, tree.Spec()) + require.NoError(t, err) + bz3, err := compactProof3.Marshal() + require.NoError(t, err) + require.NotNil(t, bz3) + require.Greater(t, len(bz3), 0) + uCproof3 := new(SparseCompactMerkleProof) + require.NoError(t, uCproof3.Unmarshal(bz3)) + require.Equal(t, compactProof3, uCproof3) + uproof3, err := DecompactProof(uCproof3, tree.Spec()) + require.NoError(t, err) + require.Equal(t, proof3, uproof3) +} + +func setupTree(t *testing.T) *SMT { + t.Helper() + + db, err := NewKVStore("") + require.NoError(t, err) + t.Cleanup(func() { + db.Stop() + }) + + tree := NewSparseMerkleTree(db, sha256.New()) + require.NoError(t, tree.Update([]byte("key"), []byte("value"))) + require.NoError(t, tree.Update([]byte("key2"), []byte("value2"))) + require.NoError(t, tree.Update([]byte("key3"), []byte("value3"))) + + return tree +} + func randomiseProof(proof *SparseMerkleProof) *SparseMerkleProof { sideNodes := make([][]byte, len(proof.SideNodes)) for i := range sideNodes { From ebef5401bbc689b5a600b32fb437ac1b7abe7e58 Mon Sep 17 00:00:00 2001 From: harry <53987565+h5law@users.noreply.github.com> Date: Tue, 5 Sep 2023 19:00:38 +0400 Subject: [PATCH 3/4] chore: address comments --- proofs.go | 6 ++---- proofs_test.go | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/proofs.go b/proofs.go index c741f0a..ae65a5e 100644 --- a/proofs.go +++ b/proofs.go @@ -30,8 +30,7 @@ type SparseMerkleProof struct { func (proof *SparseMerkleProof) Marshal() ([]byte, error) { buf := bytes.NewBuffer(nil) enc := json.NewEncoder(buf) - err := enc.Encode(proof) - if err != nil { + if err := enc.Encode(proof); err != nil { return nil, err } return buf.Bytes(), nil @@ -101,8 +100,7 @@ type SparseCompactMerkleProof struct { func (proof *SparseCompactMerkleProof) Marshal() ([]byte, error) { buf := bytes.NewBuffer(nil) enc := json.NewEncoder(buf) - err := enc.Encode(proof) - if err != nil { + if err := enc.Encode(proof); err != nil { return nil, err } return buf.Bytes(), nil diff --git a/proofs_test.go b/proofs_test.go index 3f4ed30..4ab9fd7 100644 --- a/proofs_test.go +++ b/proofs_test.go @@ -154,7 +154,7 @@ func setupTree(t *testing.T) *SMT { db, err := NewKVStore("") require.NoError(t, err) t.Cleanup(func() { - db.Stop() + require.NoError(t, db.Stop()) }) tree := NewSparseMerkleTree(db, sha256.New()) From ba147fb414e095ec06734e4b320a156b8b55690f Mon Sep 17 00:00:00 2001 From: harry <53987565+h5law@users.noreply.github.com> Date: Wed, 6 Sep 2023 19:51:58 +0400 Subject: [PATCH 4/4] feat: replace json with gob --- proofs.go | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/proofs.go b/proofs.go index ae65a5e..2874657 100644 --- a/proofs.go +++ b/proofs.go @@ -3,33 +3,38 @@ package smt import ( "bytes" "encoding/binary" - "encoding/json" + "encoding/gob" "errors" "math" ) +func init() { + gob.Register(SparseMerkleProof{}) + gob.Register(SparseCompactMerkleProof{}) +} + // ErrBadProof is returned when an invalid Merkle proof is supplied. var ErrBadProof = errors.New("bad proof") // SparseMerkleProof is a Merkle proof for an element in a SparseMerkleTree. type SparseMerkleProof struct { // SideNodes is an array of the sibling nodes leading up to the leaf of the proof. - SideNodes [][]byte `json:"side_nodes"` + SideNodes [][]byte // NonMembershipLeafData is the data of the unrelated leaf at the position // of the key being proven, in the case of a non-membership proof. For // membership proofs, is nil. - NonMembershipLeafData []byte `json:"non_membership_leaf_data"` + NonMembershipLeafData []byte // SiblingData is the data of the sibling node to the leaf being proven, // required for updatable proofs. For unupdatable proofs, is nil. - SiblingData []byte `json:"sibling_data"` + SiblingData []byte } // Marshal serialises the SparseMerkleProof to bytes func (proof *SparseMerkleProof) Marshal() ([]byte, error) { buf := bytes.NewBuffer(nil) - enc := json.NewEncoder(buf) + enc := gob.NewEncoder(buf) if err := enc.Encode(proof); err != nil { return nil, err } @@ -39,7 +44,7 @@ func (proof *SparseMerkleProof) Marshal() ([]byte, error) { // Unmarshal deserialises the SparseMerkleProof from bytes func (proof *SparseMerkleProof) Unmarshal(bz []byte) error { buf := bytes.NewBuffer(bz) - dec := json.NewDecoder(buf) + dec := gob.NewDecoder(buf) return dec.Decode(proof) } @@ -75,31 +80,31 @@ func (proof *SparseMerkleProof) sanityCheck(spec *TreeSpec) bool { // SparseCompactMerkleProof is a compact Merkle proof for an element in a SparseMerkleTree. type SparseCompactMerkleProof struct { // SideNodes is an array of the sibling nodes leading up to the leaf of the proof. - SideNodes [][]byte `json:"side_nodes"` + SideNodes [][]byte // NonMembershipLeafData is the data of the unrelated leaf at the position // of the key being proven, in the case of a non-membership proof. For // membership proofs, is nil. - NonMembershipLeafData []byte `json:"non_membership_leaf_data"` + NonMembershipLeafData []byte // BitMask, in the case of a compact proof, is a bit mask of the sidenodes // of the proof where an on-bit indicates that the sidenode at the bit's // index is a placeholder. This is only set if the proof is compact. - BitMask []byte `json:"bit_mask"` + BitMask []byte // NumSideNodes, in the case of a compact proof, indicates the number of // sidenodes in the proof when decompacted. This is only set if the proof is compact. - NumSideNodes int `json:"num_side_nodes"` + NumSideNodes int // SiblingData is the data of the sibling node to the leaf being proven, // required for updatable proofs. For unupdatable proofs, is nil. - SiblingData []byte `json:"sibling_data"` + SiblingData []byte } // Marshal serialises the SparseCompactMerkleProof to bytes func (proof *SparseCompactMerkleProof) Marshal() ([]byte, error) { buf := bytes.NewBuffer(nil) - enc := json.NewEncoder(buf) + enc := gob.NewEncoder(buf) if err := enc.Encode(proof); err != nil { return nil, err } @@ -109,7 +114,7 @@ func (proof *SparseCompactMerkleProof) Marshal() ([]byte, error) { // Unmarshal deserialises the SparseCompactMerkleProof from bytes func (proof *SparseCompactMerkleProof) Unmarshal(bz []byte) error { buf := bytes.NewBuffer(bz) - dec := json.NewDecoder(buf) + dec := gob.NewDecoder(buf) return dec.Decode(proof) }