Add a token prefix column, generate the token_prefix and the token_digest, removing the plaintext token from use.

app/models/concerns/token_prefixable.rb

@@ -0,0 +1,53 @@
+module TokenPrefixable
+  extend ActiveSupport::Concern
+
+  class_methods do
+    # Compute HMAC prefix from plaintext token
+    # Returns first 8 chars of Base64url-encoded HMAC
+    # Does NOT reveal anything about the token
+    def compute_token_prefix(plaintext_token)
+      return nil if plaintext_token.blank?
+
+      hmac = OpenSSL::HMAC.digest('SHA256', TokenHmac::KEY, plaintext_token)
+      Base64.urlsafe_encode64(hmac)[0..7]
+    end
+
+    # Find token using HMAC prefix lookup (fast, indexed)
+    def find_by_token(plaintext_token)
+      return nil if plaintext_token.blank?
+
+      prefix = compute_token_prefix(plaintext_token)
+
+      # Fast indexed lookup by HMAC prefix
+      where(token_prefix: prefix).find_each do |token|
+        return token if token.token_matches?(plaintext_token)
+      end
+
+      nil
+    end
+  end
+
+  # Check if a plaintext token matches the hashed token
+  def token_matches?(plaintext_token)
+    return false if plaintext_token.blank? || token_digest.blank?
+
+    BCrypt::Password.new(token_digest) == plaintext_token
+  rescue BCrypt::Errors::InvalidHash
+    false
+  end
+
+  # Generate new token with HMAC prefix
+  # Sets both virtual attribute (for returning to client) and digest (for storage)
+  def generate_token_with_prefix
+    plaintext = SecureRandom.urlsafe_base64(48)
+    self.token_prefix = self.class.compute_token_prefix(plaintext)
+    self.token_digest = BCrypt::Password.create(plaintext)
+
+    # Set the virtual attribute - different models use different names
+    if respond_to?(:plaintext_token=)
+      self.plaintext_token = plaintext  # OidcAccessToken
+    elsif respond_to?(:token=)
+      self.token = plaintext  # OidcRefreshToken
+    end
+  end
+end

app/models/oidc_access_token.rb

@@ -1,12 +1,15 @@
 class OidcAccessToken < ApplicationRecord
+  include TokenPrefixable
+
   belongs_to :application
   belongs_to :user
   has_many :oidc_refresh_tokens, dependent: :destroy
 
-  before_validation :generate_token, on: :create
+  before_validation :generate_token_with_prefix, on: :create
   before_validation :set_expiry, on: :create
 
-  validates :token, uniqueness: true, presence: true
+  validates :token_digest, presence: true
+  validates :token_prefix, presence: true
 
   scope :valid, -> { where("expires_at > ?", Time.current).where(revoked_at: nil) }
   scope :expired, -> { where("expires_at <= ?", Time.current) }
@@ -33,50 +36,11 @@ class OidcAccessToken < ApplicationRecord
     oidc_refresh_tokens.each(&:revoke!)
   end
 
-  # Check if a plaintext token matches the hashed token
-  def token_matches?(plaintext_token)
-    return false if plaintext_token.blank?
-
-    # Use BCrypt to compare if token_digest exists
-    if token_digest.present?
-      BCrypt::Password.new(token_digest) == plaintext_token
-    # Fall back to direct comparison for backward compatibility
-    elsif token.present?
-      token == plaintext_token
-    else
-      false
-    end
-  end
-
-  # Find by token (validates and checks if revoked)
-  def self.find_by_token(plaintext_token)
-    return nil if plaintext_token.blank?
-
-    # Find all non-revoked, non-expired tokens
-    valid.find_each do |access_token|
-      # Use BCrypt to compare (if token_digest exists) or direct comparison
-      if access_token.token_digest.present?
-        return access_token if BCrypt::Password.new(access_token.token_digest) == plaintext_token
-      elsif access_token.token == plaintext_token
-        return access_token
-      end
-    end
-    nil
-  end
+  # find_by_token, token_matches?, and generate_token_with_prefix
+  # are now provided by TokenPrefixable concern
 
   private
 
-  def generate_token
-    return if token.present?
-
-    # Generate opaque access token
-    plaintext = SecureRandom.urlsafe_base64(48)
-    self.plaintext_token = plaintext  # Store temporarily for returning to client
-    self.token_digest = BCrypt::Password.create(plaintext)
-    # Keep token column for backward compatibility during migration
-    self.token = plaintext
-  end
-
   def set_expiry
     self.expires_at ||= application.access_token_expiry
   end

app/models/oidc_refresh_token.rb

@@ -1,10 +1,12 @@
 class OidcRefreshToken < ApplicationRecord
+  include TokenPrefixable
+
   belongs_to :application
   belongs_to :user
   belongs_to :oidc_access_token
   has_many :oidc_access_tokens, foreign_key: :oidc_access_token_id, dependent: :nullify
 
-  before_validation :generate_token, on: :create
+  before_validation :generate_token_with_prefix, on: :create
   before_validation :set_expiry, on: :create
   before_validation :set_token_family_id, on: :create
 
@@ -43,37 +45,11 @@ class OidcRefreshToken < ApplicationRecord
     OidcRefreshToken.in_family(token_family_id).update_all(revoked_at: Time.current)
   end
 
-  # Verify a plaintext token against the stored digest
-  def self.find_by_token(plaintext_token)
-    return nil if plaintext_token.blank?
-
-    # Try to find tokens that could match (we can't search by hash directly)
-    # This is less efficient but necessary with BCrypt
-    # In production, you might want to add a token prefix or other optimization
-    all.find do |refresh_token|
-      refresh_token.token_matches?(plaintext_token)
-    end
-  end
-
-  def token_matches?(plaintext_token)
-    return false if plaintext_token.blank? || token_digest.blank?
-
-    BCrypt::Password.new(token_digest) == plaintext_token
-  rescue BCrypt::Errors::InvalidHash
-    false
-  end
+  # find_by_token, token_matches?, and generate_token_with_prefix
+  # are now provided by TokenPrefixable concern
 
   private
 
-  def generate_token
-    # Generate a secure random token
-    plaintext = SecureRandom.urlsafe_base64(48)
-    self.token = plaintext  # Store temporarily for returning to client
-
-    # Hash it with BCrypt for storage
-    self.token_digest = BCrypt::Password.create(plaintext)
-  end
-
   def set_expiry
     # Use application's configured refresh token TTL
     self.expires_at ||= application.refresh_token_expiry

config/initializers/token_hmac.rb

@@ -0,0 +1,6 @@
+# Token HMAC key derivation
+# This key is used to compute HMAC-based token prefixes for fast lookup
+# Derived from SECRET_KEY_BASE - no storage needed, deterministic output
+module TokenHmac
+  KEY = Rails.application.key_generator.generate_key('oidc_token_prefix', 32)
+end

db/migrate/20251229220739_add_token_prefix_to_tokens.rb

@@ -0,0 +1,42 @@
+class AddTokenPrefixToTokens < ActiveRecord::Migration[8.1]
+  def up
+    add_column :oidc_access_tokens, :token_prefix, :string, limit: 8
+    add_column :oidc_refresh_tokens, :token_prefix, :string, limit: 8
+
+    # Backfill existing tokens with prefix and digest
+    say_with_time "Backfilling token prefixes and digests..." do
+      [OidcAccessToken, OidcRefreshToken].each do |klass|
+        klass.reset_column_information  # Ensure Rails knows about new column
+
+        klass.where(token_prefix: nil).find_each do |token|
+          next unless token.token.present?
+
+          updates = {}
+
+          # Compute HMAC prefix
+          prefix = klass.compute_token_prefix(token.token)
+          updates[:token_prefix] = prefix if prefix.present?
+
+          # Backfill digest if missing
+          if token.token_digest.nil?
+            updates[:token_digest] = BCrypt::Password.create(token.token)
+          end
+
+          token.update_columns(updates) if updates.any?
+        end
+
+        say "  #{klass.name}: #{klass.where.not(token_prefix: nil).count} tokens backfilled"
+      end
+    end
+
+    add_index :oidc_access_tokens, :token_prefix
+    add_index :oidc_refresh_tokens, :token_prefix
+  end
+
+  def down
+    remove_index :oidc_access_tokens, :token_prefix
+    remove_index :oidc_refresh_tokens, :token_prefix
+    remove_column :oidc_access_tokens, :token_prefix
+    remove_column :oidc_refresh_tokens, :token_prefix
+  end
+end

db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[8.1].define(version: 2025_11_25_081147) do
+ActiveRecord::Schema[8.1].define(version: 2025_12_29_220739) do
   create_table "active_storage_attachments", force: :cascade do |t|
     t.bigint "blob_id", null: false
     t.datetime "created_at", null: false
@@ -102,6 +102,7 @@ ActiveRecord::Schema[8.1].define(version: 2025_11_25_081147) do
     t.string "scope"
     t.string "token"
     t.string "token_digest"
+    t.string "token_prefix", limit: 8
     t.datetime "updated_at", null: false
     t.integer "user_id", null: false
     t.index ["application_id", "user_id"], name: "index_oidc_access_tokens_on_application_id_and_user_id"
@@ -110,6 +111,7 @@ ActiveRecord::Schema[8.1].define(version: 2025_11_25_081147) do
     t.index ["revoked_at"], name: "index_oidc_access_tokens_on_revoked_at"
     t.index ["token"], name: "index_oidc_access_tokens_on_token", unique: true
     t.index ["token_digest"], name: "index_oidc_access_tokens_on_token_digest", unique: true
+    t.index ["token_prefix"], name: "index_oidc_access_tokens_on_token_prefix"
     t.index ["user_id"], name: "index_oidc_access_tokens_on_user_id"
   end
 
@@ -143,6 +145,7 @@ ActiveRecord::Schema[8.1].define(version: 2025_11_25_081147) do
     t.string "scope"
     t.string "token_digest", null: false
     t.integer "token_family_id"
+    t.string "token_prefix", limit: 8
     t.datetime "updated_at", null: false
     t.integer "user_id", null: false
     t.index ["application_id", "user_id"], name: "index_oidc_refresh_tokens_on_application_id_and_user_id"
@@ -152,6 +155,7 @@ ActiveRecord::Schema[8.1].define(version: 2025_11_25_081147) do
     t.index ["revoked_at"], name: "index_oidc_refresh_tokens_on_revoked_at"
     t.index ["token_digest"], name: "index_oidc_refresh_tokens_on_token_digest", unique: true
     t.index ["token_family_id"], name: "index_oidc_refresh_tokens_on_token_family_id"
+    t.index ["token_prefix"], name: "index_oidc_refresh_tokens_on_token_prefix"
     t.index ["user_id"], name: "index_oidc_refresh_tokens_on_user_id"
   end