From d851eee065785e3798aff1967ea85e69ec020e2d Mon Sep 17 00:00:00 2001
From: Paulo Soares <psoares@americastg.com>
Date: Thu, 5 Nov 2015 08:24:43 -0200
Subject: [PATCH] limiting the number of concurrent devices

---
 app/models/devise_token_auth/concerns/user.rb       |  6 ++++++
 lib/devise_token_auth/engine.rb                     |  2 ++
 .../templates/devise_token_auth.rb                  |  4 ++++
 test/controllers/demo_user_controller_test.rb       | 13 +++++++++++++
 test/test_helper.rb                                 | 12 ++++++++----
 5 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/app/models/devise_token_auth/concerns/user.rb b/app/models/devise_token_auth/concerns/user.rb
index 4ec8053be..e8416d41e 100644
--- a/app/models/devise_token_auth/concerns/user.rb
+++ b/app/models/devise_token_auth/concerns/user.rb
@@ -178,6 +178,12 @@ def create_new_auth_token(client_id=nil)
       last_token: last_token,
       updated_at: Time.now
     }
+	
+    max_clients = DeviseTokenAuth.max_number_of_devices
+    while self.tokens.keys.length > 0 and max_clients < self.tokens.keys.length
+      oldest_token = self.tokens.min_by { |cid, v| v[:expiry] || v["expiry"] }
+      self.tokens.delete(oldest_token.first)
+    end	
 
     self.save!
 
diff --git a/lib/devise_token_auth/engine.rb b/lib/devise_token_auth/engine.rb
index aa0441c11..a88bbaa54 100644
--- a/lib/devise_token_auth/engine.rb
+++ b/lib/devise_token_auth/engine.rb
@@ -10,6 +10,7 @@ class Engine < ::Rails::Engine
   end
 
   mattr_accessor :change_headers_on_each_request,
+                 :max_number_of_devices,
                  :token_lifespan,
                  :batch_request_buffer_throttle,
                  :omniauth_prefix,
@@ -20,6 +21,7 @@ class Engine < ::Rails::Engine
                  :enable_standard_devise_support
 
   self.change_headers_on_each_request       = true
+  self.max_number_of_devices                = 10
   self.token_lifespan                       = 2.weeks
   self.batch_request_buffer_throttle        = 5.seconds
   self.omniauth_prefix                      = '/omniauth'
diff --git a/lib/generators/devise_token_auth/templates/devise_token_auth.rb b/lib/generators/devise_token_auth/templates/devise_token_auth.rb
index 6169a8bf7..5c2018df9 100644
--- a/lib/generators/devise_token_auth/templates/devise_token_auth.rb
+++ b/lib/generators/devise_token_auth/templates/devise_token_auth.rb
@@ -9,6 +9,10 @@
   # determines how long tokens will remain valid after they are issued.
   # config.token_lifespan = 2.weeks
 
+  # Sets the max number of concurrent devices per user, which is 10 by default.
+  # After this limit is reached, the oldest tokens will be removed.
+  # config.max_number_of_devices = 10
+
   # Sometimes it's necessary to make several requests to the API at the same
   # time. In this case, each request in the batch will need to share the same
   # auth token. This setting determines how far apart the requests can be while
diff --git a/test/controllers/demo_user_controller_test.rb b/test/controllers/demo_user_controller_test.rb
index 899d11a84..d10b12a25 100644
--- a/test/controllers/demo_user_controller_test.rb
+++ b/test/controllers/demo_user_controller_test.rb
@@ -322,6 +322,19 @@ class DemoUserControllerTest < ActionDispatch::IntegrationTest
           it 'should not define current_mang' do
             refute_equal @resource, @controller.current_mang
           end
+		  
+		  
+          it 'should increase the number of tokens by a factor of 2 up to 11' do
+            @first_token = @resource.tokens.keys.first
+
+            DeviseTokenAuth.max_number_of_devices = 11
+            (1..10).each do |n|
+              assert_equal [11, 2*n].min, @resource.reload.tokens.keys.length
+              get '/demo/members_only', {}, nil
+            end
+
+            assert_not_includes @resource.reload.tokens.keys, @first_token
+          end
         end
 
         it 'should return success status' do
diff --git a/test/test_helper.rb b/test/test_helper.rb
index da3a528b5..3b9f88871 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -40,13 +40,17 @@ class ActiveSupport::TestCase
   # Add more helper methods to be used by all tests here...
 
   def age_token(user, client_id)
-    user.tokens[client_id]['updated_at'] = Time.now - (DeviseTokenAuth.batch_request_buffer_throttle + 10.seconds)
-    user.save!
+    if user.tokens[client_id]
+      user.tokens[client_id]['updated_at'] = Time.now - (DeviseTokenAuth.batch_request_buffer_throttle + 10.seconds)
+      user.save!
+    end
   end
 
   def expire_token(user, client_id)
-    user.tokens[client_id]['expiry'] = (Time.now - (DeviseTokenAuth.token_lifespan.to_f + 10.seconds)).to_i
-    user.save!
+    if user.tokens[client_id]
+      user.tokens[client_id]['expiry'] = (Time.now - (DeviseTokenAuth.token_lifespan.to_f + 10.seconds)).to_i
+      user.save!
+    end
   end
 end