niko_trust/internal/server/pow_test.go

package server_test

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"git.n1ko.dev/Niko/niko_trust/internal/identity/signer"
	"git.n1ko.dev/Niko/niko_trust/internal/pow"
	"git.n1ko.dev/Niko/niko_trust/internal/protocol"
	"git.n1ko.dev/Niko/niko_trust/internal/server"
	"git.n1ko.dev/Niko/niko_trust/internal/tce"
)

// newPowServer builds a relay with both PoW tiers at a difficulty cheap
// enough to solve inline in tests.
func newPowServer(t *testing.T) *httptest.Server {
	t.Helper()
	srv := server.New(server.Config{Audience: "trust.n1ko.dev"}, server.WithPow(8, 8))
	return httptest.NewServer(srv.Handler())
}

func getPoWChallenge(t *testing.T, baseURL, purpose string) (keyHex string, difficulty int) {
	t.Helper()
	body := ""
	if purpose != "" {
		body = `{"purpose":"` + purpose + `"}`
	}
	resp, err := http.Post(baseURL+"/v1/pow/challenge", "application/json", strings.NewReader(body))
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(resp.Body)
		t.Fatalf("challenge status %d: %s", resp.StatusCode, raw)
	}
	var out struct {
		Key        string `json:"key"`
		Difficulty int    `json:"difficulty"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		t.Fatal(err)
	}
	return out.Key, out.Difficulty
}

func putWithPoW(t *testing.T, baseURL string, tceBytes, sig []byte, keyHex string, counter uint32) *http.Response {
	t.Helper()
	body := map[string]any{"tce": tceBytes, "signature": sig}
	if keyHex != "" {
		body["pow"] = map[string]any{"key": keyHex, "counter": counter}
	}
	buf := &bytes.Buffer{}
	if err := json.NewEncoder(buf).Encode(body); err != nil {
		t.Fatal(err)
	}
	resp, err := http.Post(baseURL+"/v1/objects", "application/json", buf)
	if err != nil {
		t.Fatal(err)
	}
	return resp
}

func makeClaim(t *testing.T) (tceBytes, sig []byte) {
	t.Helper()
	issuer, _ := signer.Generate()
	subject, _ := signer.Generate()
	c := &protocol.Claim{
		Issuer:    issuer.Public(),
		Subject:   subject.Public(),
		Claims:    map[string]tce.Value{"pow.test": tce.Bool(true)},
		CreatedAt: uint64(time.Now().Unix()),
		Serial:    1,
		Nonce:     bytes.Repeat([]byte{0x02}, tce.NonceSize),
	}
	b, err := protocol.EncodeClaim(c)
	if err != nil {
		t.Fatal(err)
	}
	return b, issuer.Sign(b)
}

func TestPutRequiresProofOfWork(t *testing.T) {
	ts := newPowServer(t)
	defer ts.Close()

	b, sig := makeClaim(t)
	resp := putWithPoW(t, ts.URL, b, sig, "", 0)
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusTooManyRequests {
		t.Fatalf("status %d, want 429", resp.StatusCode)
	}
	var out struct {
		Error string `json:"error"`
	}
	_ = json.NewDecoder(resp.Body).Decode(&out)
	if !strings.Contains(out.Error, "proof of work required") {
		t.Fatalf("unexpected error %q", out.Error)
	}
}

func TestPutWithValidProofSucceedsAndChallengeIsSingleUse(t *testing.T) {
	ts := newPowServer(t)
	defer ts.Close()

	keyHex, diff := getPoWChallenge(t, ts.URL, "put")
	if diff != 8 {
		t.Fatalf("difficulty %d, want 8", diff)
	}

	b, sig := makeClaim(t)
	target := [32]byte(tce.ComputeID(b))
	k, err := pow.ParseKey(keyHex)
	if err != nil {
		t.Fatal(err)
	}
	counter, ok := pow.Solve(k, target, diff)
	if !ok {
		t.Fatal("solve failed")
	}

	resp := putWithPoW(t, ts.URL, b, sig, keyHex, counter)
	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(resp.Body)
		t.Fatalf("status %d: %s", resp.StatusCode, raw)
	}
	resp.Body.Close()

	// The same challenge must not open a second submission.
	b2, sig2 := makeClaim(t)
	resp2 := putWithPoW(t, ts.URL, b2, sig2, keyHex, counter)
	defer resp2.Body.Close()
	if resp2.StatusCode != http.StatusTooManyRequests {
		t.Fatalf("replayed challenge status %d, want 429", resp2.StatusCode)
	}
}

func TestPutWithWrongCounterFails(t *testing.T) {
	ts := newPowServer(t)
	defer ts.Close()

	b, sig := makeClaim(t)
	target := [32]byte(tce.ComputeID(b))

	// At low difficulty the successor of the minimal solution occasionally
	// also meets the target (~2^-8); retry challenges until the wrong
	// counter is provably wrong, keeping the test deterministic.
	for attempt := 0; attempt < 8; attempt++ {
		keyHex, diff := getPoWChallenge(t, ts.URL, "put")
		k, err := pow.ParseKey(keyHex)
		if err != nil {
			t.Fatal(err)
		}
		counter, ok := pow.Solve(k, target, diff)
		if !ok {
			t.Fatal("solve failed")
		}
		if pow.Verify(k, target, diff, counter+1) {
			continue // unlucky key: successor is also valid, try another
		}

		resp := putWithPoW(t, ts.URL, b, sig, keyHex, counter+1)
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusTooManyRequests {
			t.Fatalf("status %d, want 429", resp.StatusCode)
		}
		return
	}
	t.Skip("no challenge yielded an invalid successor within budget")
}

func TestChallengePurposeIsEnforced(t *testing.T) {
	ts := newPowServer(t)
	defer ts.Close()

	// An auth-purpose challenge must not unlock object storage.
	keyHex, _ := getPoWChallenge(t, ts.URL, "auth")
	b, sig := makeClaim(t)
	resp := putWithPoW(t, ts.URL, b, sig, keyHex, 7)
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusTooManyRequests {
		t.Fatalf("cross-purpose challenge accepted with status %d", resp.StatusCode)
	}
}

func TestAuthChallengeRequiresProofOfWork(t *testing.T) {
	ts := newPowServer(t)
	defer ts.Close()

	// Without a proof the auth challenge is refused.
	resp, err := http.Post(ts.URL+"/v1/auth/challenge", "application/json", nil)
	if err != nil {
		t.Fatal(err)
	}
	resp.Body.Close()
	if resp.StatusCode != http.StatusTooManyRequests {
		t.Fatalf("status %d, want 429", resp.StatusCode)
	}

	// With a valid proof it succeeds and returns an auth challenge.
	keyHex, diff := getPoWChallenge(t, ts.URL, "auth")
	k, _ := pow.ParseKey(keyHex)
	var zero [32]byte
	counter, ok := pow.Solve(k, zero, diff)
	if !ok {
		t.Fatal("solve failed")
	}
	buf := &bytes.Buffer{}
	json.NewEncoder(buf).Encode(map[string]any{"pow": map[string]any{"key": keyHex, "counter": counter}})
	resp2, err := http.Post(ts.URL+"/v1/auth/challenge", "application/json", buf)
	if err != nil {
		t.Fatal(err)
	}
	defer resp2.Body.Close()
	if resp2.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(resp2.Body)
		t.Fatalf("status %d: %s", resp2.StatusCode, raw)
	}
	var out struct {
		Challenge string `json:"challenge"`
	}
	if err := json.NewDecoder(resp2.Body).Decode(&out); err != nil || out.Challenge == "" {
		t.Fatalf("no challenge returned (err=%v)", err)
	}
}

func TestUnknownChallengeRejected(t *testing.T) {
	ts := newPowServer(t)
	defer ts.Close()

	fake := pow.NewKey()
	b, sig := makeClaim(t)
	resp := putWithPoW(t, ts.URL, b, sig, fake.String(), 3)
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusTooManyRequests {
		t.Fatalf("status %d, want 429", resp.StatusCode)
	}
}

func TestPoWChallengeRateLimitedPerIP(t *testing.T) {
	ts := newPowServer(t)
	defer ts.Close()

	for i := 0; i < 30; i++ {
		resp, err := http.Post(ts.URL+"/v1/pow/challenge", "application/json", nil)
		if err != nil {
			t.Fatal(err)
		}
		io.Copy(io.Discard, resp.Body)
		resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			t.Fatalf("request %d: status %d, want 200", i, resp.StatusCode)
		}
	}
	resp, err := http.Post(ts.URL+"/v1/pow/challenge", "application/json", nil)
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusTooManyRequests {
		t.Fatalf("status %d, want 429 on request 31", resp.StatusCode)
	}
}

func TestPoWMetricsExposed(t *testing.T) {
	ts := newPowServer(t)
	defer ts.Close()

	keyHex, diff := getPoWChallenge(t, ts.URL, "put")
	k, _ := pow.ParseKey(keyHex)
	b, sig := makeClaim(t)
	counter, _ := pow.Solve(k, [32]byte(tce.ComputeID(b)), diff)
	resp := putWithPoW(t, ts.URL, b, sig, keyHex, counter)
	resp.Body.Close()

	mresp, err := http.Get(ts.URL + "/v1/metrics")
	if err != nil {
		t.Fatal(err)
	}
	defer mresp.Body.Close()
	raw, _ := io.ReadAll(mresp.Body)
	for _, want := range []string{"trust_pow_ok 1", "trust_pow_failed"} {
		if !strings.Contains(string(raw), want) {
			t.Errorf("metrics missing %q:\n%s", want, raw)
		}
	}
	fmt.Fprint(io.Discard, raw)
}