niko_trust/internal/pow/pow_test.go

package pow_test

import (
	"bytes"
	"encoding/binary"
	"math"
	"testing"

	blake3 "lukechampine.com/blake3"

	"git.n1ko.dev/Niko/niko_trust/internal/pow"
)

func testKey(seed byte) pow.Key {
	var k pow.Key
	for i := range k {
		k[i] = seed + byte(i)
	}
	return k
}

func testTarget(seed byte) [32]byte {
	var t [32]byte
	for i := range t {
		t[i] = seed * byte(i+1)
	}
	return t
}

func TestSolveAndVerifyRoundTrip(t *testing.T) {
	key := testKey(0x10)
	target := testTarget(0x20)
	const difficulty = 12

	counter, ok := pow.Solve(key, target, difficulty)
	if !ok {
		t.Fatal("solve failed")
	}
	if !pow.Verify(key, target, difficulty, counter) {
		t.Fatal("verify rejected own solution")
	}

	sum := pow.Sum(key, target, counter)
	if n := pow.LeadingZeroBits(sum); n < difficulty {
		t.Fatalf("solution has %d leading zero bits, want >= %d", n, difficulty)
	}
}

func TestVerifyRejectsWrongInputs(t *testing.T) {
	key := testKey(0x01)
	target := testTarget(0x02)
	const difficulty = 10

	counter, ok := pow.Solve(key, target, difficulty)
	if !ok {
		t.Fatal("solve failed")
	}

	cases := []struct {
		name       string
		key        pow.Key
		target     [32]byte
		difficulty int
		counter    uint32
	}{
		{"wrong key", testKey(0xFF), target, difficulty, counter},
		{"wrong target", key, testTarget(0xFE), difficulty, counter},
		{"wrong counter", key, target, difficulty, counter + 1},
		{"difficulty above max", key, target, pow.MaxDifficulty + 1, counter},
	}
	for _, tc := range cases {
		if pow.Verify(tc.key, tc.target, tc.difficulty, tc.counter) {
			t.Errorf("%s: verify accepted invalid proof", tc.name)
		}
	}
}

func TestSolveDeterministicSmallestCounter(t *testing.T) {
	key := testKey(0x33)
	target := testTarget(0x44)
	first, ok := pow.Solve(key, target, 8)
	if !ok {
		t.Fatal("solve failed")
	}
	second, _ := pow.Solve(key, target, 8)
	if first != second {
		t.Fatalf("solver nondeterministic: %d vs %d", first, second)
	}
	// The smallest counter is a real lower bound: nothing below verifies.
	for c := uint32(0); c < first; c++ {
		if pow.Verify(key, target, 8, c) {
			t.Fatalf("counter %d verifies but solver returned %d", c, first)
		}
	}
}

func TestLeadingZeroBitsBoundaries(t *testing.T) {
	cases := []struct {
		b    byte
		want int
	}{
		{0x80, 0},
		{0x40, 1},
		{0x01, 7},
		{0x00, 8},
	}
	var sum [32]byte
	for _, tc := range cases {
		for i := range sum {
			sum[i] = 0xFF
		}
		sum[0] = tc.b
		if got := pow.LeadingZeroBits(sum); got != tc.want {
			t.Errorf("byte %02x: got %d leading zero bits, want %d", tc.b, got, tc.want)
		}
	}

	for i := range sum {
		sum[i] = 0
	}
	if got := pow.LeadingZeroBits(sum); got != 256 {
		t.Errorf("all-zero sum: got %d, want 256", got)
	}
}

func TestMeetsTargetDifficultySemantics(t *testing.T) {
	var zero [32]byte
	var full [32]byte
	for i := range full {
		full[i] = 0xFF
	}

	if !pow.MeetsTarget(zero, 256) {
		t.Error("all-zero hash must meet any representable difficulty")
	}
	if pow.MeetsTarget(full, 1) {
		t.Error("all-ff hash must not meet difficulty 1")
	}
	if !pow.MeetsTarget(full, 0) || !pow.MeetsTarget(full, -5) {
		t.Error("difficulty <= 0 must always pass (disabled)")
	}
}

func TestSumMatchesManualConstruction(t *testing.T) {
	// Recompute the hash input independently and compare against Sum.
	key := testKey(0x55)
	target := testTarget(0x66)
	const counter = 0x01020304

	h := blake3.New(32, key[:])
	h.Write([]byte(pow.Domain))
	h.Write(target[:])
	var cb [4]byte
	binary.BigEndian.PutUint32(cb[:], counter)
	h.Write(cb[:])
	var want [32]byte
	copy(want[:], h.Sum(nil))

	got := pow.Sum(key, target, counter)
	if !bytes.Equal(got[:], want[:]) {
		t.Fatalf("Sum mismatch:\n got %x\nwant %x", got, want)
	}
}

func TestKeyParseAndEqual(t *testing.T) {
	key := testKey(0x77)
	parsed, err := pow.ParseKey(key.String())
	if err != nil {
		t.Fatal(err)
	}
	proof := pow.Proof{Key: parsed, KeyHex: parsed.String(), Counter: 1}
	same, _ := pow.ParseKey(key.String())
	if !proof.Equal(pow.Proof{Key: same, KeyHex: same.String(), Counter: 1}) {
		t.Error("equal proofs compared unequal")
	}
	if proof.Equal(pow.Proof{Key: parsed, KeyHex: parsed.String(), Counter: 2}) {
		t.Error("different counters compared equal")
	}

	bad := []string{"", "00", "ZZ", key.String() + "00"}
	for _, s := range bad {
		if _, err := pow.ParseKey(s); err == nil {
			t.Errorf("ParseKey(%q) accepted malformed key", s)
		}
	}
}

func TestNewKeyUnique(t *testing.T) {
	seen := make(map[pow.Key]struct{})
	for i := 0; i < 100; i++ {
		k := pow.NewKey()
		if _, dup := seen[k]; dup {
			t.Fatal("CSPRNG produced duplicate key")
		}
		seen[k] = struct{}{}
	}
}

// Solve at MaxDifficulty must terminate within the uint32 space with
// overwhelming probability; this bounds runtime while exercising deep loops.
func TestSolveHighDifficultyTerminates(t *testing.T) {
	key := testKey(0x99)
	target := testTarget(0xAA)
	_, ok := pow.Solve(key, target, 20)
	if !ok {
		t.Skip("counter space exhausted for this key/target pair")
	}
}

func BenchmarkVerify(b *testing.B) {
	key := testKey(0xAB)
	target := testTarget(0xCD)
	for i := 0; b.Loop(); i++ {
		pow.Verify(key, target, 22, uint32(i%math.MaxUint32))
	}
}