veza/veza-backend-api/internal/services/jwt_revocation_service.go

package services

// JWTRevocationService — DB-backed jti revocation ledger.
// v1.0.10 sécu item 6.
//
// Companion to TokenBlacklist (Redis) :
//   * Redis = fast path, low-latency lookup on every request.
//   * DB    = durable record that survives Redis restarts, plus
//             admin-targetable revocation ("invalidate session jti=X").
//
// Write path : both written. Read path : Redis first, DB on miss.
// The middleware uses IsRevoked(ctx, jti) which encapsulates this.
//
// The blacklist exists too — these don't supersede each other.
// TokenBlacklist hashes the whole token (works without parsing claims) ;
// this service uses the parsed jti (works after claims have been
// validated). Both are checked ; either match = reject.

import (
	"context"
	"errors"
	"time"

	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"
	"gorm.io/gorm"
)

// JWTRevocation is the GORM model for a single revocation row.
type JWTRevocation struct {
	JTI        string     `gorm:"primaryKey;size:64" json:"jti"`
	UserID     uuid.UUID  `gorm:"type:uuid;not null;index" json:"user_id"`
	ExpiresAt  time.Time  `gorm:"not null;index" json:"expires_at"`
	RevokedAt  time.Time  `gorm:"not null;default:now()" json:"revoked_at"`
	Reason     string     `gorm:"size:255" json:"reason,omitempty"`
	RevokedBy  *uuid.UUID `gorm:"type:uuid" json:"revoked_by,omitempty"`
}

// TableName overrides default pluralisation.
func (JWTRevocation) TableName() string { return "jwt_revocations" }

// JWTRevocationService writes + checks per-jti revocations.
type JWTRevocationService struct {
	db        *gorm.DB
	redis     *redis.Client
	logger    *zap.Logger
	keyPrefix string
}

// NewJWTRevocationService constructs the service. redisClient may be
// nil — when absent, IsRevoked falls through to a DB-only check (still
// correct, just slower on the hot path).
func NewJWTRevocationService(db *gorm.DB, redisClient *redis.Client, logger *zap.Logger) *JWTRevocationService {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &JWTRevocationService{
		db:        db,
		redis:     redisClient,
		logger:    logger,
		keyPrefix: "jwt_revoked:",
	}
}

// Revoke records a revocation. Writes Redis first (so an immediately-
// following request rejects) then DB (durable record). A failure on
// either side is logged but doesn't undo the other — the read path
// tolerates one being stale relative to the other.
//
// expiresAt should be the token's natural expiry. The cleanup cron
// drops rows where expiresAt < now() - safety margin so the table
// doesn't grow unboundedly.
func (s *JWTRevocationService) Revoke(
	ctx context.Context,
	jti string,
	userID uuid.UUID,
	expiresAt time.Time,
	reason string,
	revokedBy *uuid.UUID,
) error {
	if jti == "" {
		return errors.New("jti required")
	}

	// Redis first (best-effort cache).
	if s.redis != nil {
		ttl := time.Until(expiresAt)
		if ttl > 0 {
			if err := s.redis.Set(ctx, s.keyPrefix+jti, "1", ttl).Err(); err != nil {
				s.logger.Warn("Redis revocation cache write failed; relying on DB",
					zap.String("jti", jti), zap.Error(err))
			}
		}
	}

	// DB write — the durable record. Idempotent via PRIMARY KEY (jti) :
	// a second Revoke call updates the existing row's reason instead
	// of erroring on UNIQUE.
	row := JWTRevocation{
		JTI:        jti,
		UserID:     userID,
		ExpiresAt:  expiresAt,
		RevokedAt:  time.Now().UTC(),
		Reason:     reason,
		RevokedBy:  revokedBy,
	}
	res := s.db.WithContext(ctx).
		Where("jti = ?", jti).
		FirstOrCreate(&row)
	if res.Error != nil {
		return res.Error
	}
	s.logger.Info("JWT revoked",
		zap.String("jti", jti),
		zap.String("user_id", userID.String()),
		zap.String("reason", reason),
	)
	return nil
}

// IsRevoked reports whether the given jti has been revoked. Redis-
// fast-path, DB-fallback. A DB error is treated as "not revoked"
// (fail-open) because the alternative — every request fails on a DB
// blip — is worse for availability than a brief revocation gap.
// The TokenBlacklist (token-hash check) provides a second line of
// defense in the middleware, so a missed jti revocation here doesn't
// mean the token is fully accepted.
func (s *JWTRevocationService) IsRevoked(ctx context.Context, jti string) (bool, error) {
	if jti == "" {
		return false, nil
	}
	// Redis fast path.
	if s.redis != nil {
		exists, err := s.redis.Exists(ctx, s.keyPrefix+jti).Result()
		if err == nil {
			if exists > 0 {
				return true, nil
			}
			// Cache says no — but Redis cache may have evicted on
			// memory pressure or restart. Fall through to DB.
		}
		// Redis error : log + fall through. Don't fail the request
		// on a Redis blip.
	}
	// DB lookup. Only checks rows whose tokens haven't expired yet —
	// an expired-but-still-in-table row has nothing to revoke.
	var count int64
	err := s.db.WithContext(ctx).
		Model(&JWTRevocation{}).
		Where("jti = ? AND expires_at > ?", jti, time.Now().UTC()).
		Count(&count).Error
	if err != nil {
		s.logger.Warn("DB revocation check failed; treating as not-revoked (fail-open)",
			zap.String("jti", jti), zap.Error(err))
		return false, nil
	}
	if count > 0 {
		// Backfill Redis so the next request hits the cache. Non-
		// fatal on failure.
		if s.redis != nil {
			// Re-derive TTL from the row.
			var row JWTRevocation
			if dberr := s.db.WithContext(ctx).Where("jti = ?", jti).First(&row).Error; dberr == nil {
				ttl := time.Until(row.ExpiresAt)
				if ttl > 0 {
					_ = s.redis.Set(ctx, s.keyPrefix+jti, "1", ttl).Err()
				}
			}
		}
		return true, nil
	}
	return false, nil
}

// ListByUser returns the active revocations for a user (rows whose
// underlying token wouldn't have naturally expired yet). Used by the
// "your active sessions" admin/user UI.
func (s *JWTRevocationService) ListByUser(ctx context.Context, userID uuid.UUID, limit int) ([]JWTRevocation, error) {
	if limit <= 0 || limit > 200 {
		limit = 50
	}
	var rows []JWTRevocation
	err := s.db.WithContext(ctx).
		Where("user_id = ? AND expires_at > ?", userID, time.Now().UTC()).
		Order("revoked_at DESC").
		Limit(limit).
		Find(&rows).Error
	return rows, err
}

// PurgeExpired drops rows whose underlying tokens have already
// expired. Run from a daily cron / Forgejo workflow. Returns the
// number of rows deleted so the caller can log / metric it.
func (s *JWTRevocationService) PurgeExpired(ctx context.Context, safetyMargin time.Duration) (int64, error) {
	cutoff := time.Now().UTC().Add(-safetyMargin)
	res := s.db.WithContext(ctx).
		Where("expires_at < ?", cutoff).
		Delete(&JWTRevocation{})
	if res.Error != nil {
		return 0, res.Error
	}
	if res.RowsAffected > 0 {
		s.logger.Info("Purged expired JWT revocations",
			zap.Int64("rows", res.RowsAffected),
			zap.Time("cutoff", cutoff),
		)
	}
	return res.RowsAffected, nil
}