mpc/pkg/api/auth.go

package api

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"slices"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"github.com/hanzoai/orm"
	"golang.org/x/crypto/bcrypt"

	"github.com/luxfi/mpc/pkg/db"
)

type JWTClaims struct {
	UserID string `json:"user_id"`
	OrgID  string `json:"org_id"`
	Role   string `json:"role"`
	jwt.RegisteredClaims
}

func (s *Server) generateJWT(userID, orgID, role string) (string, error) {
	claims := JWTClaims{
		UserID: userID,
		OrgID:  orgID,
		Role:   role,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    "mpc.lux.network",
			Audience:  jwt.ClaimStrings{"mpc-api"},
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(24 * time.Hour)),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
		},
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString(s.jwtSecret)
}

func (s *Server) generateRefreshToken(userID, orgID, role string) (string, error) {
	claims := JWTClaims{
		UserID: userID,
		OrgID:  orgID,
		Role:   role,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    "mpc.lux.network",
			Audience:  jwt.ClaimStrings{"mpc-refresh"},
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(7 * 24 * time.Hour)),
			IssuedAt:  jwt.NewNumericDate(time.Now()),
		},
	}
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString(s.jwtSecret)
}

func (s *Server) validateJWT(tokenStr string) (*JWTClaims, error) {
	token, err := jwt.ParseWithClaims(tokenStr, &JWTClaims{}, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return s.jwtSecret, nil
	})
	if err != nil {
		return nil, err
	}
	claims, ok := token.Claims.(*JWTClaims)
	if !ok || !token.Valid {
		return nil, fmt.Errorf("invalid token")
	}
	if claims.Issuer != "mpc.lux.network" {
		return nil, fmt.Errorf("invalid token issuer")
	}
	if !slices.Contains([]string(claims.Audience), "mpc-api") {
		return nil, fmt.Errorf("invalid token audience")
	}
	return claims, nil
}

func (s *Server) validateRefreshToken(tokenStr string) (*JWTClaims, error) {
	token, err := jwt.ParseWithClaims(tokenStr, &JWTClaims{}, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return s.jwtSecret, nil
	})
	if err != nil {
		return nil, err
	}
	claims, ok := token.Claims.(*JWTClaims)
	if !ok || !token.Valid {
		return nil, fmt.Errorf("invalid token")
	}
	if claims.Issuer != "mpc.lux.network" {
		return nil, fmt.Errorf("invalid token issuer")
	}
	if !slices.Contains([]string(claims.Audience), "mpc-refresh") {
		return nil, fmt.Errorf("invalid token audience")
	}
	return claims, nil
}

func hashPassword(password string) (string, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return "", err
	}
	return string(hash), nil
}

func checkPassword(hash, password string) bool {
	return bcrypt.CompareHashAndPassword([]byte(hash), []byte(password)) == nil
}

func generateAPIKeyToken() (string, string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", "", err
	}
	key := "lux_" + hex.EncodeToString(b)
	hash := sha256.Sum256([]byte(key))
	return key, hex.EncodeToString(hash[:]), nil
}

func (s *Server) validateAPIKey(ctx context.Context, key string) (*db.APIKey, error) {
	hash := sha256.Sum256([]byte(key))
	keyHash := hex.EncodeToString(hash[:])
	prefix := key
	if len(prefix) > 8 {
		prefix = prefix[:8]
	}

	apiKey, err := orm.TypedQuery[db.APIKey](s.db.ORM).
		Filter("keyHash=", keyHash).
		Filter("keyPrefix=", prefix).
		First()
	if err != nil {
		return nil, fmt.Errorf("invalid api key")
	}

	// Check expiry
	if apiKey.ExpiresAt != nil && apiKey.ExpiresAt.Before(time.Now()) {
		return nil, fmt.Errorf("api key expired")
	}

	// Update last used
	now := time.Now()
	apiKey.LastUsedAt = &now
	apiKey.Update()

	return apiKey, nil
}