keccak: add independent hasher cloning

keccak.go

@@ -9,7 +9,13 @@ package keccak
 //go:generate go run gen_keccakf_bmi2.go
 //go:generate go run github.com/klauspost/asmfmt/cmd/asmfmt@v1.3.2 -w keccakf_amd64_bmi2.s
 
-import "hash"
+import (
+	"encoding"
+	"errors"
+	"hash"
+
+	"golang.org/x/crypto/sha3"
+)
 
 // KeccakState wraps the keccak hasher. In addition to the usual hash methods, it also supports
 // Read to get a variable amount of data from the hash state. Read is faster than Sum
@@ -28,3 +34,29 @@ var _ KeccakState = (*Hasher)(nil)
 func NewFastKeccak() *Hasher {
 	return &Hasher{}
 }
+
+var errCloneUnsupported = errors.New("keccak: fallback hasher does not support cloning")
+
+func cloneXCrypto(state KeccakState) (KeccakState, error) {
+	marshaler, ok := state.(encoding.BinaryMarshaler)
+	if !ok {
+		return nil, errCloneUnsupported
+	}
+	encoded, err := marshaler.MarshalBinary()
+	if err != nil {
+		return nil, err
+	}
+
+	clone, ok := sha3.NewLegacyKeccak256().(KeccakState)
+	if !ok {
+		return nil, errCloneUnsupported
+	}
+	unmarshaler, ok := clone.(encoding.BinaryUnmarshaler)
+	if !ok {
+		return nil, errCloneUnsupported
+	}
+	if err := unmarshaler.UnmarshalBinary(encoded); err != nil {
+		return nil, err
+	}
+	return clone, nil
+}

keccak_asm.go

@@ -149,8 +149,9 @@ func sum256XCrypto(data []byte) [32]byte {
 	return out
 }
 
-// Hasher is a streaming Keccak-256 hasher.
-// Uses platform assembly when available, x/crypto/sha3 otherwise.
+// Hasher is a streaming Keccak-256 hasher. It uses platform assembly when
+// available and x/crypto/sha3 otherwise. A Hasher must not be copied after
+// first use; use Clone to create an independent copy.
 type Hasher struct {
 	sponge
 	xc KeccakState // x/crypto fallback
@@ -221,6 +222,20 @@ func (h *Hasher) Read(out []byte) (int, error) {
 	return h.sponge.Read(out)
 }
 
+// Clone returns an independent copy of h in its current state. It carries
+// both sub-states, so the copy dispatches the same way the original does for
+// any value of useASM.
+func (h *Hasher) Clone() (*Hasher, error) {
+	if h.xc == nil {
+		return &Hasher{sponge: h.sponge}, nil
+	}
+	state, err := cloneXCrypto(h.xc)
+	if err != nil {
+		return nil, err
+	}
+	return &Hasher{sponge: h.sponge, xc: state}, nil
+}
+
 // xorIn XORs data into the first len(data) bytes of state using uint64 loads.
 func xorIn(state *[200]byte, data []byte) {
 	for i := 0; i+8 <= len(data); i += 8 {

keccak_asm_test.go

@@ -5,6 +5,8 @@ package keccak
 import (
 	"math/rand"
 	"testing"
+
+	"golang.org/x/crypto/sha3"
 )
 
 // TestXorAndPermute verifies the fused XOR+permute assembly entry point
@@ -34,3 +36,42 @@ func TestXorAndPermute(t *testing.T) {
 		}
 	}
 }
+
+func xcDigest(state KeccakState) [32]byte {
+	var out [32]byte
+	state.Sum(out[:0])
+	return out
+}
+
+// TestHasherCloneFallbackState covers Clone on a receiver whose state lives in
+// the x/crypto fallback while the package dispatches to assembly. Building the
+// receiver directly keeps the test off the useASM global, which is only safe
+// to write while nothing else in the package runs concurrently.
+func TestHasherCloneFallbackState(t *testing.T) {
+	const prefix = "shared prefix"
+	var original Hasher
+	original.Write([]byte("sponge side"))
+	original.xc = sha3.NewLegacyKeccak256().(KeccakState)
+	original.xc.Write([]byte(prefix))
+
+	clone, err := original.Clone()
+	if err != nil {
+		t.Fatalf("Clone: %v", err)
+	}
+	if got, want := clone.Sum256(), original.Sum256(); got != want {
+		t.Fatalf("clone digest = %x, original = %x: the copy does not dispatch like the original", got, want)
+	}
+	if clone.sponge != original.sponge {
+		t.Fatal("Clone dropped the sponge state")
+	}
+
+	original.xc.Write([]byte(" original"))
+	clone.xc.Write([]byte(" clone"))
+
+	if got, want := xcDigest(original.xc), Sum256([]byte(prefix+" original")); got != want {
+		t.Fatalf("original digest = %x, want %x", got, want)
+	}
+	if got, want := xcDigest(clone.xc), Sum256([]byte(prefix+" clone")); got != want {
+		t.Fatalf("clone digest = %x, want %x", got, want)
+	}
+}

keccak_clone_test.go

@@ -0,0 +1,50 @@
+package keccak
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestHasherClone(t *testing.T) {
+	var zero Hasher
+	zeroClone, err := zero.Clone()
+	if err != nil {
+		t.Fatalf("Clone zero value: %v", err)
+	}
+	if got, want := zeroClone.Sum256(), Sum256(nil); got != want {
+		t.Fatalf("zero-value clone digest = %x, want %x", got, want)
+	}
+
+	prefix := []byte("shared prefix")
+	var original Hasher
+	original.Write(prefix)
+
+	clone, err := original.Clone()
+	if err != nil {
+		t.Fatalf("Clone: %v", err)
+	}
+	original.Write([]byte(" original"))
+	clone.Write([]byte(" clone"))
+
+	if got, want := original.Sum256(), Sum256(append(bytes.Clone(prefix), " original"...)); got != want {
+		t.Fatalf("original digest = %x, want %x", got, want)
+	}
+	if got, want := clone.Sum256(), Sum256(append(bytes.Clone(prefix), " clone"...)); got != want {
+		t.Fatalf("clone digest = %x, want %x", got, want)
+	}
+
+	var squeezing Hasher
+	squeezing.Write(prefix)
+	squeezing.Read(make([]byte, 17))
+	squeezingClone, err := squeezing.Clone()
+	if err != nil {
+		t.Fatalf("Clone while squeezing: %v", err)
+	}
+	want := make([]byte, 200)
+	got := make([]byte, len(want))
+	squeezing.Read(want)
+	squeezingClone.Read(got)
+	if !bytes.Equal(got, want) {
+		t.Fatal("clone did not preserve an independent squeeze position")
+	}
+}

keccak_default.go

@@ -16,7 +16,8 @@ func Sum256(data []byte) [32]byte {
 	return out
 }
 
-// Hasher is a streaming Keccak-256 hasher wrapping x/crypto/sha3.
+// Hasher is a streaming Keccak-256 hasher wrapping x/crypto/sha3. A Hasher
+// must not be copied after first use; use Clone to create an independent copy.
 type Hasher struct {
 	h KeccakState
 }
@@ -69,3 +70,15 @@ func (h *Hasher) Read(out []byte) (int, error) {
 	h.init()
 	return h.h.Read(out)
 }
+
+// Clone returns an independent copy of h in its current state.
+func (h *Hasher) Clone() (*Hasher, error) {
+	if h.h == nil {
+		return &Hasher{}, nil
+	}
+	state, err := cloneXCrypto(h.h)
+	if err != nil {
+		return nil, err
+	}
+	return &Hasher{h: state}, nil
+}