Arkie-Library-Backend/internal/handlers/wallet_auth.go

package handlers

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/arkie/ark-database/internal/auth"
	"github.com/ethereum/go-ethereum/accounts"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/jackc/pgx/v5/pgxpool"
)

var ethAddrRe = regexp.MustCompile(`^0x[0-9a-fA-F]{40}$`)

func EnsureWalletAuthSchema(ctx context.Context, pool *pgxpool.Pool) error {
	_, err := pool.Exec(ctx, `
		CREATE TABLE IF NOT EXISTS wallet_auth_nonces (
			address TEXT PRIMARY KEY,
			nonce TEXT NOT NULL,
			expires_at TIMESTAMPTZ NOT NULL
		);
		CREATE INDEX IF NOT EXISTS idx_wallet_auth_nonces_expires ON wallet_auth_nonces(expires_at);
	`)
	return err
}

func normalizeAddr(a string) string {
	return strings.ToLower(strings.TrimSpace(a))
}

type walletNonceReq struct {
	Address string `json:"address"`
}

type walletVerifyReq struct {
	Address   string `json:"address"`
	Message   string `json:"message"`
	Signature string `json:"signature"`
}

func WalletNonce() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var req walletNonceReq
		if err := jsonDecode(r, &req); err != nil {
			http.Error(w, "bad json", http.StatusBadRequest)
			return
		}
		addr := normalizeAddr(req.Address)
		if !ethAddrRe.MatchString(addr) {
			http.Error(w, "invalid address", http.StatusBadRequest)
			return
		}
		var b [16]byte
		if _, err := rand.Read(b[:]); err != nil {
			http.Error(w, "rng", http.StatusInternalServerError)
			return
		}
		nonce := hex.EncodeToString(b[:])
		pool := poolFrom(r)
		ctx := r.Context()
		_, _ = pool.Exec(ctx, `DELETE FROM wallet_auth_nonces WHERE expires_at < NOW()`)
		_, err := pool.Exec(ctx, `
			INSERT INTO wallet_auth_nonces (address, nonce, expires_at) VALUES ($1, $2, NOW() + INTERVAL '15 minutes')
			ON CONFLICT (address) DO UPDATE SET nonce = EXCLUDED.nonce, expires_at = EXCLUDED.expires_at`,
			addr, nonce)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		checksum := common.HexToAddress(addr).Hex()
		message := "ARK Database — wallet sign-in\n\n" +
			"Wallet: " + checksum + "\n" +
			"One-time code: " + nonce + "\n\n" +
			"Sign this message to log in. No transaction or gas fee."
		writeJSON(w, map[string]any{"nonce": nonce, "message": message})
	}
}

func WalletVerify(jwtSecret string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		var req walletVerifyReq
		if err := jsonDecode(r, &req); err != nil {
			http.Error(w, "bad json", http.StatusBadRequest)
			return
		}
		addr := normalizeAddr(req.Address)
		if !ethAddrRe.MatchString(addr) {
			http.Error(w, "invalid address", http.StatusBadRequest)
			return
		}
		sig := strings.TrimSpace(req.Signature)
		if !strings.HasPrefix(sig, "0x") {
			sig = "0x" + sig
		}
		msg := req.Message
		if msg == "" {
			http.Error(w, "message required", http.StatusBadRequest)
			return
		}
		pool := poolFrom(r)
		ctx := r.Context()
		var storedNonce string
		err := pool.QueryRow(ctx, `SELECT nonce FROM wallet_auth_nonces WHERE address = $1 AND expires_at > NOW()`, addr).Scan(&storedNonce)
		if err != nil {
			http.Error(w, "nonce expired or missing — request a new code", http.StatusUnauthorized)
			return
		}
		if !strings.Contains(msg, storedNonce) {
			http.Error(w, "message does not match nonce", http.StatusUnauthorized)
			return
		}
		recovered, err := recoverPersonalSign(msg, sig)
		if err != nil {
			http.Error(w, "bad signature", http.StatusUnauthorized)
			return
		}
		if normalizeAddr(recovered.Hex()) != addr {
			http.Error(w, "signer mismatch", http.StatusUnauthorized)
			return
		}
		_, _ = pool.Exec(ctx, `DELETE FROM wallet_auth_nonces WHERE address = $1`, addr)
		tok, err := auth.SignUserWallet(jwtSecret, common.HexToAddress(addr).Hex(), 30*24*time.Hour)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		writeJSON(w, map[string]any{"token": tok, "wallet": common.HexToAddress(addr).Hex()})
	}
}

func recoverPersonalSign(message, sigHex string) (common.Address, error) {
	sig, err := hex.DecodeString(strings.TrimPrefix(sigHex, "0x"))
	if err != nil || len(sig) != 65 {
		return common.Address{}, err
	}
	if sig[64] >= 27 {
		sig[64] -= 27
	}
	hash := accounts.TextHash([]byte(message))
	pub, err := crypto.SigToPub(hash, sig)
	if err != nil {
		return common.Address{}, err
	}
	return crypto.PubkeyToAddress(*pub), nil
}

func WalletMe(jwtSecret string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		h := r.Header.Get("Authorization")
		if !strings.HasPrefix(strings.ToLower(h), "bearer ") {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		tok := strings.TrimSpace(h[7:])
		claims, err := auth.ParseUserWallet(jwtSecret, tok)
		if err != nil {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		writeJSON(w, map[string]any{"wallet": claims.Wallet, "role": claims.Role})
	}
}