From 19b21432a13a855707124af03eb36a5f7d6936f6 Mon Sep 17 00:00:00 2001
From: Alexis Sellier <alexis@cloudhead.io>
Date: Mon, 4 Dec 2017 13:12:08 +0100
Subject: [PATCH] Fix key proof serialization

---
 proof_key.go      | 38 ++++++++++++++++++++++++++++++++------
 proof_key_test.go | 16 ++++++++--------
 testutils_test.go |  4 ++--
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/proof_key.go b/proof_key.go
index 4588fc209..cf7531d3e 100644
--- a/proof_key.go
+++ b/proof_key.go
@@ -22,6 +22,12 @@ type KeyProof interface {
 	Bytes() []byte
 }
 
+const (
+	// Used for serialization of proofs.
+	keyExistsMagicNumber = 0x50
+	keyAbsentMagicNumber = 0x51
+)
+
 // KeyExistsProof represents a proof of existence of a single key.
 type KeyExistsProof struct {
 	RootHash data.Bytes `json:"root_hash"`
@@ -47,11 +53,11 @@ func (proof *KeyExistsProof) Verify(key []byte, value []byte, root []byte) error
 
 // Bytes returns a go-wire binary serialization
 func (proof *KeyExistsProof) Bytes() []byte {
-	return wire.BinaryBytes(proof)
+	return append([]byte{keyExistsMagicNumber}, wire.BinaryBytes(proof)...)
 }
 
-// ReadKeyExistsProof will deserialize a KeyExistsProof from bytes.
-func ReadKeyExistsProof(data []byte) (*KeyExistsProof, error) {
+// readKeyExistsProof will deserialize a KeyExistsProof from bytes.
+func readKeyExistsProof(data []byte) (*KeyExistsProof, error) {
 	proof := new(KeyExistsProof)
 	err := wire.ReadBinaryBytes(data, &proof)
 	return proof, err
@@ -94,12 +100,32 @@ func (proof *KeyAbsentProof) Verify(key, value []byte, root []byte) error {
 
 // Bytes returns a go-wire binary serialization
 func (proof *KeyAbsentProof) Bytes() []byte {
-	return wire.BinaryBytes(proof)
+	return append([]byte{keyAbsentMagicNumber}, wire.BinaryBytes(proof)...)
 }
 
-// ReadKeyAbsentProof will deserialize a KeyAbsentProof from bytes.
-func ReadKeyAbsentProof(data []byte) (*KeyAbsentProof, error) {
+// readKeyAbsentProof will deserialize a KeyAbsentProof from bytes.
+func readKeyAbsentProof(data []byte) (*KeyAbsentProof, error) {
 	proof := new(KeyAbsentProof)
 	err := wire.ReadBinaryBytes(data, &proof)
 	return proof, err
 }
+
+func ReadKeyProof(data []byte) (KeyProof, error) {
+	var err error
+	var n int
+
+	buf := bytes.NewBuffer(data)
+
+	b := wire.ReadByte(buf, &n, &err)
+	if err != nil {
+		return nil, err
+	}
+
+	switch b {
+	case keyExistsMagicNumber:
+		return readKeyExistsProof(buf.Bytes())
+	case keyAbsentMagicNumber:
+		return readKeyAbsentProof(buf.Bytes())
+	}
+	return nil, errors.New("unrecognized proof")
+}
diff --git a/proof_key_test.go b/proof_key_test.go
index 027add234..953c4d4f1 100644
--- a/proof_key_test.go
+++ b/proof_key_test.go
@@ -23,12 +23,15 @@ func TestSerializeProofs(t *testing.T) {
 	val, proof, err := tree.GetWithProof(key)
 	require.Nil(err, "%+v", err)
 	require.NotNil(val)
+
 	bin := proof.Bytes()
-	eproof, err := ReadKeyExistsProof(bin)
+	proof2, err := ReadKeyProof(bin)
 	require.Nil(err, "%+v", err)
-	require.NoError(eproof.Verify(key, val, root))
-	_, err = ReadKeyAbsentProof(bin)
-	require.NotNil(err)
+	require.NoError(proof2.Verify(key, val, root))
+
+	if _, ok := proof2.(*KeyExistsProof); !ok {
+		require.FailNow("Proof should be *KeyExistsProof")
+	}
 
 	// test with key absent
 	key = []byte{0x38}
@@ -36,10 +39,7 @@ func TestSerializeProofs(t *testing.T) {
 	require.Nil(err, "%+v", err)
 	require.Nil(val)
 	bin = proof.Bytes()
-	// I think this is ugly it works this way, but without type-bytes nothing we can do :(
-	// eproof, err = ReadKeyExistsProof(bin)
-	// require.NotNil(err)
-	aproof, err := ReadKeyAbsentProof(bin)
+	aproof, err := ReadKeyProof(bin)
 	require.Nil(err, "%+v", err)
 	require.NoError(aproof.Verify(key, val, root))
 }
diff --git a/testutils_test.go b/testutils_test.go
index 9c4340063..d9d01baca 100644
--- a/testutils_test.go
+++ b/testutils_test.go
@@ -125,14 +125,14 @@ func testProof(t *testing.T, proof *KeyExistsProof, keyBytes, valueBytes, rootHa
 
 	// Write/Read then verify.
 	proofBytes := proof.Bytes()
-	proof2, err := ReadKeyExistsProof(proofBytes)
+	proof2, err := ReadKeyProof(proofBytes)
 	require.Nil(t, err, "Failed to read KeyExistsProof from bytes: %v", err)
 	require.NoError(t, proof2.Verify(keyBytes, valueBytes, proof.RootHash))
 
 	// Random mutations must not verify
 	for i := 0; i < 10; i++ {
 		badProofBytes := MutateByteSlice(proofBytes)
-		badProof, err := ReadKeyExistsProof(badProofBytes)
+		badProof, err := ReadKeyProof(badProofBytes)
 		// may be invalid... errors are okay
 		if err == nil {
 			assert.Error(t, badProof.Verify(keyBytes, valueBytes, rootHashBytes),