kalzu-value-store/auth/jwt.go

package auth

import (
	"fmt"
	"time"

	"github.com/golang-jwt/jwt/v4"
)

// JWT signing key (should be configurable in production)
var jwtSigningKey = []byte("your-secret-signing-key-change-this-in-production")

// JWTClaims represents the custom claims for our JWT tokens
type JWTClaims struct {
	UserUUID string   `json:"user_uuid"`
	Scopes   []string `json:"scopes"`
	jwt.RegisteredClaims
}

// GenerateJWT creates a new JWT token for a user with specified scopes
func GenerateJWT(userUUID string, scopes []string, expirationHours int) (string, int64, error) {
	if expirationHours <= 0 {
		expirationHours = 1 // Default to 1 hour
	}

	now := time.Now()
	expiresAt := now.Add(time.Duration(expirationHours) * time.Hour)

	claims := JWTClaims{
		UserUUID: userUUID,
		Scopes:   scopes,
		RegisteredClaims: jwt.RegisteredClaims{
			IssuedAt:  jwt.NewNumericDate(now),
			ExpiresAt: jwt.NewNumericDate(expiresAt),
			Issuer:    "kvs-server",
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenString, err := token.SignedString(jwtSigningKey)
	if err != nil {
		return "", 0, err
	}

	return tokenString, expiresAt.Unix(), nil
}

// ValidateJWT validates a JWT token and returns the claims if valid
func ValidateJWT(tokenString string) (*JWTClaims, error) {
	token, err := jwt.ParseWithClaims(tokenString, &JWTClaims{}, func(token *jwt.Token) (interface{}, error) {
		// Validate signing method
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		return jwtSigningKey, nil
	})

	if err != nil {
		return nil, err
	}

	if claims, ok := token.Claims.(*JWTClaims); ok && token.Valid {
		return claims, nil
	}

	return nil, fmt.Errorf("invalid token")
}