From 2b1401d62aaaf7e74d0fd9b0f787dc79bef58200 Mon Sep 17 00:00:00 2001
From: Evgenii Baidakov <evgenii@nspcc.io>
Date: Wed, 28 Feb 2024 10:51:13 +0400
Subject: [PATCH] tz: Add BinaryMarshaler/BinaryUnmarshaler implementation to
 hash

Closes #47.

Signed-off-by: Evgenii Baidakov <evgenii@nspcc.io>
---
 tz/digest.go      | 44 ++++++++++++++++++++++++++++
 tz/digest_test.go | 74 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 118 insertions(+)
 create mode 100644 tz/digest_test.go

diff --git a/tz/digest.go b/tz/digest.go
index 5d7066a..844f5f7 100644
--- a/tz/digest.go
+++ b/tz/digest.go
@@ -1,6 +1,8 @@
 package tz
 
 import (
+	"errors"
+	"fmt"
 	"hash"
 
 	"github.com/nspcc-dev/tzhash/gf127"
@@ -22,6 +24,8 @@ type digest struct {
 }
 
 // New returns a new [hash.Hash] computing the Tillich-Zémor checksum.
+// The Hash also implements [encoding.BinaryMarshaler] and [encoding.BinaryUnmarshaler]
+// to marshal and unmarshal the internal state of the hash.
 func New() hash.Hash {
 	d := new(digest)
 	d.Reset()
@@ -122,3 +126,43 @@ func mulBitRightGeneric(c00, c10, c01, c11 *GF127, bit bool, tmp *GF127) {
 		*c11 = *tmp
 	}
 }
+
+// MarshalBinary implements [encoding.BinaryMarshaler].
+func (d *digest) MarshalBinary() ([]byte, error) {
+	var (
+		b = make([]byte, 0, Size)
+	)
+
+	for _, a := range d.x {
+		state, err := a.MarshalBinary()
+		if err != nil {
+			return nil, err
+		}
+
+		b = append(b, state...)
+	}
+
+	return b, nil
+}
+
+// UnmarshalBinary implements [encoding.BinaryUnmarshaler].
+func (d *digest) UnmarshalBinary(b []byte) error {
+	if len(b) != Size {
+		return errors.New("tz: invalid hash state size")
+	}
+
+	var (
+		start, end int
+	)
+
+	for i := 0; i < 4; i++ {
+		start = gf127.Size * i
+		end = start + gf127.Size
+
+		if err := d.x[i].UnmarshalBinary(b[start:end]); err != nil {
+			return fmt.Errorf("gf127 unmarshal: %w", err)
+		}
+	}
+
+	return nil
+}
diff --git a/tz/digest_test.go b/tz/digest_test.go
new file mode 100644
index 0000000..45a3b65
--- /dev/null
+++ b/tz/digest_test.go
@@ -0,0 +1,74 @@
+package tz
+
+import (
+	"encoding/binary"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func newDigest() *digest {
+	d := &digest{}
+	d.Reset()
+	return d
+}
+
+func Test_digest_marshaling(t *testing.T) {
+	var (
+		d              = newDigest()
+		marshaledState []byte
+		hashSum        []byte
+	)
+
+	for i := byte(0); i < 10; i++ {
+		n, err := d.Write([]byte{i})
+		require.NoError(t, err)
+		require.Equal(t, 1, n)
+
+		t.Run("marshal", func(t *testing.T) {
+			marshaledState, err = d.MarshalBinary()
+
+			require.NoError(t, err)
+			require.Len(t, marshaledState, Size)
+
+			hashSum = d.Sum(nil)
+			require.Len(t, hashSum, Size)
+		})
+
+		t.Run("unmarshal", func(t *testing.T) {
+			unmarshalDigest := newDigest()
+			err = unmarshalDigest.UnmarshalBinary(marshaledState)
+			require.NoError(t, err)
+
+			unmarshalDigestHash := unmarshalDigest.Sum(nil)
+			require.Len(t, unmarshalDigestHash, Size)
+
+			require.Equal(t, hashSum, unmarshalDigestHash)
+		})
+	}
+
+	t.Run("invalid length", func(t *testing.T) {
+		unmarshalDigest := newDigest()
+		state := []byte{1, 2, 3}
+
+		err := unmarshalDigest.UnmarshalBinary(state)
+		require.Error(t, err)
+	})
+
+	t.Run("invalid state data", func(t *testing.T) {
+		someDigest := newDigest()
+		_, err := d.Write([]byte{1, 2, 3})
+		require.NoError(t, err)
+
+		state, err := someDigest.MarshalBinary()
+		require.NoError(t, err)
+
+		a := uint64(1) << 63
+		// broke state data.
+		binary.BigEndian.PutUint64(state[0:8], a)
+		unmarshalDigest := newDigest()
+
+		err = unmarshalDigest.UnmarshalBinary(state)
+		require.Error(t, err)
+	})
+}