Prevent reuse of OTP code with ROTP after option (#99)

ROTP reference:
https://github.com/mdp/rotp/blob/2be11551ac4982670d526ceef4a5db985d0f626a/README.md#preventing-reuse-of-time-based-otps

README.md

@@ -98,9 +98,31 @@ sleep 30 # lets wait again
 user.authenticate_otp('186522', drift: 60) # => true
 ```
 
+### Preventing reuse of Time based OTP's
+
+By keeping track of the last time a user's OTP was verified, we can prevent token reuse during the interval window (default 30 seconds). It is useful with SMS, that is commonly used in combination with `drift` to extend the life of the code.
+
+```ruby
+rails g migration AddLastOtpAtToUsers last_otp_at:integer
+=>
+      invoke  active_record
+      create    db/migrate/20220407010931_add_last_otp_at_to_users.rb
+```
+
+```ruby
+class User < ApplicationRecord
+  has_one_time_password after_column_name: :last_otp_at
+end
+```
+
+```ruby
+user.authenticate_otp('186522') # => true
+user.authenticate_otp('186522') # => false
+```
+
 ## Counter based OTP
 
-An additonal counter field is required in our ``User`` Model
+An additional counter field is required in our ``User`` Model
 
 ```ruby
 rails g migration AddCounterForOtpToUsers otp_counter:integer
@@ -213,7 +235,7 @@ user.provisioning_uri(nil, issuer: 'MYAPP') #=> 'otpauth://totp/hello@heapsource
 
 This can then be rendered as a QR Code which can be scanned and added to the users list of OTP credentials.
 
-### Setting up a customer interval 
+### Setting up a customer interval
 
 If you define a custom interval for TOTP codes, just as `has_one_time_password interval: 10` (for example), remember to include the interval also in `provisioning_uri` method. If not defined, the default value is 30 seconds (according to ROTP gem: https://github.com/mdp/rotp/blob/master/lib/rotp/totp.rb#L9)
 

lib/active_model/one_time_password.rb

@@ -13,7 +13,7 @@ module OneTimePassword
     module ClassMethods
       def has_one_time_password(options = {})
         cattr_accessor :otp_column_name, :otp_counter_column_name,
-                       :otp_backup_codes_column_name
+                       :otp_backup_codes_column_name, :otp_after_column_name
         class_attribute :otp_digits, :otp_counter_based,
                         :otp_backup_codes_count, :otp_one_time_backup_codes,
                         :otp_interval
@@ -29,6 +29,7 @@ def has_one_time_password(options = {})
           options[:counter_column_name] || OTP_DEFAULT_COUNTER_COLUMN_NAME
         ).to_s
         self.otp_interval = options[:interval]
+        self.otp_after_column_name = options[:after_column_name]
         self.otp_backup_codes_column_name = (
           options[:backup_codes_column_name] ||
           OTP_DEFAULT_BACKUP_CODES_COLUMN_NAME
@@ -167,13 +168,20 @@ def authenticate_totp(code, options = {})
           digits: otp_digits,
           interval: otp_interval
         )
-        if (drift = options[:drift])
-          totp.verify(code, drift_behind: drift)
-        else
-          totp.verify(code)
+        otp_after = if otp_after_column_name_enabled?
+                      public_send(otp_after_column_name)
+                    end
+        totp.verify(code, drift_behind: options[:drift] || 0, after: otp_after)
+            .tap do |updated_last_otp_at|
+          updated_last_otp_at && otp_after_column_name_enabled? &&
+            update(otp_after_column_name => updated_last_otp_at)
         end
       end
 
+      def otp_after_column_name_enabled?
+        otp_after_column_name && respond_to?(otp_after_column_name)
+      end
+
       def hotp_code(options = {})
         if options[:auto_increment]
           self.otp_counter += 1

test/models/after_user.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class AfterUser < ActiveRecord::Base
+  has_one_time_password after_column_name: :last_otp_at
+end

test/one_time_password_test.rb

@@ -23,6 +23,10 @@ def setup
     @opt_in = OptInTwoFactor.new
     @opt_in.email = '[email protected]'
     @opt_in.run_callbacks :create
+
+    @after_user = AfterUser.new
+    @after_user.email = '[email protected]'
+    @after_user.run_callbacks :create
   end
 
   def test_authenticate_with_otp
@@ -88,6 +92,18 @@ def test_authenticate_with_otp_when_drift_is_allowed
     assert_equal true, @visitor.authenticate_otp(code, drift: 60)
   end
 
+  def test_authenticate_with_otp_when_after_is_allowed
+    code = @user.otp_code
+    assert_equal true, @user.authenticate_otp(code)
+    assert_equal true, @user.authenticate_otp(code)
+
+    code = @after_user.otp_code
+    assert_equal true, @after_user.authenticate_otp(code)
+    assert_equal false, @after_user.authenticate_otp(code)
+    assert_equal false, @after_user.authenticate_otp('1111111')
+    assert_equal false, @after_user.authenticate_otp(code)
+  end
+
   def test_authenticate_with_backup_code
     backup_code = @user.public_send(@user.otp_backup_codes_column_name).first
     assert_equal true, @user.authenticate_otp(backup_code)

test/schema.rb

@@ -24,4 +24,13 @@
     t.string :otp_secret_key
     t.timestamps
   end
+
+  create_table :after_users, force: true do |t|
+    t.string :key
+    t.string :email
+    t.integer :otp_counter
+    t.string :otp_secret_key
+    t.integer :last_otp_at
+    t.timestamps
+  end
 end