customer/internal/utils/jwt.go

package utils

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/golang-jwt/jwt/v5"
	"penahub.gitlab.yandexcloud.net/pena-services/customer/internal/models"
	jsonUtil "penahub.gitlab.yandexcloud.net/pena-services/customer/pkg/json"
)

type JWT[T any] struct {
	privateKey []byte
	publicKey  []byte
	algorithm  *jwt.SigningMethodRSA
	expiresIn  time.Duration
	issuer     string
	audience   string
}

func NewJWT[T any](configuration *models.JWTConfiguration) *JWT[T] {
	return &JWT[T]{
		privateKey: []byte(configuration.PrivateKey),
		publicKey:  []byte(configuration.PublicKey),
		algorithm:  &configuration.Algorithm,
		expiresIn:  configuration.ExpiresIn,
		issuer:     configuration.Issuer,
		audience:   configuration.Audience,
	}
}

func (receiver *JWT[T]) Create(content *T) (string, error) {
	privateKey, err := jwt.ParseRSAPrivateKeyFromPEM(receiver.privateKey)
	if err != nil {
		return "", fmt.Errorf("failed to parse private key on <Create> of <JWT>: %w", err)
	}

	encoded, err := json.Marshal(content)
	if err != nil {
		return "", fmt.Errorf("failed to encode content to json on <Create> of <JWT>: %w", err)
	}

	now := time.Now().UTC()

	claims := jwt.MapClaims{
		"dat": string(encoded),                    // Our custom data.
		"exp": now.Add(receiver.expiresIn).Unix(), // The expiration time after which the token must be disregarded.
		"aud": receiver.audience,                  // Audience
		"iss": receiver.issuer,                    // Issuer
	}

	token, err := jwt.NewWithClaims(receiver.algorithm, claims).SignedString(privateKey)
	if err != nil {
		return "", fmt.Errorf("failed to sing on <Create> of <JWT>: %w", err)
	}

	return token, nil
}

func (receiver *JWT[T]) Validate(tokenString string) (*T, error) {
	key, err := jwt.ParseRSAPublicKeyFromPEM(receiver.publicKey)
	if err != nil {
		return nil, fmt.Errorf("failed to parse rsa public key on <Validate> of <JWT>: %w", err)
	}

	parseCallback := func(token *jwt.Token) (any, error) {
		if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("unexpected signing method: %s", token.Header["alg"])
		}

		return key, nil
	}

	token, err := jwt.Parse(
		tokenString,
		parseCallback,
		jwt.WithAudience(receiver.audience),
		jwt.WithIssuer(receiver.issuer),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to parse jwt token on <Validate> of <JWT>: %w", err)
	}

	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok || !token.Valid {
		return nil, errors.New("token is invalid on <Validate> of <JWT>")
	}

	data, ok := claims["dat"].(string)
	if !ok {
		return nil, errors.New("data is empty or not a string on <Validate> of <JWT>")
	}

	parsedData, err := jsonUtil.Parse[T](strings.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("failed to parse data on <Validate> of <JWT>: %w", err)
	}

	return parsedData, nil
}