byom-core-subscription/internal/common/jwt/jwt.go

package jwt

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"net/url"
	"time"

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

// JWTClient handles JWT operations for email verification
type JWTClient struct {
	secret []byte
}

// EmailVerificationClaims contains claims specific to email verification tokens
type EmailVerificationClaims struct {
	Email string `json:"email"`
	jwt.RegisteredClaims
}

func NewJWTClient(secret []byte) *JWTClient {
	return &JWTClient{secret: secret}
}

// GenerateEmailVerificationToken creates a JWT token for email verification
// The token includes the email in both the claims and subject for additional security
// and expires after 24 hours
func (j *JWTClient) GenerateEmailVerificationToken(email string) (string, error) {
	uniqueID, err := generateUniqueID()
	if err != nil {
		return "", fmt.Errorf("failed to generate unique ID: %w", err)
	}

	claims := EmailVerificationClaims{
		Email: email,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().UTC().Add(24 * time.Hour)),
			IssuedAt:  jwt.NewNumericDate(time.Now().UTC()),
			NotBefore: jwt.NewNumericDate(time.Now().UTC()),
			Subject:   email,
			ID:        uniqueID,
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString(j.secret)
}

func (j *JWTClient) GenerateEmailVerificationURL(email, baseURL string) (string, error) {
	token, err := j.GenerateEmailVerificationToken(email)
	if err != nil {
		return "", fmt.Errorf("failed to generate token: %w", err)
	}

	// Ensure the token is URL-safe
	return fmt.Sprintf("%s?token=%s", baseURL, url.QueryEscape(token)), nil
}

func (j *JWTClient) VerifyEmailToken(tokenString string) (string, error) {
	token, err := jwt.ParseWithClaims(
		tokenString,
		&EmailVerificationClaims{},
		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 j.secret, nil
		},
	)

	if err != nil {
		return "", fmt.Errorf("invalid token: %w", err)
	}

	if !token.Valid {
		return "", errors.New("token is not valid")
	}

	claims, ok := token.Claims.(*EmailVerificationClaims)
	if !ok {
		return "", errors.New("invalid claims type")
	}

	if claims.Subject != claims.Email {
		return "", errors.New("token subject does not match email")
	}

	return claims.Email, nil
}

func generateUniqueID() (string, error) {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	return hex.EncodeToString(bytes), nil
}