Refactor custom token attributes & some other code

CHANGELOG.md

@@ -8,6 +8,8 @@ User-visible changes worth mentioning.
 ## main
 
 - [#ID] Add your PR description here.
+- [#1602] Allow custom data to be stored inside access grants/tokens.
+- [#1603] Code refactoring for custom token attributes.
 
 # 5.6.4
 

app/controllers/doorkeeper/authorizations_controller.rb

@@ -88,7 +88,7 @@ def pre_auth_params
     end
 
     def pre_auth_param_fields
-      Doorkeeper.configuration.custom_access_token_attributes + %i[
+      custom_access_token_attributes + %i[
         client_id
         code_challenge
         code_challenge_method
@@ -100,6 +100,10 @@ def pre_auth_param_fields
       ]
     end
 
+    def custom_access_token_attributes
+      Doorkeeper.config.custom_access_token_attributes.map(&:to_sym)
+    end
+
     def authorization
       @authorization ||= strategy.request
     end

lib/doorkeeper/config/validations.rb

@@ -11,6 +11,7 @@ def validate!
         validate_reuse_access_token_value
         validate_token_reuse_limit
         validate_secret_strategies
+        validate_custom_access_token_attributes
       end
 
       private
@@ -49,14 +50,16 @@ def validate_token_reuse_limit
         @token_reuse_limit = 100
       end
 
+      # Validate that the access_token and access_grant models
+      # both respond to all of the custom attributes
       def validate_custom_access_token_attributes
-        # Validate that the access_token and access_grant models
-        # both respond to all of the custom attributes
-        Doorkeeper.config.custom_access_token_attributes.each do |attribute_name|
-          [Doorkeeper.config.access_token_model, Doorkeeper.config.access_grant_model].each do |model|
-            unless model.has_attribute?(attribute_name)
-              raise NotImplementedError, "#{model} does not recognize custom attribute: #{attribute_name}."
-            end
+        return if custom_access_token_attributes.blank?
+
+        custom_access_token_attributes.each do |attribute_name|
+          [access_token_model, access_grant_model].each do |model|
+            next if model.has_attribute?(attribute_name)
+
+            raise Doorkeeper::Errors::ConfigError, "#{model} does not recognize custom attribute: #{attribute_name}."
           end
         end
       end

lib/doorkeeper/errors.rb

@@ -44,6 +44,7 @@ def initialize(response)
     UnableToGenerateToken = Class.new(DoorkeeperError)
     TokenGeneratorNotFound = Class.new(DoorkeeperError)
     NoOrmCleaner = Class.new(DoorkeeperError)
+    ConfigError = Class.new(DoorkeeperError)
 
     InvalidToken = Class.new(BaseResponseError)
     TokenExpired = Class.new(InvalidToken)

lib/doorkeeper/oauth/authorization/code.rb

@@ -45,13 +45,13 @@ def access_grant_attributes
             attributes[:resource_owner_id] = resource_owner.id
           end
 
+          pkce_attributes.merge(attributes).merge(custom_attributes)
+        end
+
+        def custom_attributes
           # Custom access token attributes are saved into the access grant,
           # and then included in subsequently generated access tokens.
-          Doorkeeper.config.custom_access_token_attributes.each do |attribute_name|
-            attributes[attribute_name] = @pre_auth.custom_access_token_attributes[attribute_name]
-          end
-
-          pkce_attributes.merge(attributes)
+          @pre_auth.custom_access_token_attributes.to_h.with_indifferent_access
         end
 
         def pkce_attributes

lib/doorkeeper/oauth/authorization_code_request.rb

@@ -56,11 +56,12 @@ def pkce_supported?
       end
 
       def validate_params
-        @missing_param = if grant&.uses_pkce? && code_verifier.blank?
-                           :code_verifier
-                         elsif redirect_uri.blank?
-                           :redirect_uri
-                         end
+        @missing_param =
+          if grant&.uses_pkce? && code_verifier.blank?
+            :code_verifier
+          elsif redirect_uri.blank?
+            :redirect_uri
+          end
 
         @missing_param.nil?
       end
@@ -98,11 +99,15 @@ def validate_code_verifier
       end
 
       def generate_code_challenge(code_verifier)
-        server_config.access_grant_model.generate_code_challenge(code_verifier)
+        Doorkeeper.config.access_grant_model.generate_code_challenge(code_verifier)
       end
 
-      private def custom_token_attributes_with_data
-        grant.attributes.with_indifferent_access.slice(*Doorkeeper.config.custom_access_token_attributes).symbolize_keys
+      def custom_token_attributes_with_data
+        grant
+          .attributes
+          .with_indifferent_access
+          .slice(*Doorkeeper.config.custom_access_token_attributes)
+          .symbolize_keys
       end
     end
   end

lib/doorkeeper/oauth/base_request.rb

@@ -28,8 +28,7 @@ def scopes
 
       def find_or_create_access_token(client, resource_owner, scopes, custom_attributes, server)
         context = Authorization::Token.build_context(client, grant_type, scopes, resource_owner)
-        token_model = server_config.access_token_model
-        application = client.is_a?(server_config.application_model) ? client : client&.application
+        application = client.is_a?(Doorkeeper.config.application_model) ? client : client&.application
 
         token_attributes = {
           application: application,
@@ -39,19 +38,16 @@ def find_or_create_access_token(client, resource_owner, scopes, custom_attribute
           use_refresh_token: Authorization::Token.refresh_token_enabled?(server, context),
         }
 
-        @access_token = token_model.find_or_create_for(**token_attributes.merge(custom_attributes))
+        @access_token =
+          Doorkeeper.config.access_token_model.find_or_create_for(**token_attributes.merge(custom_attributes))
       end
 
       def before_successful_response
-        server_config.before_successful_strategy_response.call(self)
+        Doorkeeper.config.before_successful_strategy_response.call(self)
       end
 
       def after_successful_response
-        server_config.after_successful_strategy_response.call(self, @response)
-      end
-
-      def server_config
-        Doorkeeper.config
+        Doorkeeper.config.after_successful_strategy_response.call(self, @response)
       end
 
       private

lib/doorkeeper/oauth/client_credentials/creator.rb

@@ -9,12 +9,12 @@ def call(client, scopes, attributes = {})
 
           if lookup_existing_token?
             existing_token = find_active_existing_token_for(client, scopes)
-            return existing_token if server_config.reuse_access_token && existing_token&.reusable?
+            return existing_token if Doorkeeper.config.reuse_access_token && existing_token&.reusable?
           end
 
           with_revocation(existing_token: existing_token) do
-            application = client.is_a?(server_config.application_model) ? client : client&.application
-            server_config.access_token_model.create_for(
+            application = client.is_a?(Doorkeeper.config.application_model) ? client : client&.application
+            Doorkeeper.config.access_token_model.create_for(
               application: application,
               resource_owner: nil,
               scopes: scopes,
@@ -26,7 +26,7 @@ def call(client, scopes, attributes = {})
         private
 
         def with_revocation(existing_token:)
-          if existing_token && server_config.revoke_previous_client_credentials_token?
+          if existing_token && Doorkeeper.config.revoke_previous_client_credentials_token?
             existing_token.with_lock do
               raise Errors::DoorkeeperError, :invalid_token_reuse if existing_token.revoked?
 
@@ -40,16 +40,12 @@ def with_revocation(existing_token:)
         end
 
         def lookup_existing_token?
-          server_config.reuse_access_token ||
-            server_config.revoke_previous_client_credentials_token?
+          Doorkeeper.config.reuse_access_token ||
+            Doorkeeper.config.revoke_previous_client_credentials_token?
         end
 
         def find_active_existing_token_for(client, scopes)
-          server_config.access_token_model.matching_token_for(client, nil, scopes, include_expired: false)
-        end
-
-        def server_config
-          Doorkeeper.config
+          Doorkeeper.config.access_token_model.matching_token_for(client, nil, scopes, include_expired: false)
         end
       end
     end

lib/doorkeeper/oauth/password_access_token_request.rb

@@ -68,7 +68,7 @@ def validate_client
       end
 
       def validate_client_supports_grant_flow
-        server_config.allow_grant_flow_for_client?(grant_type, client&.application)
+        Doorkeeper.config.allow_grant_flow_for_client?(grant_type, client&.application)
       end
     end
   end

lib/doorkeeper/oauth/pre_authorization.rb

@@ -21,21 +21,17 @@ class PreAuthorization
                   :authorization_response_flow, :response_mode, :custom_access_token_attributes
 
       def initialize(server, parameters = {}, resource_owner = nil)
-        @server                = server
-        @client_id             = parameters[:client_id]
-        @response_type         = parameters[:response_type]
-        @response_mode         = parameters[:response_mode]
-        @redirect_uri          = parameters[:redirect_uri]
-        @scope                 = parameters[:scope]
-        @state                 = parameters[:state]
-        @code_challenge        = parameters[:code_challenge]
+        @server = server
+        @client_id = parameters[:client_id]
+        @response_type = parameters[:response_type]
+        @response_mode = parameters[:response_mode]
+        @redirect_uri = parameters[:redirect_uri]
+        @scope = parameters[:scope]
+        @state = parameters[:state]
+        @code_challenge = parameters[:code_challenge]
         @code_challenge_method = parameters[:code_challenge_method]
-        @resource_owner        = resource_owner
-
-        @custom_access_token_attributes = {}
-        Doorkeeper.config.custom_access_token_attributes.each do |field|
-          @custom_access_token_attributes[field] = parameters[field]
-        end
+        @resource_owner = resource_owner
+        @custom_access_token_attributes = parameters.slice(*Doorkeeper.config.custom_access_token_attributes)
       end
 
       def authorizable?

lib/doorkeeper/oauth/refresh_token_request.rb

@@ -26,7 +26,7 @@ def initialize(server, refresh_token, credentials, parameters = {})
       private
 
       def load_client(credentials)
-        server_config.application_model.by_uid_and_secret(credentials.uid, credentials.secret)
+        Doorkeeper.config.application_model.by_uid_and_secret(credentials.uid, credentials.secret)
       end
 
       def before_successful_response
@@ -41,7 +41,7 @@ def before_successful_response
       end
 
       def refresh_token_revoked_on_use?
-        server_config.access_token_model.refresh_token_revoked_on_use?
+        Doorkeeper.config.access_token_model.refresh_token_revoked_on_use?
       end
 
       def default_scopes
@@ -75,7 +75,7 @@ def create_access_token
         # Here we assume that TTL of the token received after refreshing should be
         # the same as that of the original token.
         #
-        @access_token = server_config.access_token_model.create_for(
+        @access_token = Doorkeeper.config.access_token_model.create_for(
           application: refresh_token.application,
           resource_owner: resource_owner,
           scopes: scopes,

spec/dummy/db/migrate/20230205064514_add_custom_attributes.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+class AddCustomAttributes < ActiveRecord::Migration[4.2]
+  def change
+    add_column :oauth_access_grants, :tenant_name, :string
+    add_column :oauth_access_tokens, :tenant_name, :string
+  end
+end

spec/dummy/db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 20180210183654) do
+ActiveRecord::Schema.define(version: 20230205064514) do
 
   create_table "oauth_access_grants", force: :cascade do |t|
     t.integer "resource_owner_id", null: false
@@ -22,6 +22,7 @@
     t.datetime "created_at", null: false
     t.datetime "revoked_at"
     t.string "scopes"
+    t.string "tenant_name"
     unless ENV["WITHOUT_PKCE"]
       t.string   "code_challenge"
       t.string   "code_challenge_method"
@@ -40,6 +41,7 @@
     t.datetime "created_at", null: false
     t.string "scopes"
     t.string "previous_refresh_token", default: "", null: false
+    t.string "tenant_name"
     t.index ["refresh_token"], name: "index_oauth_access_tokens_on_refresh_token", unique: true
     t.index ["resource_owner_id"], name: "index_oauth_access_tokens_on_resource_owner_id"
     t.index ["token"], name: "index_oauth_access_tokens_on_token", unique: true

spec/lib/config_spec.rb

@@ -529,9 +529,18 @@ class ApplicationWithOwner < ActiveRecord::Base
     it "can change the value" do
       Doorkeeper.configure do
         orm DOORKEEPER_ORM
-        custom_access_token_attributes [:added_field_1, :added_field_2]
+        custom_access_token_attributes [:tenant_name]
       end
-      expect(config.custom_access_token_attributes).to eq([:added_field_1, :added_field_2])
+      expect(config.custom_access_token_attributes).to eq([:tenant_name])
+    end
+
+    it "validates custom attributes to be present in the models" do
+      expect do
+        Doorkeeper.configure do
+          orm DOORKEEPER_ORM
+          custom_access_token_attributes [:none_existing_column]
+        end
+      end.to raise_error(Doorkeeper::Errors::ConfigError)
     end
   end
 

spec/lib/oauth/code_response_spec.rb

@@ -3,17 +3,18 @@
 require "spec_helper"
 
 RSpec.describe Doorkeeper::OAuth::CodeResponse do
+  let(:application) { FactoryBot.create(:application, scopes: "") }
+  let(:owner) { FactoryBot.build_stubbed(:resource_owner) }
   let(:pre_auth) do
-    application = FactoryBot.create(:application, scopes: "")
     double(
       :pre_auth,
       client: application,
       redirect_uri: "http://tst.com/cb",
       state: "state",
       scopes: Doorkeeper::OAuth::Scopes.from_string("public"),
+      custom_access_token_attributes: {},
     )
   end
-  let(:owner) { FactoryBot.build_stubbed(:resource_owner) }
 
   describe "#body" do
     subject(:body) { described_class.new(pre_auth, auth).body }

spec/requests/flows/authorization_code_spec.rb

@@ -156,7 +156,8 @@ def authorize(redirect_url)
 
     authorization_code = Doorkeeper::AccessGrant.first.token
     page.driver.post token_endpoint_url(
-      code: authorization_code, client_id: @client.uid,
+      code: authorization_code,
+      client_id: @client.uid,
       redirect_uri: @client.redirect_uri,
     )
 
@@ -174,7 +175,8 @@ def authorize(redirect_url)
 
     authorization_code = Doorkeeper::AccessGrant.first.token
     page.driver.post token_endpoint_url(
-      code: authorization_code, client_secret: @client.secret,
+      code: authorization_code,
+      client_secret: @client.secret,
       redirect_uri: @client.redirect_uri,
     )
 
@@ -550,4 +552,31 @@ def authorize(redirect_url)
       end
     end
   end
+
+  context "when custom_access_token_attributes are configured" do
+    let(:resource_owner) { FactoryBot.create(:resource_owner) }
+    let(:client) { client_exists }
+    let(:grant) do
+      authorization_code_exists(
+         application: client,
+         resource_owner_id: resource_owner.id,
+         resource_owner_type: resource_owner.class.name,
+         tenant_name: "Tenant 1",
+       )
+    end
+
+    before do
+      Doorkeeper.configure do
+        orm DOORKEEPER_ORM
+        custom_access_token_attributes [:tenant_name]
+      end
+    end
+
+    it "copies custom attributes from the grant into the token" do
+      page.driver.post token_endpoint_url(code: grant.token, client: client)
+
+      access_token = Doorkeeper::AccessToken.find_by(token: json_response["access_token"])
+      expect(access_token.tenant_name).to eq("Tenant 1")
+    end
+  end
 end