From d28ca8c4c9d3a36ac9d722161ed6001a0e3f0062 Mon Sep 17 00:00:00 2001
From: Markus Koller <markus-koller@gmx.ch>
Date: Tue, 25 Oct 2016 15:38:03 +0200
Subject: [PATCH] feature: Store and return nonces in IdToken responses

- Add migration generator for nonce table
- Abort if Doorkeeper ORM isn't set to ActiveRecord
- Use prepend to extend Doorkeeper classes
---
 lib/doorkeeper/openid_connect.rb              | 80 ++++---------------
 lib/doorkeeper/openid_connect/config.rb       |  6 ++
 .../openid_connect/helpers/controller.rb      |  2 +
 .../openid_connect/models/id_token.rb         |  8 +-
 .../oauth/authorization/code.rb               | 20 +++++
 .../oauth/authorization_code_request.rb       | 17 ++++
 .../oauth/password_access_token_request.rb    | 28 +++++++
 .../openid_connect/oauth/pre_authorization.rb | 20 +++++
 .../openid_connect/oauth/token_response.rb    | 25 ++++++
 .../openid_connect/orm/active_record.rb       | 21 +++++
 .../orm/active_record/access_grant.rb         | 16 ++++
 .../openid_connect/orm/active_record/nonce.rb | 17 ++++
 .../openid_connect/migration_generator.rb     | 15 ++++
 .../openid_connect/templates/migration.rb     | 14 ++++
 .../discovery_controller_spec.rb              |  0
 .../userinfo_controller_spec.rb               |  0
 ...create_doorkeeper_openid_connect_tables.rb | 14 ++++
 spec/dummy/db/schema.rb                       | 23 ++++--
 spec/factories.rb                             | 21 +++--
 .../openid_connect => }/config_spec.rb        | 21 ++++-
 spec/lib/oauth/authorization/code_spec.rb     | 35 ++++++++
 .../oauth/authorization_code_request_spec.rb  | 22 +++++
 .../password_access_token_request_spec.rb     | 27 +++++++
 spec/lib/oauth/pre_authorization_spec.rb      | 13 +++
 spec/lib/oauth/token_response_spec.rb         | 33 ++++++++
 .../{doorkeeper => }/openid_connect_spec.rb   |  0
 .../openid_connect => }/routes_spec.rb        |  0
 spec/models/access_grant_spec.rb              | 15 ++++
 spec/models/id_token_spec.rb                  | 12 ++-
 spec/models/nonce_spec.rb                     | 42 ++++++++++
 30 files changed, 480 insertions(+), 87 deletions(-)
 create mode 100644 lib/doorkeeper/openid_connect/oauth/authorization/code.rb
 create mode 100644 lib/doorkeeper/openid_connect/oauth/authorization_code_request.rb
 create mode 100644 lib/doorkeeper/openid_connect/oauth/password_access_token_request.rb
 create mode 100644 lib/doorkeeper/openid_connect/oauth/pre_authorization.rb
 create mode 100644 lib/doorkeeper/openid_connect/oauth/token_response.rb
 create mode 100644 lib/doorkeeper/openid_connect/orm/active_record.rb
 create mode 100644 lib/doorkeeper/openid_connect/orm/active_record/access_grant.rb
 create mode 100644 lib/doorkeeper/openid_connect/orm/active_record/nonce.rb
 create mode 100644 lib/generators/doorkeeper/openid_connect/migration_generator.rb
 create mode 100644 lib/generators/doorkeeper/openid_connect/templates/migration.rb
 rename spec/controllers/{doorkeeper/openid_connect => }/discovery_controller_spec.rb (100%)
 rename spec/controllers/{doorkeeper/openid_connect => }/userinfo_controller_spec.rb (100%)
 create mode 100644 spec/dummy/db/migrate/20161031135842_create_doorkeeper_openid_connect_tables.rb
 rename spec/lib/{doorkeeper/openid_connect => }/config_spec.rb (78%)
 create mode 100644 spec/lib/oauth/authorization/code_spec.rb
 create mode 100644 spec/lib/oauth/authorization_code_request_spec.rb
 create mode 100644 spec/lib/oauth/password_access_token_request_spec.rb
 create mode 100644 spec/lib/oauth/pre_authorization_spec.rb
 create mode 100644 spec/lib/oauth/token_response_spec.rb
 rename spec/lib/{doorkeeper => }/openid_connect_spec.rb (100%)
 rename spec/lib/{doorkeeper/openid_connect => }/routes_spec.rb (100%)
 create mode 100644 spec/models/access_grant_spec.rb
 create mode 100644 spec/models/nonce_spec.rb

diff --git a/lib/doorkeeper/openid_connect.rb b/lib/doorkeeper/openid_connect.rb
index 2d90ede..c9f2b91 100644
--- a/lib/doorkeeper/openid_connect.rb
+++ b/lib/doorkeeper/openid_connect.rb
@@ -1,5 +1,10 @@
-require 'doorkeeper/openid_connect/version'
+require 'doorkeeper'
+require 'json/jwt'
+
+require 'doorkeeper/openid_connect/claims_builder'
+require 'doorkeeper/openid_connect/config'
 require 'doorkeeper/openid_connect/engine'
+require 'doorkeeper/openid_connect/version'
 
 require 'doorkeeper/openid_connect/helpers/controller'
 
@@ -8,80 +13,25 @@
 require 'doorkeeper/openid_connect/models/claims/claim'
 require 'doorkeeper/openid_connect/models/claims/normal_claim'
 
-require 'doorkeeper/openid_connect/claims_builder'
-require 'doorkeeper/openid_connect/config'
+require 'doorkeeper/openid_connect/oauth/authorization/code'
+require 'doorkeeper/openid_connect/oauth/authorization_code_request'
+require 'doorkeeper/openid_connect/oauth/password_access_token_request'
+require 'doorkeeper/openid_connect/oauth/pre_authorization'
+require 'doorkeeper/openid_connect/oauth/token_response'
 
-require 'doorkeeper/openid_connect/rails/routes'
+require 'doorkeeper/openid_connect/orm/active_record'
 
-require 'doorkeeper'
-require 'json/jwt'
+require 'doorkeeper/openid_connect/rails/routes'
 
 module Doorkeeper
+  singleton_class.send :prepend, OpenidConnect::DoorkeeperConfiguration
+
   module OpenidConnect
     # TODO: make this configurable
     SIGNING_ALGORITHM = 'RS256'
 
-    def self.configured?
-      @config.present?
-    end
-
-    def self.installed?
-      configured?
-    end
-
     def self.signing_key
       JSON::JWK.new(OpenSSL::PKey.read(configuration.jws_private_key))
     end
   end
 end
-
-module Doorkeeper
-  class << self
-    prepend ::Doorkeeper::OpenidConnect::DoorkeeperConfiguration
-  end
-
-  module Helpers::Controller
-    prepend ::Doorkeeper::OpenidConnect::Helpers::Controller
-  end
-end
-
-module Doorkeeper
-  module OAuth
-    class PasswordAccessTokenRequest
-      private
-
-      def after_successful_response
-        id_token = Doorkeeper::OpenidConnect::Models::IdToken.new(access_token)
-        @response.id_token = id_token
-      end
-    end
-  end
-end
-
-module Doorkeeper
-  module OAuth
-    class AuthorizationCodeRequest
-      private
-
-      def after_successful_response
-        id_token = Doorkeeper::OpenidConnect::Models::IdToken.new(access_token)
-        @response.id_token = id_token
-      end
-    end
-  end
-end
-
-module Doorkeeper
-  module OAuth
-    class TokenResponse
-      attr_accessor :id_token
-      alias_method :original_body, :body
-
-      def body
-        original_body.
-          merge({:id_token => id_token.try(:as_jws_token)}).
-          reject { |_, value| value.blank? }
-      end
-    end
-  end
-end
diff --git a/lib/doorkeeper/openid_connect/config.rb b/lib/doorkeeper/openid_connect/config.rb
index 3c8bfd2..046e491 100644
--- a/lib/doorkeeper/openid_connect/config.rb
+++ b/lib/doorkeeper/openid_connect/config.rb
@@ -3,10 +3,16 @@ module OpenidConnect
     module DoorkeeperConfiguration
       def configure(&block)
         super(&block)
+
+        if configuration.orm != :active_record
+          fail ConfigurationError, 'Doorkeeper OpenID Connect currently only supports the ActiveRecord ORM adapter'
+        end
+
         configuration.optional_scopes.add :openid
       end
     end
 
+    class ConfigurationError < StandardError; end
     class MissingConfiguration < StandardError
       def initialize
         super('Configuration for Doorkeeper OpenID Connect missing. Do you have doorkeeper_openid_connect initializer?')
diff --git a/lib/doorkeeper/openid_connect/helpers/controller.rb b/lib/doorkeeper/openid_connect/helpers/controller.rb
index f6b3938..e2d9bfa 100644
--- a/lib/doorkeeper/openid_connect/helpers/controller.rb
+++ b/lib/doorkeeper/openid_connect/helpers/controller.rb
@@ -27,4 +27,6 @@ def prompt_values
       end
     end
   end
+
+  Helpers::Controller.send :prepend, OpenidConnect::Helpers::Controller
 end
diff --git a/lib/doorkeeper/openid_connect/models/id_token.rb b/lib/doorkeeper/openid_connect/models/id_token.rb
index 42174ec..d7719f8 100644
--- a/lib/doorkeeper/openid_connect/models/id_token.rb
+++ b/lib/doorkeeper/openid_connect/models/id_token.rb
@@ -4,8 +4,11 @@ module Models
       class IdToken
         include ActiveModel::Validations
 
-        def initialize(access_token)
+        attr_reader :nonce
+
+        def initialize(access_token, nonce = nil)
           @access_token = access_token
+          @nonce = nonce
           @resource_owner = access_token.instance_eval(&Doorkeeper::OpenidConnect.configuration.resource_owner_from_access_token)
           @issued_at = Time.now
         end
@@ -16,7 +19,8 @@ def claims
             sub: subject,
             aud: audience,
             exp: expiration,
-            iat: issued_at
+            iat: issued_at,
+            nonce: nonce
           }
         end
 
diff --git a/lib/doorkeeper/openid_connect/oauth/authorization/code.rb b/lib/doorkeeper/openid_connect/oauth/authorization/code.rb
new file mode 100644
index 0000000..7fec7d2
--- /dev/null
+++ b/lib/doorkeeper/openid_connect/oauth/authorization/code.rb
@@ -0,0 +1,20 @@
+module Doorkeeper
+  module OpenidConnect
+    module OAuth
+      module Authorization
+        module Code
+          def issue_token
+            super.tap do |access_grant|
+              ::Doorkeeper::OpenidConnect::Nonce.create!(
+                access_grant: access_grant,
+                nonce: pre_auth.nonce
+              )
+            end
+          end
+        end
+      end
+    end
+  end
+
+  OAuth::Authorization::Code.send :prepend, OpenidConnect::OAuth::Authorization::Code
+end
diff --git a/lib/doorkeeper/openid_connect/oauth/authorization_code_request.rb b/lib/doorkeeper/openid_connect/oauth/authorization_code_request.rb
new file mode 100644
index 0000000..27f0890
--- /dev/null
+++ b/lib/doorkeeper/openid_connect/oauth/authorization_code_request.rb
@@ -0,0 +1,17 @@
+module Doorkeeper
+  module OpenidConnect
+    module OAuth
+      module AuthorizationCodeRequest
+        private
+
+        def after_successful_response
+          super
+          id_token = Doorkeeper::OpenidConnect::Models::IdToken.new(access_token, grant.openid_connect_nonce.use!)
+          @response.id_token = id_token
+        end
+      end
+    end
+  end
+
+  OAuth::AuthorizationCodeRequest.send :prepend, OpenidConnect::OAuth::AuthorizationCodeRequest
+end
diff --git a/lib/doorkeeper/openid_connect/oauth/password_access_token_request.rb b/lib/doorkeeper/openid_connect/oauth/password_access_token_request.rb
new file mode 100644
index 0000000..2043a1c
--- /dev/null
+++ b/lib/doorkeeper/openid_connect/oauth/password_access_token_request.rb
@@ -0,0 +1,28 @@
+module Doorkeeper
+  module OpenidConnect
+    module OAuth
+      module PasswordAccessTokenRequest
+        def self.prepended(base)
+          base.class_eval do
+            attr_reader :nonce
+          end
+        end
+
+        def initialize(server, client, resource_owner, parameters = {})
+          super
+          @nonce = parameters[:nonce]
+        end
+
+        private
+
+        def after_successful_response
+          super
+          id_token = Doorkeeper::OpenidConnect::Models::IdToken.new(access_token, nonce)
+          @response.id_token = id_token
+        end
+      end
+    end
+  end
+
+  OAuth::PasswordAccessTokenRequest.send :prepend, OpenidConnect::OAuth::PasswordAccessTokenRequest
+end
diff --git a/lib/doorkeeper/openid_connect/oauth/pre_authorization.rb b/lib/doorkeeper/openid_connect/oauth/pre_authorization.rb
new file mode 100644
index 0000000..cdd6ac3
--- /dev/null
+++ b/lib/doorkeeper/openid_connect/oauth/pre_authorization.rb
@@ -0,0 +1,20 @@
+module Doorkeeper
+  module OpenidConnect
+    module OAuth
+      module PreAuthorization
+        def self.prepended(base)
+          base.class_eval do
+            attr_reader :nonce
+          end
+        end
+
+        def initialize(server, client, attrs = {})
+          super
+          @nonce = attrs[:nonce]
+        end
+      end
+    end
+  end
+
+  OAuth::PreAuthorization.send :prepend, OpenidConnect::OAuth::PreAuthorization
+end
diff --git a/lib/doorkeeper/openid_connect/oauth/token_response.rb b/lib/doorkeeper/openid_connect/oauth/token_response.rb
new file mode 100644
index 0000000..9049822
--- /dev/null
+++ b/lib/doorkeeper/openid_connect/oauth/token_response.rb
@@ -0,0 +1,25 @@
+module Doorkeeper
+  module OpenidConnect
+    module OAuth
+      module TokenResponse
+        def self.prepended(base)
+          base.class_eval do
+            attr_accessor :id_token
+          end
+        end
+
+        def body
+          if token.includes_scope? 'openid'
+            super.
+              merge({:id_token => id_token.try(:as_jws_token)}).
+              reject { |_, value| value.blank? }
+          else
+            super
+          end
+        end
+      end
+    end
+  end
+
+  OAuth::TokenResponse.send :prepend, OpenidConnect::OAuth::TokenResponse
+end
diff --git a/lib/doorkeeper/openid_connect/orm/active_record.rb b/lib/doorkeeper/openid_connect/orm/active_record.rb
new file mode 100644
index 0000000..e759af3
--- /dev/null
+++ b/lib/doorkeeper/openid_connect/orm/active_record.rb
@@ -0,0 +1,21 @@
+module Doorkeeper
+  module OpenidConnect
+    module Orm
+      module ActiveRecord
+        def initialize_models!
+          super
+          require 'doorkeeper/openid_connect/orm/active_record/access_grant'
+          require 'doorkeeper/openid_connect/orm/active_record/nonce'
+
+          if Doorkeeper.configuration.active_record_options[:establish_connection]
+            [Doorkeeper::OpenidConnect::Nonce].each do |c|
+              c.send :establish_connection, Doorkeeper.configuration.active_record_options[:establish_connection]
+            end
+          end
+        end
+      end
+    end
+  end
+
+  Orm::ActiveRecord.singleton_class.send :prepend, OpenidConnect::Orm::ActiveRecord
+end
diff --git a/lib/doorkeeper/openid_connect/orm/active_record/access_grant.rb b/lib/doorkeeper/openid_connect/orm/active_record/access_grant.rb
new file mode 100644
index 0000000..1476abf
--- /dev/null
+++ b/lib/doorkeeper/openid_connect/orm/active_record/access_grant.rb
@@ -0,0 +1,16 @@
+module Doorkeeper
+  module OpenidConnect
+    module AccessGrant
+      def self.prepended(base)
+        base.class_eval do
+          has_one :openid_connect_nonce,
+            class_name: 'Doorkeeper::OpenidConnect::Nonce',
+            inverse_of: :access_grant,
+            dependent: :delete
+        end
+      end
+    end
+  end
+
+  AccessGrant.send :prepend, OpenidConnect::AccessGrant
+end
diff --git a/lib/doorkeeper/openid_connect/orm/active_record/nonce.rb b/lib/doorkeeper/openid_connect/orm/active_record/nonce.rb
new file mode 100644
index 0000000..ebc28a9
--- /dev/null
+++ b/lib/doorkeeper/openid_connect/orm/active_record/nonce.rb
@@ -0,0 +1,17 @@
+module Doorkeeper
+  module OpenidConnect
+    class Nonce < ActiveRecord::Base
+      self.table_name = "#{table_name_prefix}oauth_openid_connect_nonces#{table_name_suffix}".to_sym
+
+      validates :access_grant_id, :nonce, presence: true
+      belongs_to :access_grant,
+        class_name: 'Doorkeeper::AccessGrant',
+        inverse_of: :openid_connect_nonce
+
+      def use!
+        destroy!
+        nonce
+      end
+    end
+  end
+end
diff --git a/lib/generators/doorkeeper/openid_connect/migration_generator.rb b/lib/generators/doorkeeper/openid_connect/migration_generator.rb
new file mode 100644
index 0000000..0f5e481
--- /dev/null
+++ b/lib/generators/doorkeeper/openid_connect/migration_generator.rb
@@ -0,0 +1,15 @@
+require 'rails/generators/active_record'
+
+class Doorkeeper::OpenidConnect::MigrationGenerator < ::Rails::Generators::Base
+  include Rails::Generators::Migration
+  source_root File.expand_path('../templates', __FILE__)
+  desc 'Installs Doorkeeper OpenID Connect migration file.'
+
+  def install
+    migration_template 'migration.rb', 'db/migrate/create_doorkeeper_openid_connect_tables.rb'
+  end
+
+  def self.next_migration_number(dirname)
+    ActiveRecord::Generators::Base.next_migration_number(dirname)
+  end
+end
diff --git a/lib/generators/doorkeeper/openid_connect/templates/migration.rb b/lib/generators/doorkeeper/openid_connect/templates/migration.rb
new file mode 100644
index 0000000..7dc6424
--- /dev/null
+++ b/lib/generators/doorkeeper/openid_connect/templates/migration.rb
@@ -0,0 +1,14 @@
+class CreateDoorkeeperOpenidConnectTables < ActiveRecord::Migration
+  def change
+    create_table :oauth_openid_connect_nonces do |t|
+      t.integer :access_grant_id, null: false
+      t.string :nonce, null: false
+    end
+
+    add_foreign_key(
+      :oauth_openid_connect_nonces,
+      :oauth_access_grants,
+      column: :access_grant_id
+    )
+  end
+end
diff --git a/spec/controllers/doorkeeper/openid_connect/discovery_controller_spec.rb b/spec/controllers/discovery_controller_spec.rb
similarity index 100%
rename from spec/controllers/doorkeeper/openid_connect/discovery_controller_spec.rb
rename to spec/controllers/discovery_controller_spec.rb
diff --git a/spec/controllers/doorkeeper/openid_connect/userinfo_controller_spec.rb b/spec/controllers/userinfo_controller_spec.rb
similarity index 100%
rename from spec/controllers/doorkeeper/openid_connect/userinfo_controller_spec.rb
rename to spec/controllers/userinfo_controller_spec.rb
diff --git a/spec/dummy/db/migrate/20161031135842_create_doorkeeper_openid_connect_tables.rb b/spec/dummy/db/migrate/20161031135842_create_doorkeeper_openid_connect_tables.rb
new file mode 100644
index 0000000..7dc6424
--- /dev/null
+++ b/spec/dummy/db/migrate/20161031135842_create_doorkeeper_openid_connect_tables.rb
@@ -0,0 +1,14 @@
+class CreateDoorkeeperOpenidConnectTables < ActiveRecord::Migration
+  def change
+    create_table :oauth_openid_connect_nonces do |t|
+      t.integer :access_grant_id, null: false
+      t.string :nonce, null: false
+    end
+
+    add_foreign_key(
+      :oauth_openid_connect_nonces,
+      :oauth_access_grants,
+      column: :access_grant_id
+    )
+  end
+end
diff --git a/spec/dummy/db/schema.rb b/spec/dummy/db/schema.rb
index 1ba2e84..21d6dff 100644
--- a/spec/dummy/db/schema.rb
+++ b/spec/dummy/db/schema.rb
@@ -1,3 +1,4 @@
+# encoding: UTF-8
 # This file is auto-generated from the current state of the database. Instead
 # of editing this file, please use the migrations feature of Active Record to
 # incrementally modify your database, and then regenerate this schema definition.
@@ -10,7 +11,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 20160320211015) do
+ActiveRecord::Schema.define(version: 20161031135842) do
 
   create_table "oauth_access_grants", force: :cascade do |t|
     t.integer  "resource_owner_id", null: false
@@ -21,9 +22,10 @@
     t.datetime "created_at",        null: false
     t.datetime "revoked_at"
     t.string   "scopes"
-    t.index ["token"], name: "index_oauth_access_grants_on_token", unique: true
   end
 
+  add_index "oauth_access_grants", ["token"], name: "index_oauth_access_grants_on_token", unique: true
+
   create_table "oauth_access_tokens", force: :cascade do |t|
     t.integer  "resource_owner_id"
     t.integer  "application_id"
@@ -34,11 +36,12 @@
     t.datetime "created_at",                          null: false
     t.string   "scopes"
     t.string   "previous_refresh_token", default: "", null: false
-    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
   end
 
+  add_index "oauth_access_tokens", ["refresh_token"], name: "index_oauth_access_tokens_on_refresh_token", unique: true
+  add_index "oauth_access_tokens", ["resource_owner_id"], name: "index_oauth_access_tokens_on_resource_owner_id"
+  add_index "oauth_access_tokens", ["token"], name: "index_oauth_access_tokens_on_token", unique: true
+
   create_table "oauth_applications", force: :cascade do |t|
     t.string   "name",                      null: false
     t.string   "uid",                       null: false
@@ -49,8 +52,14 @@
     t.datetime "updated_at",                null: false
     t.integer  "owner_id"
     t.string   "owner_type"
-    t.index ["owner_id", "owner_type"], name: "index_oauth_applications_on_owner_id_and_owner_type"
-    t.index ["uid"], name: "index_oauth_applications_on_uid", unique: true
+  end
+
+  add_index "oauth_applications", ["owner_id", "owner_type"], name: "index_oauth_applications_on_owner_id_and_owner_type"
+  add_index "oauth_applications", ["uid"], name: "index_oauth_applications_on_uid", unique: true
+
+  create_table "oauth_openid_connect_nonces", force: :cascade do |t|
+    t.integer "access_grant_id", null: false
+    t.string  "nonce",           null: false
   end
 
   create_table "users", force: :cascade do |t|
diff --git a/spec/factories.rb b/spec/factories.rb
index 08bfe6e..5d8f7fa 100644
--- a/spec/factories.rb
+++ b/spec/factories.rb
@@ -1,14 +1,14 @@
 FactoryGirl.define do
-  # factory :access_grant, class: Doorkeeper::AccessGrant do
-  #   sequence(:resource_owner_id) { |n| n }
-  #   application
-  #   redirect_uri 'https://app.com/callback'
-  #   expires_in 100
-  #   scopes 'public write'
-  # end
+  factory :access_grant, class: Doorkeeper::AccessGrant do
+    resource_owner_id { create(:user).id }
+    application
+    redirect_uri 'https://app.com/callback'
+    expires_in 100
+    scopes 'public write'
+  end
 
   factory :access_token, class: Doorkeeper::AccessToken do
-    sequence(:resource_owner_id) { |n| n }
+    resource_owner_id { create(:user).id }
     application
     expires_in 2.hours
 
@@ -23,4 +23,9 @@
   end
 
   factory :user
+
+  factory :nonce, class: Doorkeeper::OpenidConnect::Nonce do
+    access_grant
+    sequence(:nonce) { |n| n.to_s }
+  end
 end
diff --git a/spec/lib/doorkeeper/openid_connect/config_spec.rb b/spec/lib/config_spec.rb
similarity index 78%
rename from spec/lib/doorkeeper/openid_connect/config_spec.rb
rename to spec/lib/config_spec.rb
index 75aba8a..c335a70 100644
--- a/spec/lib/doorkeeper/openid_connect/config_spec.rb
+++ b/spec/lib/config_spec.rb
@@ -3,16 +3,31 @@
 describe Doorkeeper::OpenidConnect, 'configuration' do
   subject { Doorkeeper::OpenidConnect.configuration }
 
-  after :all do
+  after :each do
+    load "#{Rails.root}/config/initializers/doorkeeper.rb"
     load "#{Rails.root}/config/initializers/doorkeeper_openid_connect.rb"
   end
 
   describe 'scopes' do
-    it' adds the openid scope to the Doorkeeper configuration' do
+    it 'adds the openid scope to the Doorkeeper configuration' do
       expect(Doorkeeper.configuration.scopes).to include 'openid'
     end
   end
 
+  describe 'orm' do
+    it 'fails if not set to :active_record' do
+      # stub ORM setup to avoid Doorkeeper exceptions
+      allow(Doorkeeper).to receive(:setup_orm_adapter)
+      allow(Doorkeeper).to receive(:setup_orm_models)
+
+      expect do
+        Doorkeeper.configure do
+          orm :mongoid
+        end
+      end.to raise_error Doorkeeper::OpenidConnect::ConfigurationError
+    end
+  end
+
   describe 'jws_private_key' do
     it 'sets the value that is accessible via jws_private_key' do
       value = 'private_key'
@@ -79,7 +94,7 @@
         claims do
         end
       end
-      expect(subject.claims).not_to be_nil
+      expect(subject.claims).to_not be_nil
     end
   end
 end
diff --git a/spec/lib/oauth/authorization/code_spec.rb b/spec/lib/oauth/authorization/code_spec.rb
new file mode 100644
index 0000000..50ea947
--- /dev/null
+++ b/spec/lib/oauth/authorization/code_spec.rb
@@ -0,0 +1,35 @@
+require 'rails_helper'
+
+describe Doorkeeper::OpenidConnect::OAuth::Authorization::Code do
+  subject { Doorkeeper::OAuth::Authorization::Code.new pre_auth, resource_owner }
+  let(:resource_owner) { create :user }
+  let(:access_grant) { create :access_grant }
+  let(:pre_auth) { double }
+  let(:client) { double }
+
+  describe '#issue_token' do
+    before do
+      allow(pre_auth).to receive(:client) { client }
+      allow(pre_auth).to receive(:redirect_uri) { 'redirect_uri' }
+      allow(pre_auth).to receive(:scopes) { 'scopes' }
+      allow(pre_auth).to receive(:nonce) { '123456' }
+      allow(client).to receive(:id) { 'client_id' }
+
+      allow(Doorkeeper::AccessGrant).to receive(:create!) { access_grant }
+      allow(Doorkeeper::OpenidConnect::Nonce).to receive(:create!)
+    end
+
+    it 'stores the nonce' do
+      subject.issue_token
+
+      expect(Doorkeeper::OpenidConnect::Nonce).to have_received(:create!).with({
+        access_grant: access_grant,
+        nonce: '123456'
+      })
+    end
+
+    it 'returns the created grant' do
+      expect(subject.issue_token).to be_a Doorkeeper::AccessGrant
+    end
+  end
+end
diff --git a/spec/lib/oauth/authorization_code_request_spec.rb b/spec/lib/oauth/authorization_code_request_spec.rb
new file mode 100644
index 0000000..f86cb00
--- /dev/null
+++ b/spec/lib/oauth/authorization_code_request_spec.rb
@@ -0,0 +1,22 @@
+require 'rails_helper'
+
+describe Doorkeeper::OpenidConnect::OAuth::AuthorizationCodeRequest do
+  subject { Doorkeeper::OAuth::AuthorizationCodeRequest.new server, grant, client }
+  let(:server) { double }
+  let(:client) { double }
+  let(:grant) { create :access_grant, openid_connect_nonce: nonce }
+  let(:nonce) { build :nonce, nonce: '123456' }
+  let(:token) { create :access_token }
+  let(:response) { Doorkeeper::OAuth::TokenResponse.new token }
+
+  describe '#after_successful_response' do
+    it 'adds the ID token to the response' do
+      subject.instance_variable_set '@response', response
+      subject.access_token = token
+      subject.send :after_successful_response
+
+      expect(response.id_token).to be_a Doorkeeper::OpenidConnect::Models::IdToken
+      expect(response.id_token.nonce).to eq '123456'
+    end
+  end
+end
diff --git a/spec/lib/oauth/password_access_token_request_spec.rb b/spec/lib/oauth/password_access_token_request_spec.rb
new file mode 100644
index 0000000..d88f150
--- /dev/null
+++ b/spec/lib/oauth/password_access_token_request_spec.rb
@@ -0,0 +1,27 @@
+require 'rails_helper'
+
+describe Doorkeeper::OpenidConnect::OAuth::PasswordAccessTokenRequest do
+  subject { Doorkeeper::OAuth::PasswordAccessTokenRequest.new server, client, resource_owner, { nonce: '123456' } }
+  let(:server) { double }
+  let(:client) { double }
+  let(:resource_owner) { create :user }
+  let(:token) { create :access_token }
+  let(:response) { Doorkeeper::OAuth::TokenResponse.new token }
+
+  describe '#initialize' do
+    it 'stores the nonce attribute' do
+      expect(subject.nonce).to eq '123456'
+    end
+  end
+
+  describe '#after_successful_response' do
+    it 'adds the ID token to the response' do
+      subject.instance_variable_set '@response', response
+      subject.access_token = token
+      subject.send :after_successful_response
+
+      expect(response.id_token).to be_a Doorkeeper::OpenidConnect::Models::IdToken
+      expect(response.id_token.nonce).to eq '123456'
+    end
+  end
+end
diff --git a/spec/lib/oauth/pre_authorization_spec.rb b/spec/lib/oauth/pre_authorization_spec.rb
new file mode 100644
index 0000000..fb47906
--- /dev/null
+++ b/spec/lib/oauth/pre_authorization_spec.rb
@@ -0,0 +1,13 @@
+require 'rails_helper'
+
+describe Doorkeeper::OpenidConnect::OAuth::PreAuthorization do
+  subject { Doorkeeper::OAuth::PreAuthorization.new server, client, { nonce: '123456' } }
+  let(:server) { double }
+  let(:client) { double }
+
+  describe '#initialize' do
+    it 'stores the nonce attribute' do
+      expect(subject.nonce).to eq '123456'
+    end
+  end
+end
diff --git a/spec/lib/oauth/token_response_spec.rb b/spec/lib/oauth/token_response_spec.rb
new file mode 100644
index 0000000..cb884a6
--- /dev/null
+++ b/spec/lib/oauth/token_response_spec.rb
@@ -0,0 +1,33 @@
+require 'rails_helper'
+
+describe Doorkeeper::OpenidConnect::OAuth::TokenResponse do
+  subject { Doorkeeper::OAuth::TokenResponse.new token }
+  let(:token) { create :access_token }
+  let(:id_token) { Doorkeeper::OpenidConnect::Models::IdToken.new token, '123456' }
+
+  describe '#body' do
+    before do
+      subject.id_token = id_token
+    end
+
+    context 'with the openid scope present' do
+      before do
+        token.scopes = 'openid email'
+      end
+
+      it 'adds the ID token to the response' do
+        expect(subject.body[:id_token]).to eq id_token.as_jws_token
+      end
+    end
+
+    context 'with the openid scope not present' do
+      before do
+        token.scopes = 'email'
+      end
+
+      it 'does not add the ID token to the response' do
+        expect(subject.body).to_not include :id_token
+      end
+    end
+  end
+end
diff --git a/spec/lib/doorkeeper/openid_connect_spec.rb b/spec/lib/openid_connect_spec.rb
similarity index 100%
rename from spec/lib/doorkeeper/openid_connect_spec.rb
rename to spec/lib/openid_connect_spec.rb
diff --git a/spec/lib/doorkeeper/openid_connect/routes_spec.rb b/spec/lib/routes_spec.rb
similarity index 100%
rename from spec/lib/doorkeeper/openid_connect/routes_spec.rb
rename to spec/lib/routes_spec.rb
diff --git a/spec/models/access_grant_spec.rb b/spec/models/access_grant_spec.rb
new file mode 100644
index 0000000..f804fee
--- /dev/null
+++ b/spec/models/access_grant_spec.rb
@@ -0,0 +1,15 @@
+require 'rails_helper'
+
+describe Doorkeeper::OpenidConnect::AccessGrant do
+  subject { Doorkeeper::AccessGrant.new }
+
+  it 'has one openid_connect_nonce' do
+    association = subject.class.reflect_on_association :openid_connect_nonce
+
+    expect(association.options).to eq({
+      class_name: 'Doorkeeper::OpenidConnect::Nonce',
+      inverse_of: :access_grant,
+      dependent: :delete,
+    })
+  end
+end
diff --git a/spec/models/id_token_spec.rb b/spec/models/id_token_spec.rb
index 7fdf0f1..d024898 100644
--- a/spec/models/id_token_spec.rb
+++ b/spec/models/id_token_spec.rb
@@ -1,9 +1,16 @@
 require 'rails_helper'
 
 describe Doorkeeper::OpenidConnect::Models::IdToken, type: :model do
-  subject { described_class.new(access_token) }
+  subject { described_class.new(access_token, nonce) }
   let(:access_token) { create :access_token, resource_owner_id: user.id }
   let(:user) { create :user }
+  let(:nonce) { '123456' }
+
+  describe '#nonce' do
+    it 'returns the stored nonce' do
+      expect(subject.nonce).to eq '123456'
+    end
+  end
 
   describe '#claims' do
     it 'returns all default claims' do
@@ -12,6 +19,7 @@
       expect(subject.claims[:aud]).to eq access_token.application.uid
       expect(subject.claims[:exp]).to eq subject.claims[:iat] + 120
       expect(subject.claims[:iat]).to be_a Integer
+      expect(subject.claims[:nonce]).to eq nonce
     end
   end
 
@@ -21,7 +29,7 @@
       json = subject.as_json
 
       expect(json).to include :aud
-      expect(json).not_to include :iss
+      expect(json).to_not include :iss
     end
   end
 
diff --git a/spec/models/nonce_spec.rb b/spec/models/nonce_spec.rb
new file mode 100644
index 0000000..a751860
--- /dev/null
+++ b/spec/models/nonce_spec.rb
@@ -0,0 +1,42 @@
+require 'rails_helper'
+
+describe Doorkeeper::OpenidConnect::Nonce do
+  describe 'validations' do
+    it 'requires an access grant' do
+      subject.access_grant_id = nil
+
+      expect(subject).to_not be_valid
+      expect(subject.errors).to include :access_grant_id
+    end
+
+    it 'requires a nonce' do
+      subject.nonce = nil
+
+      expect(subject).to_not be_valid
+      expect(subject.errors).to include :nonce
+    end
+  end
+
+  describe 'associations' do
+    it 'belongs to an access_grant' do
+      association = subject.class.reflect_on_association :access_grant
+
+      expect(association.options).to eq({
+        class_name: 'Doorkeeper::AccessGrant',
+        inverse_of: :openid_connect_nonce,
+      })
+    end
+  end
+
+  describe '#use' do
+    before do
+      subject.nonce = '123456'
+      allow(subject).to receive(:destroy!)
+    end
+
+    it 'destroys the record and returns the nonce' do
+      expect(subject.use!).to eq '123456'
+      expect(subject).to have_received(:destroy!)
+    end
+  end
+end