Add Hasher.Clone and binary state marshaling

keccak.go

@@ -9,7 +9,14 @@ 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"
+	"slices"
+
+	"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
@@ -21,10 +28,148 @@ type KeccakState interface {
 
 const rate = 136 // sponge rate for Keccak-256: (1600 - 2*256) / 8
 
-var _ KeccakState = (*Hasher)(nil)
+var (
+	_ KeccakState                = (*Hasher)(nil)
+	_ encoding.BinaryMarshaler   = (*Hasher)(nil)
+	_ encoding.BinaryAppender    = (*Hasher)(nil)
+	_ encoding.BinaryUnmarshaler = (*Hasher)(nil)
+)
 
 // NewFastKeccak returns a new Keccak-256 hasher. The zero value of Hasher is
 // equally usable and avoids the allocation.
 func NewFastKeccak() *Hasher {
 	return &Hasher{}
 }
+
+// Canonical marshaled-state format, identical no matter which implementation
+// (native assembly or the x/crypto/sha3 fallback) produced it — the active
+// implementation is a runtime CPU-feature decision, so the encoding must not
+// depend on it:
+//
+//	magic (4) || state (200) || flags (1) || pos (1)
+//
+// state always has any partially absorbed input already XORed in; flags bit 0
+// is the sponge direction (0 = absorbing, 1 = squeezing); pos is the byte
+// position within the current block (absorb offset or squeeze read offset).
+const (
+	marshalMagic  = "fk1\x01"
+	marshaledSize = len(marshalMagic) + 200 + 1 + 1
+)
+
+var errInvalidState = errors.New("keccak: invalid hash state")
+
+// appendState appends the canonical encoding of a sponge state.
+func appendState(b []byte, state *[200]byte, squeezing bool, pos int) []byte {
+	b = slices.Grow(b, marshaledSize)
+	b = append(b, marshalMagic...)
+	b = append(b, state[:]...)
+	var flags byte
+	if squeezing {
+		flags = 1
+	}
+	return append(b, flags, byte(pos))
+}
+
+// parseState validates and decodes a canonical encoding.
+//
+// pos == rate is valid only while squeezing: the x/crypto implementation
+// permutes lazily on the next Read, so a state squeezed to an exact block
+// boundary persists with pos == rate. The native decoder normalizes it by
+// permuting eagerly.
+func parseState(data []byte) (state [200]byte, squeezing bool, pos int, err error) {
+	if len(data) != marshaledSize || string(data[:len(marshalMagic)]) != marshalMagic {
+		return state, false, 0, errInvalidState
+	}
+	payload := data[len(marshalMagic):]
+	flags, posByte := payload[200], int(payload[201])
+	squeezing = flags == 1
+	if flags > 1 || posByte > rate || (!squeezing && posByte == rate) {
+		return state, false, 0, errInvalidState
+	}
+	copy(state[:], payload[:200])
+	return state, squeezing, posByte, nil
+}
+
+// x/crypto/sha3's own stable marshal layout for legacy Keccak
+// (sha3/legacy_hash.go): magic || rate || state || n || direction.
+const (
+	xcMagic = "sha\x0b"
+	xcSize  = len(xcMagic) + 1 + 200 + 1 + 1
+)
+
+var errXCFormat = errors.New("keccak: unexpected x/crypto/sha3 state format")
+
+// xcAppendState re-encodes a wrapped x/crypto/sha3 hasher's state in the
+// canonical format. x/crypto keeps partial input XORed into the state with n
+// as the block offset, which matches the canonical layout directly.
+func xcAppendState(b []byte, xc KeccakState) ([]byte, error) {
+	// On error return b unchanged, per the encoding.BinaryAppender contract,
+	// so a caller's existing prefix is not dropped.
+	m, ok := xc.(encoding.BinaryMarshaler)
+	if !ok {
+		return b, errXCFormat
+	}
+	enc, err := m.MarshalBinary()
+	if err != nil {
+		return b, err
+	}
+	if len(enc) != xcSize || string(enc[:len(xcMagic)]) != xcMagic || int(enc[4]) != rate {
+		return b, errXCFormat
+	}
+	var state [200]byte
+	copy(state[:], enc[5:205])
+	return appendState(b, &state, enc[206] == 1, int(enc[205])), nil
+}
+
+// xcFromState builds a wrapped x/crypto/sha3 hasher holding the given
+// canonical sponge state.
+func xcFromState(state *[200]byte, squeezing bool, pos int) (KeccakState, error) {
+	// Mirror the layout xcAppendState decodes: magic || rate || state || n
+	// || direction. blob does not escape (x/crypto reads it and copies),
+	// so this is stack-resident.
+	var blob [xcSize]byte
+	copy(blob[:], xcMagic)
+	blob[4] = byte(rate)
+	copy(blob[5:205], state[:])
+	blob[205] = byte(pos)
+	if squeezing {
+		blob[206] = 1
+	}
+
+	st, ok := sha3.NewLegacyKeccak256().(KeccakState)
+	if !ok {
+		return nil, errXCFormat
+	}
+	u, ok := st.(encoding.BinaryUnmarshaler)
+	if !ok {
+		return nil, errXCFormat
+	}
+	if err := u.UnmarshalBinary(blob[:]); err != nil {
+		return nil, err
+	}
+	return st, nil
+}
+
+// cloneXC deep-copies a wrapped x/crypto/sha3 hasher via a marshal round-trip.
+func cloneXC(xc KeccakState) (KeccakState, error) {
+	m, ok := xc.(encoding.BinaryMarshaler)
+	if !ok {
+		return nil, errXCFormat
+	}
+	enc, err := m.MarshalBinary()
+	if err != nil {
+		return nil, err
+	}
+	st, ok := sha3.NewLegacyKeccak256().(KeccakState)
+	if !ok {
+		return nil, errXCFormat
+	}
+	u, ok := st.(encoding.BinaryUnmarshaler)
+	if !ok {
+		return nil, errXCFormat
+	}
+	if err := u.UnmarshalBinary(enc); err != nil {
+		return nil, err
+	}
+	return st, nil
+}

keccak_asm.go

@@ -111,6 +111,7 @@ func (s *sponge) padAndSqueeze() {
 	s.state[s.absorbed] ^= 0x01
 	s.state[rate-1] ^= 0x80
 	keccakF1600(&s.state)
+	s.absorbed = 0
 	s.squeezing = true
 	s.readIdx = 0
 }
@@ -221,6 +222,82 @@ func (h *Hasher) Read(out []byte) (int, error) {
 	return h.sponge.Read(out)
 }
 
+// Clone returns a copy of the hasher in its current state. The copy and the
+// original evolve independently.
+func (h *Hasher) Clone() (*Hasher, error) {
+	if useASM {
+		return &Hasher{sponge: h.sponge}, nil
+	}
+	if h.xc == nil {
+		return &Hasher{}, nil
+	}
+	xc, err := cloneXC(h.xc)
+	if err != nil {
+		return nil, err
+	}
+	return &Hasher{xc: xc}, nil
+}
+
+// MarshalBinary implements encoding.BinaryMarshaler. The encoding is
+// canonical: the same logical state marshals to the same bytes regardless of
+// platform or active implementation, and can be restored anywhere.
+func (h *Hasher) MarshalBinary() ([]byte, error) {
+	return h.AppendBinary(make([]byte, 0, marshaledSize))
+}
+
+// AppendBinary implements encoding.BinaryAppender.
+func (h *Hasher) AppendBinary(b []byte) ([]byte, error) {
+	if !useASM {
+		if h.xc == nil {
+			// A zero-value hasher is the canonical zero state; avoid
+			// allocating (and retaining) a fallback state just to encode it.
+			var zero [200]byte
+			return appendState(b, &zero, false, 0), nil
+		}
+		return xcAppendState(b, h.xc)
+	}
+	s := &h.sponge
+	if s.squeezing {
+		return appendState(b, &s.state, true, s.readIdx), nil
+	}
+	// Canonical form keeps partial input XORed into the state.
+	state := s.state
+	xorIn(&state, s.buf[:s.absorbed])
+	return appendState(b, &state, false, s.absorbed), nil
+}
+
+// UnmarshalBinary implements encoding.BinaryUnmarshaler.
+func (h *Hasher) UnmarshalBinary(data []byte) error {
+	state, squeezing, pos, err := parseState(data)
+	if err != nil {
+		return err
+	}
+	if !useASM {
+		xc, err := xcFromState(&state, squeezing, pos)
+		if err != nil {
+			return err
+		}
+		h.xc = xc
+		return nil
+	}
+	// Partial input is already XORed into state, so buf stays zero: absorbing
+	// XORs buf into state at finalization, and XORing zeroes is a no-op.
+	s := sponge{state: state}
+	if squeezing {
+		if pos == rate {
+			// Block-boundary state from the lazy fallback implementation;
+			// the native sponge permutes eagerly.
+			keccakF1600(&s.state)
+			pos = 0
+		}
+		s.squeezing, s.readIdx = true, pos
+	} else {
+		s.absorbed = pos
+	}
+	h.sponge = s
+	return 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_default.go

@@ -69,3 +69,48 @@ func (h *Hasher) Read(out []byte) (int, error) {
 	h.init()
 	return h.h.Read(out)
 }
+
+// Clone returns a copy of the hasher in its current state. The copy and the
+// original evolve independently.
+func (h *Hasher) Clone() (*Hasher, error) {
+	if h.h == nil {
+		return &Hasher{}, nil
+	}
+	c, err := cloneXC(h.h)
+	if err != nil {
+		return nil, err
+	}
+	return &Hasher{h: c}, nil
+}
+
+// MarshalBinary implements encoding.BinaryMarshaler. The encoding is
+// canonical: the same logical state marshals to the same bytes regardless of
+// platform or active implementation, and can be restored anywhere.
+func (h *Hasher) MarshalBinary() ([]byte, error) {
+	return h.AppendBinary(make([]byte, 0, marshaledSize))
+}
+
+// AppendBinary implements encoding.BinaryAppender.
+func (h *Hasher) AppendBinary(b []byte) ([]byte, error) {
+	if h.h == nil {
+		// A zero-value hasher is the canonical zero state; avoid allocating
+		// (and retaining) a backing state just to encode it.
+		var zero [200]byte
+		return appendState(b, &zero, false, 0), nil
+	}
+	return xcAppendState(b, h.h)
+}
+
+// UnmarshalBinary implements encoding.BinaryUnmarshaler.
+func (h *Hasher) UnmarshalBinary(data []byte) error {
+	state, squeezing, pos, err := parseState(data)
+	if err != nil {
+		return err
+	}
+	st, err := xcFromState(&state, squeezing, pos)
+	if err != nil {
+		return err
+	}
+	h.h = st
+	return nil
+}