golib/jwt/generator.go

package jwt

import (
	"context"
	"database/sql"
	"errors"
	"time"

	pkgerrors "github.com/pkg/errors"
)

type TokenGenerator struct {
	accessExpireAfter  int64         // Access Token expiry time in minutes
	refreshExpireAfter int64         // Refresh Token expiry time in minutes
	freshExpireAfter   int64         // Token freshness expiry time in minutes
	trustedHost        string        // Trusted hostname to use for the tokens
	secretKey          string        // Secret key to use for token hashing
	db                 *sql.DB       // Database connection for token blacklisting
	tableConfig        TableConfig   // Table configuration
	tableManager       *TableManager // Table lifecycle manager
}

// GeneratorConfig holds configuration for creating a TokenGenerator.
type GeneratorConfig struct {
	// AccessExpireAfter is the access token expiry time in minutes.
	AccessExpireAfter int64

	// RefreshExpireAfter is the refresh token expiry time in minutes.
	RefreshExpireAfter int64

	// FreshExpireAfter is the token freshness expiry time in minutes.
	FreshExpireAfter int64

	// TrustedHost is the trusted hostname to use for the tokens.
	TrustedHost string

	// SecretKey is the secret key to use for token hashing.
	SecretKey string

	// DB is the database connection. Can be nil to disable token revocation.
	// When using ORMs like GORM or Bun, pass the underlying *sql.DB.
	DB *sql.DB

	// DBType specifies the database type and version for proper table management.
	// Only required if DB is not nil.
	DBType DatabaseType

	// TableConfig configures the blacklist table name and behavior.
	// Only required if DB is not nil.
	TableConfig TableConfig
}

// CreateGenerator creates and returns a new TokenGenerator using the provided configuration.
func CreateGenerator(config GeneratorConfig) (gen *TokenGenerator, err error) {
	if config.AccessExpireAfter <= 0 {
		return nil, errors.New("accessExpireAfter must be greater than 0")
	}
	if config.RefreshExpireAfter <= 0 {
		return nil, errors.New("refreshExpireAfter must be greater than 0")
	}
	if config.FreshExpireAfter <= 0 {
		return nil, errors.New("freshExpireAfter must be greater than 0")
	}
	if config.TrustedHost == "" {
		return nil, errors.New("trustedHost cannot be an empty string")
	}
	if config.SecretKey == "" {
		return nil, errors.New("secretKey cannot be an empty string")
	}

	var tableManager *TableManager
	if config.DB != nil {
		// Create table manager
		tableManager = NewTableManager(config.DB, config.DBType, config.TableConfig)

		// Create table if AutoCreate is enabled
		if config.TableConfig.AutoCreate {
			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			defer cancel()

			err = tableManager.CreateTable(ctx)
			if err != nil {
				return nil, pkgerrors.Wrap(err, "failed to create blacklist table")
			}
		}

		// Setup automatic cleanup if enabled
		if config.TableConfig.EnableAutoCleanup {
			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			defer cancel()

			err = tableManager.SetupAutoCleanup(ctx)
			if err != nil {
				return nil, pkgerrors.Wrap(err, "failed to setup automatic cleanup")
			}
		}
	}

	return &TokenGenerator{
		accessExpireAfter:  config.AccessExpireAfter,
		refreshExpireAfter: config.RefreshExpireAfter,
		freshExpireAfter:   config.FreshExpireAfter,
		trustedHost:        config.TrustedHost,
		secretKey:          config.SecretKey,
		db:                 config.DB,
		tableConfig:        config.TableConfig,
		tableManager:       tableManager,
	}, nil
}

// Cleanup manually removes expired tokens from the blacklist table.
// This method should be called periodically if automatic cleanup is not enabled,
// or can be called on-demand regardless of automatic cleanup settings.
func (gen *TokenGenerator) Cleanup(ctx context.Context) error {
	if gen.db == nil {
		return errors.New("No DB provided, unable to use this function")
	}

	tableName := gen.tableConfig.TableName
	currentTime := time.Now().Unix()

	query := "DELETE FROM " + tableName + " WHERE exp < ?"

	_, err := gen.db.ExecContext(ctx, query, currentTime)
	if err != nil {
		return pkgerrors.Wrap(err, "failed to cleanup expired tokens")
	}

	return nil
}