feat: add Marshal and Unmarshal methods to SM(S)T proof types

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"`
+}
+
+// Marshal serialises the SparseMerkleProof to bytes
+func (proof *SparseMerkleProof) Marshal() ([]byte, error) {
+	buf := bytes.NewBuffer(nil)
+	enc := json.NewEncoder(buf)
+	err := enc.Encode(proof)
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+// 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)
 }
 
 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"`
+}
+
+// Marshal serialises the SparseCompactMerkleProof to bytes
+func (proof *SparseCompactMerkleProof) Marshal() ([]byte, error) {
+	buf := bytes.NewBuffer(nil)
+	enc := json.NewEncoder(buf)
+	err := enc.Encode(proof)
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+// 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)
 }
 
 func (proof *SparseCompactMerkleProof) sanityCheck(spec *TreeSpec) bool {

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 {