Add ActiveRecord::Base::generates_token_for (rails#44189)

Currently, `signed_id` fulfills the role of generating tokens for e.g.
resetting a password.  However, signed IDs cannot reflect record state,
so if a token is intended to be single-use, it must be tracked in a
database at least until it expires.

With `generates_token_for`, a token can embed data from a record.  When
using the token to fetch the record, the data from the token and the
data from the record will be compared.  If the two do not match, the
token will be treated as invalid, the same as if it had expired.  For
example:

```ruby
class User < ActiveRecord::Base
  has_secure_password

  generates_token_for :password_reset do
    # BCrypt salt changes when password is updated
    BCrypt::Password.new(password_digest).salt[-10..]
  end
end

user = User.first
token = user.generate_token_for(:password_reset)

User.find_by_token_for(:password_reset, token) # => user

user.update!(password: "new password")
User.find_by_token_for(:password_reset, token) # => nil
```

Author: jonathanhefner

activerecord/lib/active_record.rb

@@ -77,6 +77,7 @@ module ActiveRecord
   autoload :TestDatabases
   autoload :TestFixtures, "active_record/fixtures"
   autoload :Timestamp
+  autoload :TokenFor
   autoload :TouchLater
   autoload :Transactions
   autoload :Translation

activerecord/lib/active_record/base.rb

@@ -325,6 +325,7 @@ class Base
     include Serialization
     include Store
     include SecureToken
+    include TokenFor
     include SignedId
     include Suppressor
     include Encryption::EncryptableRecord

activerecord/lib/active_record/railtie.rb

@@ -334,6 +334,14 @@ class Railtie < Rails::Railtie # :nodoc:
       end
     end
 
+    initializer "active_record.generated_token_verifier" do
+      config.after_initialize do |app|
+        ActiveSupport.on_load(:active_record) do
+          self.generated_token_verifier ||= app.message_verifier("active_record/token_for")
+        end
+      end
+    end
+
     initializer "active_record_encryption.configuration" do |app|
       ActiveRecord::Encryption.configure \
          primary_key: app.credentials.dig(:active_record_encryption, :primary_key),

activerecord/lib/active_record/token_for.rb

@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/object/json"
+
+module ActiveRecord
+  module TokenFor
+    extend ActiveSupport::Concern
+
+    included do
+      class_attribute :token_definitions, instance_accessor: false, instance_predicate: false, default: {}
+      class_attribute :generated_token_verifier, instance_accessor: false, instance_predicate: false
+    end
+
+    # :nodoc:
+    TokenDefinition = Struct.new(:defining_class, :purpose, :expires_in, :block) do
+      def full_purpose
+        @full_purpose ||= [defining_class.name, purpose, expires_in].join("\n")
+      end
+
+      def message_verifier
+        defining_class.generated_token_verifier
+      end
+
+      def payload_for(model)
+        block ? [model.id, model.instance_eval(&block).as_json] : [model.id]
+      end
+
+      def generate_token(model)
+        message_verifier.generate(payload_for(model), expires_in: expires_in, purpose: full_purpose)
+      end
+
+      def resolve_token(token)
+        payload = message_verifier.verified(token, purpose: full_purpose)
+        model = yield(payload[0]) if payload
+        model if model && payload_for(model) == payload
+      end
+    end
+
+    module ClassMethods
+      # Defines the behavior of tokens generated for a specific +purpose+.
+      # A token can be generated by calling TokenFor#generate_token_for on a
+      # record. Later, that record can be fetched by calling #find_by_token_for
+      # (or #find_by_token_for!) with the same purpose and token.
+      #
+      # Tokens are signed so that they are tamper-proof. Thus they can be
+      # exposed to outside world as, for example, password reset tokens.
+      #
+      # By default, tokens do not expire. They can be configured to expire by
+      # specifying a duration via the +expires_in+ option. The duration becomes
+      # part of the token's signature, so changing the value of +expires_in+
+      # will automatically invalidate previously generated tokens.
+      #
+      # A block may also be specified. When generating a token with
+      # TokenFor#generate_token_for, the block will be evaluated in the context
+      # of the record, and its return value will be embedded in the token as
+      # JSON. Later, when fetching the record with #find_by_token_for, the block
+      # will be evaluated again in the context of the fetched record. If the two
+      # JSON values do not match, the token will be treated as invalid. Note
+      # that the value returned by the block <strong>should not contain
+      # sensitive information</strong> because it will be embedded in the token
+      # as <strong>human-readable plaintext JSON</strong>.
+      #
+      # ==== Examples
+      #
+      #   class User < ActiveRecord::Base
+      #     has_secure_password
+      #
+      #     generates_token_for :password_reset, expires_in: 15.minutes do
+      #       # Last 10 characters of password salt, which changes when password is updated:
+      #       BCrypt::Password.new(password_digest).salt[-10..]
+      #     end
+      #   end
+      #
+      #   user = User.first
+      #
+      #   token = user.generate_token_for(:password_reset)
+      #   User.find_by_token_for(:password_reset, token) # => user
+      #   # 16 minutes later...
+      #   User.find_by_token_for(:password_reset, token) # => nil
+      #
+      #   token = user.generate_token_for(:password_reset)
+      #   User.find_by_token_for(:password_reset, token) # => user
+      #   user.update!(password: "new password")
+      #   User.find_by_token_for(:password_reset, token) # => nil
+      def generates_token_for(purpose, expires_in: nil, &block)
+        self.token_definitions = token_definitions.merge(purpose => TokenDefinition.new(self, purpose, expires_in, block))
+      end
+
+      # Finds a record using a given +token+ for a predefined +purpose+. Returns
+      # +nil+ if the token is invalid or the record was not found.
+      def find_by_token_for(purpose, token)
+        raise UnknownPrimaryKey.new(self) unless primary_key
+        token_definitions.fetch(purpose).resolve_token(token) { |id| find_by(primary_key => id) }
+      end
+
+      # Finds a record using a given +token+ for a predefined +purpose+. Raises
+      # ActiveSupport::MessageVerifier::InvalidSignature if the token is invalid
+      # (e.g. expired, bad format, etc). Raises ActiveRecord::RecordNotFound if
+      # the token is valid but the record was not found.
+      def find_by_token_for!(purpose, token)
+        token_definitions.fetch(purpose).resolve_token(token) { |id| find(id) } ||
+          (raise ActiveSupport::MessageVerifier::InvalidSignature)
+      end
+    end
+
+    # Generates a token for a predefined +purpose+.
+    #
+    # Use ClassMethods::generates_token_for to define a token purpose and
+    # behavior.
+    def generate_token_for(purpose)
+      self.class.token_definitions.fetch(purpose).generate_token(self)
+    end
+  end
+end

activerecord/test/cases/token_for_test.rb

@@ -0,0 +1,156 @@
+# frozen_string_literal: true
+
+require "cases/helper"
+require "models/matey"
+require "models/user"
+require "active_support/message_verifier"
+
+class TokenForTest < ActiveRecord::TestCase
+  class User < ::User
+    generates_token_for :lookup
+
+    generates_token_for :password_reset, expires_in: 15.minutes do
+      password_digest.to_s[-(31 + 22), 10] # first 10 characters of BCrypt salt
+    end
+
+    generates_token_for :snapshot do
+      { updated_at: updated_at }
+    end
+  end
+
+  setup do
+    @original_verifier = ActiveRecord::Base.generated_token_verifier
+    ActiveRecord::Base.generated_token_verifier = ActiveSupport::MessageVerifier.new("secret")
+
+    @user = User.create!(password_digest: "$2a$4$#{"x" * 22}#{"y" * 31}")
+    @lookup_token = @user.generate_token_for(:lookup)
+    @password_reset_token = @user.generate_token_for(:password_reset)
+  end
+
+  teardown do
+    ActiveRecord::Base.generated_token_verifier = @original_verifier
+  end
+
+  test "finds record by token" do
+    assert_equal @user, User.find_by_token_for(:lookup, @lookup_token)
+    assert_equal @user, User.find_by_token_for!(:lookup, @lookup_token)
+  end
+
+  test "returns nil when record is not found" do
+    @user.destroy
+    assert_nil User.find_by_token_for(:lookup, @lookup_token)
+  end
+
+  test "raises on bang when record is not found" do
+    @user.destroy
+    assert_raises(ActiveRecord::RecordNotFound) do
+      User.find_by_token_for!(:lookup, @lookup_token)
+    end
+  end
+
+  test "raises when token definition does not exist" do
+    assert_raises { User.find_by_token_for(:bad, @lookup_token) }
+  end
+
+  test "does not find record when token is invalid" do
+    assert_nil User.find_by_token_for(:lookup, "bad")
+    assert_raises(ActiveSupport::MessageVerifier::InvalidSignature) do
+      User.find_by_token_for!(:lookup, "bad")
+    end
+  end
+
+  test "does not find record when token is for a different purpose" do
+    assert_nil User.find_by_token_for(:password_reset, @lookup_token)
+    assert_raises(ActiveSupport::MessageVerifier::InvalidSignature) do
+      User.find_by_token_for!(:password_reset, @lookup_token)
+    end
+  end
+
+  test "finds record when token has not expired and embedded data has not changed" do
+    assert_equal @user, User.find_by_token_for(:password_reset, @password_reset_token)
+  end
+
+  test "does not find record when token has expired" do
+    travel 1.day
+    assert_nil User.find_by_token_for(:password_reset, @password_reset_token)
+    assert_raises(ActiveSupport::MessageVerifier::InvalidSignature) do
+      User.find_by_token_for!(:password_reset, @password_reset_token)
+    end
+  end
+
+  test "tokens do not expire by default" do
+    travel 1000.years
+    assert_equal @user, User.find_by_token_for(:lookup, @lookup_token)
+  end
+
+  test "does not find record when expires_in is different" do
+    User.generates_token_for :lookup, expires_in: 1.year
+
+    assert_nil User.find_by_token_for(:lookup, @lookup_token)
+    new_lookup_token = @user.generate_token_for(:lookup)
+    assert_equal @user, User.find_by_token_for(:lookup, new_lookup_token)
+  ensure
+    User.generates_token_for :lookup
+  end
+
+  test "does not find record when embedded data is different" do
+    @user.update!(password: "new password")
+    assert_nil User.find_by_token_for(:password_reset, @password_reset_token)
+    assert_raises(ActiveSupport::MessageVerifier::InvalidSignature) do
+      User.find_by_token_for!(:password_reset, @password_reset_token)
+    end
+  end
+
+  test "supports JSON-serializable embedded data" do
+    snapshot_token = @user.generate_token_for(:snapshot)
+    assert_equal @user, User.find_by_token_for(:snapshot, snapshot_token)
+    @user.touch
+    assert_nil User.find_by_token_for(:snapshot, snapshot_token)
+  end
+
+  test "finds record through relation" do
+    assert_equal @user, User.where("1=1").find_by_token_for(:lookup, @lookup_token)
+    assert_nil User.where("1=0").find_by_token_for(:lookup, @lookup_token)
+  end
+
+  test "finds record through subclass" do
+    subclass = Class.new(User)
+    subclassed_user = subclass.find_by_token_for(:lookup, @lookup_token)
+
+    assert_instance_of subclass, subclassed_user
+    assert_equal @user.id, subclassed_user.id
+  end
+
+  test "subclasses can redefine tokens" do
+    subclass = Class.new(User) do
+      generates_token_for :lookup
+    end
+    subclassed_user = subclass.find(@user.id)
+    subclassed_lookup_token = subclassed_user.generate_token_for(:lookup)
+
+    assert_equal subclassed_user, subclass.find_by_token_for(:lookup, subclassed_lookup_token)
+    assert_nil subclass.find_by_token_for(:lookup, @lookup_token)
+    assert_nil User.find_by_token_for(:lookup, subclassed_lookup_token)
+  end
+
+  test "finds record with a custom primary key" do
+    custom_pk = Class.new(User) do
+      self.primary_key = "auth_token"
+    end
+    custom_pk_user = custom_pk.find(@user.auth_token)
+    custom_pk_lookup_token = custom_pk_user.generate_token_for(:lookup)
+
+    assert_equal custom_pk_user, custom_pk.find_by_token_for(:lookup, custom_pk_lookup_token)
+    assert_nil custom_pk.find_by_token_for(:lookup, @lookup_token)
+  end
+
+  test "raises when no primary key has been declared" do
+    no_pk = Class.new(Matey) do
+      generates_token_for :parley
+    end
+
+    assert_raises(ActiveRecord::UnknownPrimaryKey) do
+      no_pk.find_by_token_for(:parley, "this token will not be checked")
+    end
+  end
+end

activerecord/test/schema/schema.rb

@@ -1291,6 +1291,7 @@
     t.string :auth_token
     t.string :password_digest
     t.string :recovery_password_digest
+    t.timestamps null: true
   end
 
   create_table :test_with_keyword_column_name, force: true do |t|