dumpsterChat/internal/auth/webauthn.go

package auth

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"log/slog"
	"net/http"
	"sync"

	"git.dustin.coffee/hobokenchicken/dumpsterChat/internal/config"
	"git.dustin.coffee/hobokenchicken/dumpsterChat/internal/middleware"
	"github.com/go-webauthn/webauthn/protocol"
	"github.com/go-webauthn/webauthn/webauthn"
)

// challengeStore temporarily stores WebAuthn session data between begin and finish.
type challengeStore struct {
	mu    sync.Mutex
	data  map[string]*webauthn.SessionData
	order []string
}

func newChallengeStore() *challengeStore {
	return &challengeStore{data: make(map[string]*webauthn.SessionData)}
}

func (s *challengeStore) Set(key string, sd *webauthn.SessionData) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.data[key] = sd
	s.order = append(s.order, key)
	// Evict old entries (keep max 100)
	for len(s.order) > 100 {
		delete(s.data, s.order[0])
		s.order = s.order[1:]
	}
}

func (s *challengeStore) Get(key string) (*webauthn.SessionData, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	sd, ok := s.data[key]
	if ok {
		delete(s.data, key)
	}
	return sd, ok
}

type WebAuthnHandler struct {
	db         *sql.DB
	wa         *webauthn.WebAuthn
	sessions   *SessionStore
	challenges *challengeStore
	logger     *slog.Logger
}

func NewWebAuthnHandler(db *sql.DB, cfg *config.Config, sessions *SessionStore, logger *slog.Logger) (*WebAuthnHandler, error) {
	origins := []string{"https://" + cfg.Host, "http://" + cfg.Host + ":" + cfg.Port}
	if cfg.Host == "localhost" {
		origins = append(origins, "http://localhost:"+cfg.Port)
	}

	wa, err := webauthn.New(&webauthn.Config{
		RPDisplayName: "Dumpster",
		RPID:          cfg.Host,
		RPOrigins:     origins,
	})
	if err != nil {
		return nil, err
	}

	return &WebAuthnHandler{
		db:         db,
		wa:         wa,
		sessions:   sessions,
		challenges: newChallengeStore(),
		logger:     logger,
	}, nil
}

func (h *WebAuthnHandler) RegisterRoutes(r *http.ServeMux) {
	r.HandleFunc("POST /auth/webauthn/register/begin", h.RegisterBegin)
	r.HandleFunc("POST /auth/webauthn/register/finish", h.RegisterFinish)
	r.HandleFunc("POST /auth/webauthn/login/begin", h.LoginBegin)
	r.HandleFunc("POST /auth/webauthn/login/finish", h.LoginFinish)
}

// RegisterBegin starts passkey registration for an authenticated user.
func (h *WebAuthnHandler) RegisterBegin(w http.ResponseWriter, r *http.Request) {
	userID, ok := middleware.UserIDFromContext(r.Context())
	if !ok {
		http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized)
		return
	}

	user, err := h.loadUser(r.Context(), userID)
	if err != nil {
		http.Error(w, `{"error":"user not found"}`, http.StatusNotFound)
		return
	}

	options, sessionData, err := h.wa.BeginRegistration(user)
	if err != nil {
		h.logger.Error("webauthn register begin failed", "error", err)
		http.Error(w, `{"error":"registration failed"}`, http.StatusInternalServerError)
		return
	}

	// Store challenge keyed by user ID
	h.challenges.Set(userID, sessionData)

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(options)
}

// RegisterFinish completes passkey registration.
func (h *WebAuthnHandler) RegisterFinish(w http.ResponseWriter, r *http.Request) {
	userID, ok := middleware.UserIDFromContext(r.Context())
	if !ok {
		http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized)
		return
	}

	sessionData, ok := h.challenges.Get(userID)
	if !ok {
		http.Error(w, `{"error":"no pending registration"}`, http.StatusBadRequest)
		return
	}

	user, err := h.loadUser(r.Context(), userID)
	if err != nil {
		http.Error(w, `{"error":"user not found"}`, http.StatusNotFound)
		return
	}

	credential, err := h.wa.FinishRegistration(user, *sessionData, r)
	if err != nil {
		h.logger.Error("webauthn register finish failed", "error", err)
		http.Error(w, `{"error":"registration verification failed"}`, http.StatusBadRequest)
		return
	}

	// Store credential in database
	_, err = h.db.ExecContext(r.Context(),
		`INSERT INTO webauthn_credentials (user_id, credential_id, public_key, aaguid, sign_count)
		 VALUES ($1, $2, $3, $4, $5)`,
		userID,
		credential.ID,
		credential.PublicKey,
		credential.Authenticator.AAGUID,
		credential.Authenticator.SignCount,
	)
	if err != nil {
		h.logger.Error("failed to store credential", "error", err)
		http.Error(w, `{"error":"failed to store credential"}`, http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"status": "registered"})
}

// LoginBegin starts passkey authentication (no auth required).
func (h *WebAuthnHandler) LoginBegin(w http.ResponseWriter, r *http.Request) {
	// Get all registered credentials for the login ceremony
	rows, err := h.db.QueryContext(r.Context(),
		`SELECT credential_id, user_id FROM webauthn_credentials LIMIT 1000`)
	if err != nil {
		h.logger.Error("failed to query credentials", "error", err)
		http.Error(w, `{"error":"failed"}`, http.StatusInternalServerError)
		return
	}
	defer rows.Close()

	var descriptors []protocol.CredentialDescriptor
	for rows.Next() {
		var credID []byte
		var userID string
		if err := rows.Scan(&credID, &userID); err != nil {
			continue
		}
		descriptors = append(descriptors, protocol.CredentialDescriptor{
			Type:         protocol.PublicKeyCredentialType,
			CredentialID: credID,
		})
	}

	options, sessionData, err := h.wa.BeginLogin(nil, webauthn.WithAllowedCredentials(descriptors))
	if err != nil {
		h.logger.Error("webauthn login begin failed", "error", err)
		http.Error(w, `{"error":"login failed"}`, http.StatusInternalServerError)
		return
	}

	// Store challenge with a random key
	challengeKey := randomChallenge()
	h.challenges.Set(challengeKey, sessionData)

	// Set cookie to link challenge to user
	http.SetCookie(w, &http.Cookie{
		Name:     "webauthn_challenge",
		Value:    challengeKey,
		Path:     "/",
		HttpOnly: true,
		MaxAge:   300, // 5 minutes
		SameSite: http.SameSiteStrictMode,
	})

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(options)
}

// LoginFinish completes passkey authentication.
func (h *WebAuthnHandler) LoginFinish(w http.ResponseWriter, r *http.Request) {
	// Get challenge key from cookie
	cookie, err := r.Cookie("webauthn_challenge")
	if err != nil {
		http.Error(w, `{"error":"no pending login"}`, http.StatusBadRequest)
		return
	}

	sessionData, ok := h.challenges.Get(cookie.Value)
	if !ok {
		http.Error(w, `{"error":"challenge expired"}`, http.StatusBadRequest)
		return
	}

	// Parse the credential response
	var parsedResponse struct {
		ID       string `json:"id"`
		RawID    string `json:"rawId"`
		Type     string `json:"type"`
		Response struct {
			AuthenticatorData string `json:"authenticatorData"`
			ClientDataJSON    string `json:"clientDataJSON"`
			Signature         string `json:"signature"`
			UserHandle        string `json:"userHandle"`
		} `json:"response"`
	}
	if err := json.NewDecoder(r.Body).Decode(&parsedResponse); err != nil {
		http.Error(w, `{"error":"invalid request"}`, http.StatusBadRequest)
		return
	}

	// Find the credential in the database
	var credentialID []byte
	credentialID, err = base64.StdEncoding.DecodeString(parsedResponse.ID)
	if err != nil {
		// Try raw URL encoding
		credentialID, err = base64.RawURLEncoding.DecodeString(parsedResponse.ID)
		if err != nil {
			http.Error(w, `{"error":"invalid credential"}`, http.StatusBadRequest)
			return
		}
	}

	var storedUserID string
	var storedPubKey []byte
	var signCount int64
	err = h.db.QueryRowContext(r.Context(),
		`SELECT user_id, public_key, sign_count FROM webauthn_credentials WHERE credential_id = $1`,
		credentialID,
	).Scan(&storedUserID, &storedPubKey, &signCount)
	if err != nil {
		http.Error(w, `{"error":"credential not found"}`, http.StatusUnauthorized)
		return
	}

	user, err := h.loadUser(r.Context(), storedUserID)
	if err != nil {
		http.Error(w, `{"error":"user not found"}`, http.StatusNotFound)
		return
	}

	_, err = h.wa.FinishLogin(user, *sessionData, r)
	if err != nil {
		h.logger.Error("webauthn login finish failed", "error", err)
		http.Error(w, `{"error":"login verification failed"}`, http.StatusUnauthorized)
		return
	}

	// Update sign count
	h.db.ExecContext(r.Context(),
		`UPDATE webauthn_credentials SET sign_count = sign_count + 1 WHERE credential_id = $1`,
		credentialID)

	// Create session
	token, err := h.sessions.Create(r.Context(), storedUserID)
	if err != nil {
		h.logger.Error("failed to create session", "error", err)
		http.Error(w, `{"error":"session creation failed"}`, http.StatusInternalServerError)
		return
	}

	http.SetCookie(w, &http.Cookie{
		Name:     "dumpster_session",
		Value:    token,
		Path:     "/",
		HttpOnly: true,
		MaxAge:   86400 * 30, // 30 days
		SameSite: http.SameSiteStrictMode,
	})

	// Clear challenge cookie
	http.SetCookie(w, &http.Cookie{
		Name:     "webauthn_challenge",
		Value:    "",
		Path:     "/",
		HttpOnly: true,
		MaxAge:   -1,
		SameSite: http.SameSiteStrictMode,
	})

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"status": "authenticated"})
}

// loadUser loads a user and their credentials from the database.
func (h *WebAuthnHandler) loadUser(ctx context.Context, userID string) (*WebAuthnUser, error) {
	var user WebAuthnUser
	err := h.db.QueryRowContext(ctx,
		`SELECT id, username, display_name FROM users WHERE id = $1`, userID,
	).Scan(&user.ID, &user.Username, &user.DisplayName)
	if err != nil {
		return nil, err
	}

	rows, err := h.db.QueryContext(ctx,
		`SELECT credential_id, public_key, sign_count FROM webauthn_credentials WHERE user_id = $1`, userID)
	if err != nil {
		return &user, nil
	}
	defer rows.Close()

	for rows.Next() {
		var credID, pubKey []byte
		var signCount uint32
		rows.Scan(&credID, &pubKey, &signCount)
		user.credentials = append(user.credentials, webauthn.Credential{
			ID:        credID,
			PublicKey: pubKey,
			Authenticator: webauthn.Authenticator{
				SignCount: signCount,
			},
		})
	}

	return &user, nil
}

// WebAuthnUser implements the webauthn.User interface
type WebAuthnUser struct {
	ID          string
	Username    string
	DisplayName string
	credentials []webauthn.Credential
}

func (u *WebAuthnUser) WebAuthnID() []byte                         { return []byte(u.ID) }
func (u *WebAuthnUser) WebAuthnName() string                       { return u.Username }
func (u *WebAuthnUser) WebAuthnDisplayName() string                { return u.DisplayName }
func (u *WebAuthnUser) WebAuthnIcon() string                       { return "" }
func (u *WebAuthnUser) WebAuthnCredentials() []webauthn.Credential { return u.credentials }

func randomChallenge() string {
	b := make([]byte, 32)
	rand.Read(b)
	return base64.RawURLEncoding.EncodeToString(b)
}