Optimize PoW

registration/pow_verifier.go

@@ -68,7 +68,7 @@ func (v *powVerifier) Verify(poetChallenge, nodeID []byte, nonce uint64) error {
 		return fmt.Errorf("%w: invalid nodeID length: %d", ErrInvalidPow, len(nodeID))
 	}
 
-	hash := shared.CalcSubmitPowHash(v.params.Challenge, poetChallenge, nodeID, nil, nonce)
+	hash := shared.NewPowHasher(v.params.Challenge, nodeID, poetChallenge).Hash(nonce, nil)
 
 	if !shared.CheckLeadingZeroBits(hash, v.params.Difficulty) {
 		return fmt.Errorf("%w: invalid leading zero bits", ErrInvalidPow)

registration/pow_verifier_test.go

@@ -15,10 +15,11 @@ func TestVerify(t *testing.T) {
 	nodeID := make([]byte, 32)
 	poetChallenge := []byte("some data to submit to poet")
 	challenge := []byte("challenge this")
-	difficulty := uint(4)
+	difficulty := uint(6)
 
 	nonce, err := shared.FindSubmitPowNonce(context.Background(), challenge, poetChallenge, nodeID, difficulty)
 	require.NoError(t, err)
+	require.EqualValues(t, 143, nonce)
 
 	verifier := NewPowVerifier(NewPowParams(challenge, difficulty))
 	require.NoError(t, verifier.Verify(poetChallenge, nodeID, nonce))

shared/pow.go

@@ -0,0 +1,75 @@
+package shared
+
+import (
+	"context"
+	"encoding/binary"
+	"hash"
+
+	"github.com/c0mm4nd/go-ripemd"
+)
+
+// FindSubmitPowNonce finds the nonce that solves the PoW challenge.
+func FindSubmitPowNonce(
+	ctx context.Context,
+	powChallenge, poetChallenge, nodeID []byte,
+	difficulty uint,
+) (uint64, error) {
+	var hash []byte
+	p := NewPowHasher(powChallenge, nodeID, poetChallenge)
+
+	for nonce := uint64(0); ; nonce++ {
+		select {
+		case <-ctx.Done():
+			return 0, ctx.Err()
+		default:
+		}
+
+		hash = p.Hash(nonce, hash[:0])
+
+		if CheckLeadingZeroBits(hash, difficulty) {
+			return nonce, nil
+		}
+	}
+}
+
+type powHasher struct {
+	h     hash.Hash
+	input []byte
+}
+
+func NewPowHasher(inputs ...[]byte) *powHasher {
+	h := &powHasher{h: ripemd.New256(), input: []byte{}}
+	for _, in := range inputs {
+		h.input = append(h.input, in...)
+	}
+	h.input = append(h.input, make([]byte, 8)...) // placeholder for nonce
+	return h
+}
+
+func (p *powHasher) Hash(nonce uint64, output []byte) []byte {
+	nonceBytes := p.input[len(p.input)-8:]
+	binary.LittleEndian.PutUint64(nonceBytes, nonce)
+
+	p.h.Reset()
+	p.h.Write(p.input)
+	return p.h.Sum(output)
+}
+
+// CheckLeadingZeroBits checks if the first 'expected' bits of the byte array are all zero.
+func CheckLeadingZeroBits(data []byte, expected uint) bool {
+	if len(data)*8 < int(expected) {
+		return false
+	}
+	for i := 0; i < int(expected/8); i++ {
+		if data[i] != 0 {
+			return false
+		}
+	}
+	if expected%8 != 0 {
+		if data[expected/8]>>(8-expected%8) != 0 {
+			return false
+		}
+	}
+
+	return true
+}

shared/pow_test.go

@@ -0,0 +1,63 @@
+package shared_test
+
+import (
+	"context"
+	"encoding/binary"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/spacemeshos/poet/shared"
+)
+
+func TestCheckLeadingZeroBits(t *testing.T) {
+	r := require.New(t)
+
+	r.True(shared.CheckLeadingZeroBits([]byte{0x00}, 0))
+	r.True(shared.CheckLeadingZeroBits([]byte{0x00}, 8))
+
+	// Out of bounds
+	r.False(shared.CheckLeadingZeroBits([]byte{0x00}, 9))
+
+	r.True(shared.CheckLeadingZeroBits([]byte{0x0F}, 4))
+	r.False(shared.CheckLeadingZeroBits([]byte{0x0F}, 5))
+
+	r.True(shared.CheckLeadingZeroBits([]byte{0x00, 0x0F}, 5))
+	r.True(shared.CheckLeadingZeroBits([]byte{0x00, 0x0F}, 12))
+	r.False(shared.CheckLeadingZeroBits([]byte{0x00, 0x0F}, 13))
+}
+
+func BenchmarkPowHash(b *testing.B) {
+	powChallenge := make([]byte, 32)
+	poetChallenge := make([]byte, 32)
+	nodeID := make([]byte, 32)
+
+	var hash []byte
+	p := shared.NewPowHasher(powChallenge, poetChallenge, nodeID)
+	for i := 0; i < b.N; i++ {
+		hash = p.Hash(123678, hash[:0])
+	}
+}
+
+func BenchmarkFindSubmitPowNonce(b *testing.B) {
+	powChallenge := make([]byte, 32)
+	poetChallenge := make([]byte, 32)
+	nodeID := make([]byte, 32)
+	b.ResetTimer()
+	for difficulty := 20; difficulty <= 20; difficulty += 1 {
+		b.Run(fmt.Sprintf("difficulty=%v", difficulty), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				binary.LittleEndian.PutUint64(nodeID, uint64(i))
+				_, err := shared.FindSubmitPowNonce(
+					context.Background(),
+					powChallenge,
+					poetChallenge,
+					nodeID,
+					uint(difficulty),
+				)
+				require.NoError(b, err)
+			}
+		})
+	}
+}

shared/shared.go

@@ -1,11 +1,9 @@
 package shared
 
 import (
-	"context"
 	"encoding/binary"
 	"fmt"
 
-	"github.com/c0mm4nd/go-ripemd"
 	"github.com/spacemeshos/go-scale"
 	"github.com/spacemeshos/merkle-tree"
 	"github.com/spacemeshos/sha256-simd"
@@ -171,58 +169,6 @@ func DecodeSliceOfByteSliceWithLimit(d *scale.Decoder, outerLimit, innerLimit ui
 	return result, total, nil
 }
 
-// FindSubmitPowNonce finds the nonce that solves the PoW challenge.
-func FindSubmitPowNonce(
-	ctx context.Context,
-	powChallenge, poetChallenge, nodeID []byte,
-	difficulty uint,
-) (uint64, error) {
-	var hash []byte
-	for nonce := uint64(0); ; nonce++ {
-		select {
-		case <-ctx.Done():
-			return 0, ctx.Err()
-		default:
-		}
-
-		hash := CalcSubmitPowHash(powChallenge, poetChallenge, nodeID, hash, nonce)
-		if CheckLeadingZeroBits(hash, difficulty) {
-			return nonce, nil
-		}
-	}
-}
-
-// CalcSubmitPowHash calculates the hash for the Submit PoW.
-// The hash is ripemd256(powChallenge || nodeID || poetChallenge || nonce).
-func CalcSubmitPowHash(powChallenge, poetChallenge, nodeID, output []byte, nonce uint64) []byte {
-	md := ripemd.New256()
-	md.Write(powChallenge)
-	md.Write(nodeID)
-	md.Write(poetChallenge)
-	if err := binary.Write(md, binary.LittleEndian, nonce); err != nil {
-		panic(err)
-	}
-	return md.Sum(output)
-}
-
-// CheckLeadingZeroBits checks if the first 'expected' bits of the byte array are all zero.
-func CheckLeadingZeroBits(data []byte, expected uint) bool {
-	if len(data)*8 < int(expected) {
-		return false
-	}
-	for i := 0; i < int(expected/8); i++ {
-		if data[i] != 0 {
-			return false
-		}
-	}
-	if expected%8 != 0 {
-		if data[expected/8]>>(8-expected%8) != 0 {
-			return false
-		}
-	}
-	return true
-}
-
 // HashMembershipTreeNode calculates internal node of
 // the membership merkle tree.
 func HashMembershipTreeNode(buf, lChild, rChild []byte) []byte {

shared/shared_test.go

@@ -1,10 +1,8 @@
 package shared
 
 import (
-	"context"
 	"encoding/binary"
 	"encoding/hex"
-	"fmt"
 	"math"
 	"testing"
 
@@ -111,42 +109,3 @@ func BenchmarkFiatShamir(b *testing.B) {
 		FiatShamir([]byte("challenge this"), uint64(30000000000), T)
 	}
 }
-
-func TestCheckLeadingZeroBits(t *testing.T) {
-	r := require.New(t)
-
-	r.True(CheckLeadingZeroBits([]byte{0x00}, 0))
-	r.True(CheckLeadingZeroBits([]byte{0x00}, 8))
-
-	// Out of bounds
-	r.False(CheckLeadingZeroBits([]byte{0x00}, 9))
-
-	r.True(CheckLeadingZeroBits([]byte{0x0F}, 4))
-	r.False(CheckLeadingZeroBits([]byte{0x0F}, 5))
-
-	r.True(CheckLeadingZeroBits([]byte{0x00, 0x0F}, 5))
-	r.True(CheckLeadingZeroBits([]byte{0x00, 0x0F}, 12))
-	r.False(CheckLeadingZeroBits([]byte{0x00, 0x0F}, 13))
-}
-
-func BenchmarkFindSubmitPowNonce(b *testing.B) {
-	powChallenge := make([]byte, 32)
-	poetChallenge := make([]byte, 32)
-	nodeID := make([]byte, 32)
-	b.ResetTimer()
-	for difficulty := 0; difficulty <= 24; difficulty += 4 {
-		b.Run(fmt.Sprintf("difficulty=%v", difficulty), func(b *testing.B) {
-			for i := 0; i < b.N; i++ {
-				binary.LittleEndian.PutUint64(nodeID, uint64(i))
-				_, err := FindSubmitPowNonce(
-					context.Background(),
-					powChallenge,
-					poetChallenge,
-					nodeID,
-					uint(difficulty),
-				)
-				require.NoError(b, err)
-			}
-		})
-	}
-}