MiauInv/handlers/passkeys.go

package handlers

import (
	"MiauInv/auth"
	"MiauInv/models"
	"MiauInv/storage"
	utils "MiauInv/util"
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/go-webauthn/webauthn/protocol"
	"github.com/go-webauthn/webauthn/webauthn"
)

const passkeyChallengeTTL = 5 * time.Minute

type passkeyWebAuthnUser struct {
	user        models.User
	credentials []webauthn.Credential
}

func (user passkeyWebAuthnUser) WebAuthnID() []byte {
	return []byte(user.user.ID)
}

func (user passkeyWebAuthnUser) WebAuthnName() string {
	return user.user.Username
}

func (user passkeyWebAuthnUser) WebAuthnDisplayName() string {
	return user.user.Username
}

func (user passkeyWebAuthnUser) WebAuthnCredentials() []webauthn.Credential {
	return user.credentials
}

func Passkeys(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case http.MethodGet:
		listPasskeys(w, r)
	case http.MethodDelete:
		deletePasskey(w, r)
	default:
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

func PasskeyRegisterOptions(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		Name     string `json:"name"`
		Password string `json:"password"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		log.Println("POST [api/passkeys/register/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	if strings.TrimSpace(req.Password) == "" {
		http.Error(w, "Current password required", http.StatusBadRequest)
		return
	}

	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	user, err := storage.GetUserById(claims.UserID)
	if err != nil {
		log.Println("POST [api/passkeys/register/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "User not found", http.StatusNotFound)
		return
	}

	if !auth.CheckPasswordHash(req.Password, user.Password) {
		http.Error(w, "Invalid password", http.StatusUnauthorized)
		return
	}

	waUser, err := loadPasskeyWebAuthnUser(user)
	if err != nil {
		log.Println("POST [api/passkeys/register/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not load passkeys", http.StatusInternalServerError)
		return
	}

	wa, err := webAuthnForRequest(r)
	if err != nil {
		log.Println("POST [api/passkeys/register/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid WebAuthn origin", http.StatusBadRequest)
		return
	}

	creation, sessionData, err := wa.BeginRegistration(waUser,
		webauthn.WithResidentKeyRequirement(protocol.ResidentKeyRequirementRequired),
		webauthn.WithAuthenticatorSelection(protocol.AuthenticatorSelection{
			RequireResidentKey: protocol.ResidentKeyRequired(),
			ResidentKey:        protocol.ResidentKeyRequirementRequired,
			UserVerification:   protocol.VerificationRequired,
		}),
		webauthn.WithExclusions(webauthn.Credentials(waUser.WebAuthnCredentials()).CredentialDescriptors()),
		webauthn.WithExtensions(map[string]any{"credProps": true}),
	)
	if err != nil {
		log.Println("POST [api/passkeys/register/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not create passkey registration challenge", http.StatusInternalServerError)
		return
	}

	sessionToken, err := utils.GenerateOpaqueToken()
	if err != nil {
		log.Println("POST [api/passkeys/register/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not create passkey challenge", http.StatusInternalServerError)
		return
	}
	if err := storage.SavePasskeyChallenge(sessionToken, user.ID, storage.PasskeyCeremonyRegister, *sessionData, passkeyChallengeTTL); err != nil {
		log.Println("POST [api/passkeys/register/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not save passkey challenge", http.StatusInternalServerError)
		return
	}
	_ = storage.CleanupExpiredPasskeyChallenges()

	writeJSON(w, http.StatusOK, map[string]interface{}{
		"session_token": sessionToken,
		"options":       creation,
	})
	log.Println("POST [api/passkeys/register/options] " + r.RemoteAddr + ": Created passkey registration challenge")
}

func PasskeyRegisterFinish(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		SessionToken string          `json:"session_token"`
		Name         string          `json:"name"`
		Credential   json.RawMessage `json:"credential"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		log.Println("POST [api/passkeys/register/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}
	if req.SessionToken == "" || len(req.Credential) == 0 {
		http.Error(w, "Session token and credential required", http.StatusBadRequest)
		return
	}

	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	user, err := storage.GetUserById(claims.UserID)
	if err != nil {
		log.Println("POST [api/passkeys/register/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "User not found", http.StatusNotFound)
		return
	}

	challenge, sessionData, err := storage.ConsumePasskeyChallenge(req.SessionToken, storage.PasskeyCeremonyRegister)
	if err != nil {
		log.Println("POST [api/passkeys/register/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid or expired passkey registration challenge", http.StatusUnauthorized)
		return
	}
	if challenge.UserID != user.ID {
		http.Error(w, "Invalid passkey registration challenge", http.StatusUnauthorized)
		return
	}

	waUser, err := loadPasskeyWebAuthnUser(user)
	if err != nil {
		log.Println("POST [api/passkeys/register/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not load passkeys", http.StatusInternalServerError)
		return
	}

	wa, err := webAuthnForRequest(r)
	if err != nil {
		log.Println("POST [api/passkeys/register/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid WebAuthn origin", http.StatusBadRequest)
		return
	}

	credentialRequest := cloneRequestWithJSONBody(r, req.Credential)
	credential, err := wa.FinishRegistration(waUser, sessionData, credentialRequest)
	if err != nil {
		log.Println("POST [api/passkeys/register/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not verify passkey registration", http.StatusUnauthorized)
		return
	}

	name := strings.TrimSpace(req.Name)
	if name == "" {
		name = "Passkey"
	}
	if _, err := storage.AddPasskeyCredential(user.ID, name, credential); err != nil {
		log.Println("POST [api/passkeys/register/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not save passkey", http.StatusInternalServerError)
		return
	}

	if err := storage.RevokeAllRefreshTokensForUser(user.ID); err != nil {
		log.Println("POST [api/passkeys/register/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not revoke old sessions", http.StatusInternalServerError)
		return
	}

	passkeys, _ := storage.ListPasskeyCredentials(user.ID)
	issueLoginSessionWithExtra(w, r, user, map[string]interface{}{
		"passkeys_enabled": true,
		"passkey_count":    len(passkeys),
		"passkeys":         publicPasskeys(passkeys),
	})
	log.Println("POST [api/passkeys/register/finish] " + r.RemoteAddr + ": Registered passkey and revoked old sessions")
}

func PasskeyLoginOptions(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	wa, err := webAuthnForRequest(r)
	if err != nil {
		log.Println("POST [api/passkeys/login/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid WebAuthn origin", http.StatusBadRequest)
		return
	}

	assertion, sessionData, err := wa.BeginDiscoverableLogin(webauthn.WithUserVerification(protocol.VerificationRequired))
	if err != nil {
		log.Println("POST [api/passkeys/login/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not create passkey login challenge", http.StatusInternalServerError)
		return
	}

	sessionToken, err := utils.GenerateOpaqueToken()
	if err != nil {
		log.Println("POST [api/passkeys/login/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not create passkey challenge", http.StatusInternalServerError)
		return
	}
	if err := storage.SavePasskeyChallenge(sessionToken, "", storage.PasskeyCeremonyLogin, *sessionData, passkeyChallengeTTL); err != nil {
		log.Println("POST [api/passkeys/login/options] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not save passkey challenge", http.StatusInternalServerError)
		return
	}
	_ = storage.CleanupExpiredPasskeyChallenges()

	writeJSON(w, http.StatusOK, map[string]interface{}{
		"session_token": sessionToken,
		"options":       assertion,
	})
	log.Println("POST [api/passkeys/login/options] " + r.RemoteAddr + ": Created passkey login challenge")
}

func PasskeyLoginFinish(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		SessionToken string          `json:"session_token"`
		Credential   json.RawMessage `json:"credential"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		log.Println("POST [api/passkeys/login/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}
	if req.SessionToken == "" || len(req.Credential) == 0 {
		http.Error(w, "Session token and credential required", http.StatusBadRequest)
		return
	}

	challenge, sessionData, err := storage.ConsumePasskeyChallenge(req.SessionToken, storage.PasskeyCeremonyLogin)
	if err != nil {
		log.Println("POST [api/passkeys/login/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid or expired passkey login challenge", http.StatusUnauthorized)
		return
	}

	wa, err := webAuthnForRequest(r)
	if err != nil {
		log.Println("POST [api/passkeys/login/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid WebAuthn origin", http.StatusBadRequest)
		return
	}

	credentialRequest := cloneRequestWithJSONBody(r, req.Credential)

	var user models.User
	var credential *webauthn.Credential
	if challenge.UserID == "" {
		var webAuthnUser webauthn.User
		webAuthnUser, credential, err = wa.FinishPasskeyLogin(passkeyDiscoverableUserHandler, sessionData, credentialRequest)
		if err == nil {
			resolvedUser, ok := webAuthnUser.(passkeyWebAuthnUser)
			if !ok {
				err = errors.New("invalid passkey user type")
			} else {
				user = resolvedUser.user
			}
		}
	} else {
		user, err = storage.GetUserById(challenge.UserID)
		if err == nil {
			waUser, loadErr := loadPasskeyWebAuthnUser(user)
			if loadErr != nil {
				err = loadErr
			} else {
				credential, err = wa.FinishLogin(waUser, sessionData, credentialRequest)
			}
		}
	}
	if err != nil {
		log.Println("POST [api/passkeys/login/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not verify passkey login", http.StatusUnauthorized)
		return
	}

	if err := storage.UpdatePasskeyCredentialAfterLogin(user.ID, credential); err != nil {
		log.Println("POST [api/passkeys/login/finish] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not update passkey", http.StatusInternalServerError)
		return
	}

	issueLoginSession(w, r, user)
	log.Println("POST [api/passkeys/login/finish] " + r.RemoteAddr + ": Successfully logged in with passkey")
}

func PasskeyDisable(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		Password string `json:"password"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		log.Println("POST [api/passkeys/disable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	user, err := storage.GetUserById(claims.UserID)
	if err != nil {
		log.Println("POST [api/passkeys/disable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "User not found", http.StatusNotFound)
		return
	}
	if !auth.CheckPasswordHash(req.Password, user.Password) {
		http.Error(w, "Invalid password", http.StatusUnauthorized)
		return
	}

	if err := storage.DeleteAllPasskeyCredentials(user.ID); err != nil {
		log.Println("POST [api/passkeys/disable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not disable passkeys", http.StatusInternalServerError)
		return
	}
	if err := storage.RevokeAllRefreshTokensForUser(user.ID); err != nil {
		log.Println("POST [api/passkeys/disable] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not revoke old sessions", http.StatusInternalServerError)
		return
	}

	issueLoginSessionWithExtra(w, r, user, map[string]interface{}{
		"passkeys_enabled": false,
		"passkey_count":    0,
		"passkeys":         []map[string]interface{}{},
	})
	log.Println("POST [api/passkeys/disable] " + r.RemoteAddr + ": Disabled passkeys and revoked old sessions")
}

func listPasskeys(w http.ResponseWriter, r *http.Request) {
	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	passkeys, err := storage.ListPasskeyCredentials(claims.UserID)
	if err != nil {
		log.Println("GET [api/passkeys] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not load passkeys", http.StatusInternalServerError)
		return
	}

	writeJSON(w, http.StatusOK, map[string]interface{}{
		"passkeys":         publicPasskeys(passkeys),
		"passkeys_enabled": len(passkeys) > 0,
		"passkey_count":    len(passkeys),
	})
}

func deletePasskey(w http.ResponseWriter, r *http.Request) {
	var req struct {
		ID       string `json:"id"`
		Password string `json:"password"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		log.Println("DELETE [api/passkeys] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}
	if req.ID == "" || req.Password == "" {
		http.Error(w, "Passkey ID and password required", http.StatusBadRequest)
		return
	}

	claims := r.Context().Value(auth.UserContextKey).(*auth.Claims)
	user, err := storage.GetUserById(claims.UserID)
	if err != nil {
		log.Println("DELETE [api/passkeys] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "User not found", http.StatusNotFound)
		return
	}
	if !auth.CheckPasswordHash(req.Password, user.Password) {
		http.Error(w, "Invalid password", http.StatusUnauthorized)
		return
	}

	if _, err := storage.GetPasskeyCredentialByID(user.ID, req.ID); err != nil {
		http.Error(w, "Passkey not found", http.StatusNotFound)
		return
	}
	if err := storage.DeletePasskeyCredential(user.ID, req.ID); err != nil {
		log.Println("DELETE [api/passkeys] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not delete passkey", http.StatusInternalServerError)
		return
	}
	if err := storage.RevokeAllRefreshTokensForUser(user.ID); err != nil {
		log.Println("DELETE [api/passkeys] " + r.RemoteAddr + ": " + err.Error())
		http.Error(w, "Could not revoke old sessions", http.StatusInternalServerError)
		return
	}

	passkeys, _ := storage.ListPasskeyCredentials(user.ID)
	issueLoginSessionWithExtra(w, r, user, map[string]interface{}{
		"passkeys_enabled": len(passkeys) > 0,
		"passkey_count":    len(passkeys),
		"passkeys":         publicPasskeys(passkeys),
	})
	log.Println("DELETE [api/passkeys] " + r.RemoteAddr + ": Deleted passkey and revoked old sessions")
}

func loadPasskeyWebAuthnUser(user models.User) (passkeyWebAuthnUser, error) {
	rows, err := storage.ListPasskeyCredentials(user.ID)
	if err != nil {
		return passkeyWebAuthnUser{}, err
	}
	credentials, err := storage.DecodeWebAuthnCredentials(rows)
	if err != nil {
		return passkeyWebAuthnUser{}, err
	}
	return passkeyWebAuthnUser{user: user, credentials: credentials}, nil
}

func passkeyDiscoverableUserHandler(rawID, userHandle []byte) (webauthn.User, error) {
	row, err := storage.GetPasskeyCredentialByCredentialID(utils.EncodeBase64URL(rawID))
	if err != nil {
		return nil, err
	}
	if row.UserID != string(userHandle) {
		return nil, errors.New("passkey user handle mismatch")
	}
	user, err := storage.GetUserById(row.UserID)
	if err != nil {
		return nil, err
	}
	return loadPasskeyWebAuthnUser(user)
}

func publicPasskeys(passkeys []models.PasskeyCredential) []map[string]interface{} {
	out := make([]map[string]interface{}, 0, len(passkeys))
	for _, passkey := range passkeys {
		item := map[string]interface{}{
			"id":            passkey.ID,
			"name":          passkey.Name,
			"credential_id": passkey.CredentialID,
			"created_at":    passkey.CreatedAt,
		}
		if passkey.LastUsedAt > 0 {
			item["last_used_at"] = passkey.LastUsedAt
		}
		out = append(out, item)
	}
	return out
}

func webAuthnForRequest(r *http.Request) (*webauthn.WebAuthn, error) {
	origin, rpID, err := passkeyOriginAndRPID(r)
	if err != nil {
		return nil, err
	}

	return webauthn.New(&webauthn.Config{
		RPID:                        rpID,
		RPDisplayName:               "MiauInv",
		RPOrigins:                   []string{origin},
		RPTopOrigins:                []string{origin},
		RPTopOriginVerificationMode: protocol.TopOriginExplicitVerificationMode,
		AuthenticatorSelection: protocol.AuthenticatorSelection{
			RequireResidentKey: protocol.ResidentKeyRequired(),
			ResidentKey:        protocol.ResidentKeyRequirementRequired,
			UserVerification:   protocol.VerificationRequired,
		},
	})
}

func passkeyOriginAndRPID(r *http.Request) (string, string, error) {
	origin := strings.TrimSpace(r.Header.Get("Origin"))
	if origin == "" {
		scheme := strings.TrimSpace(r.Header.Get("X-Forwarded-Proto"))
		if scheme == "" {
			scheme = "https"
		}
		host := requestHost(r)
		if host == "" {
			return "", "", errors.New("missing request host")
		}
		origin = scheme + "://" + host
	}

	parsedOrigin, err := url.Parse(origin)
	if err != nil || parsedOrigin.Scheme == "" || parsedOrigin.Host == "" {
		return "", "", errors.New("invalid origin")
	}
	if parsedOrigin.Scheme != "https" && parsedOrigin.Hostname() != "localhost" {
		return "", "", errors.New("passkeys require HTTPS except for localhost")
	}

	originHost := strings.ToLower(parsedOrigin.Hostname())
	allowedHost := strings.ToLower(stripPort(requestHost(r)))
	if allowedHost != "" && originHost != allowedHost {
		return "", "", errors.New("origin host does not match request host")
	}

	return parsedOrigin.Scheme + "://" + parsedOrigin.Host, originHost, nil
}

func requestHost(r *http.Request) string {
	if forwardedHost := strings.TrimSpace(r.Header.Get("X-Forwarded-Host")); forwardedHost != "" {
		parts := strings.Split(forwardedHost, ",")
		return strings.TrimSpace(parts[0])
	}
	return strings.TrimSpace(r.Host)
}

func stripPort(host string) string {
	if host == "" {
		return ""
	}
	if parsedHost, _, err := net.SplitHostPort(host); err == nil {
		return parsedHost
	}
	if strings.Count(host, ":") == 0 {
		return host
	}
	return host
}

func cloneRequestWithJSONBody(r *http.Request, raw json.RawMessage) *http.Request {
	clone := r.Clone(r.Context())
	clone.Body = io.NopCloser(bytes.NewReader(raw))
	clone.ContentLength = int64(len(raw))
	clone.Header = r.Header.Clone()
	clone.Header.Set("Content-Type", "application/json")
	return clone
}