diff --git a/index.js b/index.js
index e9e2636..d04c3a0 100644
--- a/index.js
+++ b/index.js
@@ -15,11 +15,18 @@ const fromHapiReq = req => {
 exports.plugin = {
   pkg: require('./package.json'),
   register: async function (server, options) {
-    await Joi.validate(options, Joi.object().keys({
-      model: Joi.object().required()
-    }))
 
-    const oauthServer = new OAuth2Server(options)
+    Joi.assert(options, Joi.object().keys({
+      model: Joi.alternatives().try(Joi.object().required(), Joi.function())
+    }));
+
+    const model = typeof options.model === 'function' ? options.model(server) : options.model;
+
+    const oauthServer = new OAuth2Server({
+      ...options,
+      model: model
+    })
+
 
     const oauth = {
       authenticate: async (req, options) => {
@@ -39,6 +46,7 @@ exports.plugin = {
       }
     }
 
-    server.expose('oauth', oauth)
+    server.decorate('server', 'oauth', oauth);
+    server.decorate('request', 'oauth', oauth);
   }
 }
diff --git a/package.json b/package.json
index 2b4b7c4..54f58b9 100644
--- a/package.json
+++ b/package.json
@@ -24,16 +24,16 @@
   },
   "homepage": "https://github.com/getapper/hapi-oauth-2-server#readme",
   "devDependencies": {
-    "code": "5.2.4",
-    "hapi": "17.8.4",
-    "lab": "18.0.2",
-    "moment": "2.24.0"
+    "@hapi/code": "8.x.x",
+    "@hapi/hapi": "20.x.x",
+    "@hapi/lab": "24.x.x",
+    "date-fns": "2.21.3"
   },
   "peerDependencies": {
-    "hapi": "^17.0.0"
+    "@hapi/hapi": "^18.x.x"
   },
   "dependencies": {
-    "joi": "14.3.1",
-    "oauth2-server": "3.0.1"
+    "joi": "17.x.x",
+    "oauth2-server": "3.1.1"
   }
 }
diff --git a/test/index.js b/test/index.js
index b969e7e..b97e02b 100644
--- a/test/index.js
+++ b/test/index.js
@@ -1,7 +1,8 @@
 const build = require('./server')
-const model = require('./model')
-const code = require('code')
-const labLib = require('lab')
+const model = require('./model');
+const modelWithServer = require('./modelWithServer');
+const code = require('@hapi/code')
+const labLib = require('@hapi/lab')
 
 const lab = (exports.lab = labLib.script())
 const describe = lab.describe
@@ -9,7 +10,7 @@ const before = lab.before
 const it = lab.it
 const expect = code.expect
 
-describe('Hapi oauth2 server', () => {
+describe('Hapi oauth2 server with object model', () => {
   let server
   let res
 
@@ -43,6 +44,54 @@ describe('Hapi oauth2 server', () => {
     expect(res.statusCode, 'Wrong HTTP response status code').to.equal(200)
   })
 
+  it('should get token', async () => {
+    res = await server.inject({
+      method: 'POST',
+      url: '/token',
+      headers: {
+        Authorization: 'Bearer foobar',
+        'Content-Type': 'application/x-www-form-urlencoded'
+      },
+      payload: 'client_id=test&client_secret=test&grant_type=authorization_code&code=test'
+    })
+    expect(res.statusCode, 'Wrong HTTP response status code').to.equal(200)
+  })
+});
+
+describe('Hapi oauth2 server with model from server callback', () => {
+  let server
+  let res
+
+  before(async () => {
+    server = await build(modelWithServer);
+  })
+
+  it('should fail for no authentication provided', async () => {
+    res = await server.inject({
+      method: 'GET',
+      url: '/authenticate'
+    })
+    expect(res.statusCode, 'Wrong HTTP response status code').to.equal(401)
+  })
+
+  it('should authenticate', async () => {
+    res = await server.inject({
+      method: 'GET',
+      url: '/authenticate',
+      headers: {Authorization: 'Bearer foobar'}
+    })
+    expect(res.statusCode, 'Wrong HTTP response status code').to.equal(200)
+  })
+
+  it('should authorize', async () => {
+    res = await server.inject({
+      method: 'GET',
+      url: '/authorize?client_id=test&client_secret=test&state=test&response_type=code',
+      headers: {Authorization: 'Bearer foobar'},
+    })
+    expect(res.statusCode, 'Wrong HTTP response status code').to.equal(200)
+  })
+
   it('should get token', async () => {
     res = await server.inject({
       method: 'POST',
diff --git a/test/model.js b/test/model.js
index 5a6ff46..e26e30c 100644
--- a/test/model.js
+++ b/test/model.js
@@ -1,13 +1,14 @@
-const moment = require('moment')
+const add = require('date-fns/add');
+const toDate = require('date-fns/toDate');
 
 module.exports = {
   getAccessToken: async () => {
-    return { user: {}, accessTokenExpiresAt: moment().add(3, 'days').toDate() }
+    return { user: {}, accessTokenExpiresAt: toDate(add(new Date(), {days: 3}))}
   },
   getAuthorizationCode: async code => {
     return {
       code: 'be6e365e9f29c19e631078e8e376326bdf086576',
-      expiresAt: moment().add(3, 'days').toDate(),
+      expiresAt: toDate(add(new Date(), {days: 3})),
       scope: 'test',
       client: { id: 'test' },
       user: {}
diff --git a/test/modelWithServer.js b/test/modelWithServer.js
new file mode 100644
index 0000000..9aa443d
--- /dev/null
+++ b/test/modelWithServer.js
@@ -0,0 +1,52 @@
+const add = require('date-fns/add');
+const toDate = require('date-fns/toDate');
+
+module.exports = (server) => {
+    return {
+        getAccessToken: async () => {
+            
+            return { user: {}, accessTokenExpiresAt: toDate(add(new Date(), { days: 3 })) }
+        },
+        getAuthorizationCode: async code => {
+            return {
+                code: 'be6e365e9f29c19e631078e8e376326bdf086576',
+                expiresAt: toDate(add(new Date(), { days: 3 })),
+                scope: 'test',
+                client: { id: 'test' },
+                user: {}
+            };
+        },
+        getClient: async (clientId, clientSecret) => {
+            return {
+                id: clientId,
+                clientSecret,
+                redirectUris: ['/ok'],
+                grants: [
+                    "password",
+                    "authorization_code",
+                    "refresh_token"
+                ]
+            }
+        },
+        saveAuthorizationCode: async (code, client, user) => {
+            return {
+                ...code,
+                client: { id: 'test' },
+                user: { id: user.userId }
+            }
+        },
+        saveToken: async (token, client, user) => {
+            return {
+                ...token,
+                client: { id: 'test' },
+                user: {}
+            }
+        },
+        getUser: async (username, password) => {
+            return { user: {} }
+        },
+        revokeAuthorizationCode: async code => {
+            return true
+        }
+    }
+}
diff --git a/test/server.js b/test/server.js
index f9ed031..8a5cbe6 100644
--- a/test/server.js
+++ b/test/server.js
@@ -1,4 +1,4 @@
-const Hapi = require('hapi')
+const Hapi = require('@hapi/hapi')
 const HapiOAuth2Server = require('../')
 
 module.exports = async model => {
@@ -19,7 +19,7 @@ module.exports = async model => {
     path: '/authenticate',
     config: {
       handler: async (req, h) => {
-        const { oauth } = req.server.plugins['hapi-oauth2-server-plugin']
+        const { oauth } = req.server
         try {
           return await oauth.authenticate(req)
         } catch (e) {
@@ -34,7 +34,7 @@ module.exports = async model => {
     path: '/authorize',
     config: {
       handler: async (req, h) => {
-        const { oauth } = req.server.plugins['hapi-oauth2-server-plugin']
+        const { oauth } = req.server
         try {
           return await oauth.authorize(req)
         } catch (e) {
@@ -49,7 +49,7 @@ module.exports = async model => {
     path: '/token',
     config: {
       handler: async (req, h) => {
-        const { oauth } = req.server.plugins['hapi-oauth2-server-plugin']
+        const { oauth } = req.server
         try {
           return await oauth.token(req)
         } catch (e) {