niko_trust/internal/server/pow.go

package server

import (
	"encoding/json"
	"net/http"
	"time"

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

// powTTL is how long an issued PoW challenge stays valid.
const powTTL = 5 * time.Minute

// powChallengeLimitPerMin caps challenge issuance per IP per minute.
const powChallengeLimitPerMin = 30

// powPurpose names what a challenge will be spent on. The purpose selects the
// difficulty the client must solve for, so it is fixed at issuance.
type powPurpose string

const (
	powPurposePut  powPurpose = "put"
	powPurposeAuth powPurpose = "auth"
)

// powChallengeRecord is one issued, unspent challenge.
type powChallengeRecord struct {
	key     pow.Key
	bits    int
	purpose powPurpose
	expiry  time.Time
}

// Option configures the relay at construction time.
type Option func(*Server)

// WithPow enables proof-of-work admission control: putBits guards
// POST /v1/objects, authBits guards POST /v1/auth/challenge. Zero disables a
// tier; values above pow.MaxDifficulty are clamped.
func WithPow(putBits, authBits int) Option {
	return func(s *Server) {
		s.powPutBits = clampDifficulty(putBits)
		s.powAuthBits = clampDifficulty(authBits)
	}
}

func clampDifficulty(n int) int {
	switch {
	case n < 0:
		return 0
	case n > pow.MaxDifficulty:
		return pow.MaxDifficulty
	default:
		return n
	}
}

// wirePow is the JSON form of a solved challenge attached to a request body.
type wirePow struct {
	Key     string `json:"key"`
	Counter uint32 `json:"counter"`
}

// rateLimitPowChallenge applies the per-IP cap before issuing.
func (s *Server) rateLimitPowChallenge(w http.ResponseWriter, r *http.Request) {
	if !s.powChallengeLimiter.allow(clientIP(r)) {
		s.metrics.inc(&s.metrics.powFailed)
		writeErr(w, http.StatusTooManyRequests, "rate limited")
		return
	}
	s.handlePowChallenge(w, r)
}

// handlePowChallenge issues a fresh single-use challenge key.
//
// The request may name a purpose ("put", default, or "auth"); the response
// carries the exact difficulty to solve for that purpose. Challenges are
// bound at spend time to the submission they accompany.
func (s *Server) handlePowChallenge(w http.ResponseWriter, r *http.Request) {
	if s.powPutBits <= 0 && s.powAuthBits <= 0 {
		writeErr(w, http.StatusBadRequest, "proof of work disabled")
		return
	}
	var body struct {
		Purpose string `json:"purpose"`
	}
	if r.Body != nil {
		_ = json.NewDecoder(http.MaxBytesReader(w, r.Body, 1024)).Decode(&body)
	}
	purpose := powPurpose(body.Purpose)
	if purpose == "" {
		purpose = powPurposePut
	}
	bits := s.powPutBits
	if purpose == powPurposeAuth {
		bits = s.powAuthBits
	}
	if purpose != powPurposePut && purpose != powPurposeAuth || bits <= 0 {
		s.metrics.inc(&s.metrics.powFailed)
		writeErr(w, http.StatusBadRequest, "unknown purpose or tier disabled")
		return
	}

	key := pow.NewKey()
	s.mu.Lock()
	s.prunePowLocked()
	s.powChallenges[key.String()] = powChallengeRecord{
		key:     key,
		bits:    bits,
		purpose: purpose,
		expiry:  now().Add(powTTL),
	}
	s.mu.Unlock()
	s.metrics.inc(&s.metrics.challengesIssued)
	writeJSON(w, http.StatusOK, map[string]any{
		"key":        key.String(),
		"difficulty": bits,
		"ttl":        int(powTTL / time.Second),
	})
}

// prunePowLocked drops expired challenges. The caller holds s.mu.
func (s *Server) prunePowLocked() {
	ts := now()
	for k, rec := range s.powChallenges {
		if rec.expiry.Before(ts) {
			delete(s.powChallenges, k)
		}
	}
}

// admitPoW enforces one tier of admission control. target binds the proof to
// this exact submission. It reports whether the request may proceed and, on
// refusal, the client-facing reason.
func (s *Server) admitPoW(proof *wirePow, purpose powPurpose, target [pow.TargetSize]byte, wantBits int) (bool, string) {
	if wantBits <= 0 {
		return true, ""
	}
	if proof == nil {
		s.metrics.inc(&s.metrics.powFailed)
		return false, "proof of work required"
	}
	key, err := pow.ParseProof(proof.Key, proof.Counter)
	if err != nil {
		s.metrics.inc(&s.metrics.powFailed)
		return false, "malformed proof of work"
	}
	hexKey := key.KeyHex
	s.mu.Lock()
	rec, ok := s.powChallenges[hexKey]
	if ok {
		delete(s.powChallenges, hexKey) // single use, even on failure paths below
	}
	s.mu.Unlock()
	if !ok || rec.purpose != purpose || now().After(rec.expiry) {
		s.metrics.inc(&s.metrics.powFailed)
		return false, "unknown or expired challenge"
	}
	if !pow.Verify(rec.key, target, rec.bits, proof.Counter) {
		s.metrics.inc(&s.metrics.powFailed)
		return false, "invalid proof of work"
	}
	s.metrics.inc(&s.metrics.powOK)
	return true, ""
}

// zeroTarget is the authentication binding target: an auth challenge proves
// work against nothing but the server-chosen key itself.
func zeroTarget() [pow.TargetSize]byte {
	return [pow.TargetSize]byte{}
}