neural-hunt/internal/auth/auth.go

package auth

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"errors"
	"math/big"
	"strings"
	"sync"
	"time"
)

var b64 = base64.RawURLEncoding

type PublicJWK struct {
	Kty    string   `json:"kty"`
	Crv    string   `json:"crv"`
	X      string   `json:"x"`
	Y      string   `json:"y"`
	Ext    bool     `json:"ext,omitempty"`
	KeyOps []string `json:"key_ops,omitempty"`
	// Optional standard JWK metadata. Authentication intentionally does not
	// depend on these values; accepting them keeps WebCrypto exports portable
	// across Chrome, Firefox, Safari and other standards-compliant browsers.
	Alg string `json:"alg,omitempty"`
	Use string `json:"use,omitempty"`
	Kid string `json:"kid,omitempty"`
}

type Claims struct {
	ClientID  string `json:"cid"`
	Role      string `json:"role"`
	SessionID string `json:"sid"`
	Exp       int64  `json:"exp"`
}

type challengeEntry struct {
	ClientID  string
	ExpiresAt time.Time
	ProofBits int
}

type Manager struct {
	secret []byte

	// Login challenges are intentionally process-local ephemeral state. They do
	// not belong in SQLite: issuing a challenge is on the hot path when many
	// browsers connect at once and must not compete with guesses/presence writes
	// for SQLite's single writer lock. Challenges are keyed by the random nonce,
	// not by client ID, so two concurrent tabs can each finish their own login.
	challengeMu sync.Mutex
	challenges  map[string]challengeEntry
}

// New keeps the database parameter for source compatibility with older builds.
// Authentication challenge state no longer writes to the database.
func New(_ *sql.DB, secret string) *Manager {
	return &Manager{secret: []byte(secret), challenges: make(map[string]challengeEntry)}
}

func ClientID(j PublicJWK) (string, error) {
	if j.Kty != "EC" || j.Crv != "P-256" || j.X == "" || j.Y == "" {
		return "", errors.New("only P-256 EC JWK is supported")
	}
	s := j.Kty + "|" + j.Crv + "|" + j.X + "|" + j.Y
	h := sha256.Sum256([]byte(s))
	return b64.EncodeToString(h[:]), nil
}

func PublicKey(j PublicJWK) (*ecdsa.PublicKey, error) {
	xb, err := b64.DecodeString(j.X)
	if err != nil {
		return nil, err
	}
	yb, err := b64.DecodeString(j.Y)
	if err != nil {
		return nil, err
	}
	x, y := new(big.Int).SetBytes(xb), new(big.Int).SetBytes(yb)
	if !elliptic.P256().IsOnCurve(x, y) {
		return nil, errors.New("invalid P-256 point")
	}
	return &ecdsa.PublicKey{Curve: elliptic.P256(), X: x, Y: y}, nil
}

func VerifyRaw(pub *ecdsa.PublicKey, message string, sigB64 string) bool {
	sig, err := b64.DecodeString(sigB64)
	if err != nil || len(sig) != 64 {
		return false
	}
	h := sha256.Sum256([]byte(message))
	r := new(big.Int).SetBytes(sig[:32])
	s := new(big.Int).SetBytes(sig[32:])
	return ecdsa.Verify(pub, h[:], r, s)
}

func randomB64(n int) (string, error) {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return b64.EncodeToString(b), nil
}

func (m *Manager) NewChallenge(ctx context.Context, cid string) (string, error) {
	return m.NewChallengeWithProof(ctx, cid, 0)
}

func (m *Manager) NewChallengeWithProof(ctx context.Context, cid string, proofBits int) (string, error) {
	if err := ctx.Err(); err != nil {
		return "", err
	}
	c, err := randomB64(32)
	if err != nil {
		return "", err
	}
	now := time.Now().UTC()
	expires := now.Add(2 * time.Minute)

	m.challengeMu.Lock()
	defer m.challengeMu.Unlock()

	// Opportunistic cleanup keeps the map bounded without a maintenance goroutine.
	for nonce, entry := range m.challenges {
		if !entry.ExpiresAt.After(now) {
			delete(m.challenges, nonce)
		}
	}
	if len(m.challenges) >= 50000 {
		return "", errors.New("too many pending authentication challenges")
	}
	if proofBits < 0 {
		proofBits = 0
	}
	if proofBits > 22 {
		proofBits = 22
	}
	m.challenges[c] = challengeEntry{ClientID: cid, ExpiresAt: expires, ProofBits: proofBits}
	return c, nil
}

func (m *Manager) ConsumeChallenge(ctx context.Context, cid, challenge string) error {
	return m.ConsumeChallengeWithProof(ctx, cid, challenge, "")
}

func leadingZeroBits(b []byte) int {
	n := 0
	for _, x := range b {
		if x == 0 {
			n += 8
			continue
		}
		for mask := byte(0x80); mask != 0 && x&mask == 0; mask >>= 1 {
			n++
		}
		break
	}
	return n
}

func verifyChallengeProof(challenge, cid, counter string, bits int) bool {
	if bits <= 0 {
		return true
	}
	if len(counter) == 0 || len(counter) > 24 {
		return false
	}
	for _, r := range counter {
		if r < '0' || r > '9' {
			return false
		}
	}
	h := sha256.Sum256([]byte("nh-pow-v1|" + challenge + "|" + cid + "|" + counter))
	return leadingZeroBits(h[:]) >= bits
}

func (m *Manager) ConsumeChallengeWithProof(ctx context.Context, cid, challenge, counter string) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	m.challengeMu.Lock()
	defer m.challengeMu.Unlock()

	entry, ok := m.challenges[challenge]
	if !ok {
		return errors.New("challenge expired or unknown")
	}
	// A known challenge is one-shot even if a malformed login tries to consume it.
	delete(m.challenges, challenge)
	if !entry.ExpiresAt.After(time.Now().UTC()) {
		return errors.New("challenge expired")
	}
	if entry.ClientID != cid {
		return errors.New("challenge identity mismatch")
	}
	if !verifyChallengeProof(challenge, cid, counter, entry.ProofBits) {
		return errors.New("identity proof of work invalid")
	}
	return nil
}

func (m *Manager) Issue(cid, role string, ttl time.Duration) (string, Claims, error) {
	sid, err := randomB64(18)
	if err != nil {
		return "", Claims{}, err
	}
	c := Claims{ClientID: cid, Role: role, SessionID: sid, Exp: time.Now().Add(ttl).Unix()}
	hdr := b64.EncodeToString([]byte(`{"alg":"HS256","typ":"JWT"}`))
	p, _ := json.Marshal(c)
	payload := b64.EncodeToString(p)
	unsigned := hdr + "." + payload
	mac := hmac.New(sha256.New, m.secret)
	mac.Write([]byte(unsigned))
	sig := b64.EncodeToString(mac.Sum(nil))
	return unsigned + "." + sig, c, nil
}

func (m *Manager) Parse(token string) (Claims, error) {
	var c Claims
	parts := strings.Split(token, ".")
	if len(parts) != 3 {
		return c, errors.New("bad token")
	}
	unsigned := parts[0] + "." + parts[1]
	mac := hmac.New(sha256.New, m.secret)
	mac.Write([]byte(unsigned))
	want := mac.Sum(nil)
	got, err := b64.DecodeString(parts[2])
	if err != nil || !hmac.Equal(want, got) {
		return c, errors.New("bad token signature")
	}
	pb, err := b64.DecodeString(parts[1])
	if err != nil {
		return c, err
	}
	if err := json.Unmarshal(pb, &c); err != nil {
		return c, err
	}
	if c.Exp < time.Now().Unix() {
		return c, errors.New("token expired")
	}
	return c, nil
}