feat: async generateStateFunction and checkStateFunction (#239)

* feat: impl async function for generateStateFunction and checkStateFunction

* test: add test cases

* docs: update README.md

* docs: update README.md

Author: NotEvenANeko

README.md

@@ -222,6 +222,38 @@ When you set it, it is required to provide the function `checkStateFunction` in
   })
 ```
 
+Async functions are supported here, and the fastify instance can be accessed via `this`.
+
+```js
+  fastify.register(oauthPlugin, {
+    name: 'facebookOAuth2',
+    credentials: {
+      client: {
+        id: '<CLIENT_ID>',
+        secret: '<CLIENT_SECRET>'
+      },
+      auth: oauthPlugin.FACEBOOK_CONFIGURATION
+    },
+    // register a fastify url to start the redirect flow
+    startRedirectPath: '/login/facebook',
+    // facebook redirect here after the user login
+    callbackUri: 'http://localhost:3000/login/facebook/callback',
+    // custom function to generate the state and store it into the redis
+    generateStateFunction: async function (request) {
+      const state = request.query.customCode
+      await this.redis.set(stateKey, state)
+      return state
+    },
+    // custom function to check the state is valid
+    checkStateFunction: async function (request, callback) {
+      if (request.query.state !== request.session.state) {
+        throw new Error('Invalid state')
+      }
+      return true
+    }
+  })
+```
+
 ## Set custom callbackUri Parameters
 
 The `callbackUriParams` accepts an object that will be translated to query parameters for the callback OAUTH flow. The default value is {}.
@@ -309,12 +341,13 @@ fastify.googleOAuth2.getNewAccessTokenUsingRefreshToken(currentAccessToken, (err
 });
 ```
 
-- `generateAuthorizationUri(requestObject, replyObject)`: A function that returns the authorization uri. This is generally useful when you want to handle the redirect yourself in a specific route. The `requestObject` argument passes the request object to the `generateStateFunction`). You **do not** need to declare a `startRedirectPath` if you use this approach. Example of how you would use it:
+- `generateAuthorizationUri(requestObject, replyObject, callback)`: A function that generates the authorization uri. If the callback is not passed this function will return a Promise. The string resulting from the callback call or the resolved Promise is the authorization uri. This is generally useful when you want to handle the redirect yourself in a specific route. The `requestObject` argument passes the request object to the `generateStateFunction`). You **do not** need to declare a `startRedirectPath` if you use this approach. Example of how you would use it:
 
 ```js
-fastify.get('/external', { /* Hooks can be used here */ }, async (req, reply) => {
-  const authorizationEndpoint = fastify.oauth2CustomOAuth2.generateAuthorizationUri(req, reply);
-  reply.redirect(authorizationEndpoint)
+fastify.get('/external', { /* Hooks can be used here */ }, (req, reply) => {
+  fastify.oauth2CustomOAuth2.generateAuthorizationUri(req, reply, (err, authorizationEndpoint) => {
+    reply.redirect(authorizationEndpoint)
+  });
 });
 ```
 

index.js

@@ -19,8 +19,8 @@ const random = (bytes = 32) => randomBytes(bytes).toString('base64url')
 const codeVerifier = random
 const codeChallenge = verifier => createHash('sha256').update(verifier).digest('base64url')
 
-function defaultGenerateStateFunction () {
-  return random(16)
+function defaultGenerateStateFunction (request, callback) {
+  callback(null, random(16))
 }
 
 function defaultCheckStateFunction (request, callback) {
@@ -131,36 +131,68 @@ function fastifyOauth2 (fastify, options, next) {
     const generateCallbackUriParams = (credentials.auth && credentials.auth[kGenerateCallbackUriParams]) || defaultGenerateCallbackUriParams
     const cookieOpts = Object.assign({ httpOnly: true, sameSite: 'lax' }, options.cookie)
 
-    function generateAuthorizationUri (request, reply) {
-      const state = generateStateFunction(request)
+    const generateStateFunctionCallbacked = function (request, callback) {
+      const boundGenerateStateFunction = generateStateFunction.bind(fastify)
 
-      reply.setCookie('oauth2-redirect-state', state, cookieOpts)
+      if (generateStateFunction.length <= 1) {
+        callbackify(function (request) {
+          return Promise.resolve(boundGenerateStateFunction(request))
+        })(request, callback)
+      } else {
+        boundGenerateStateFunction(request, callback)
+      }
+    }
 
-      // when PKCE extension is used
-      let pkceParams = {}
-      if (configured.pkce) {
-        const verifier = codeVerifier()
-        const challenge = configured.pkce === 'S256' ? codeChallenge(verifier) : verifier
-        pkceParams = {
-          code_challenge: challenge,
-          code_challenge_method: configured.pkce
+    function generateAuthorizationUriCallbacked (request, reply, callback) {
+      generateStateFunctionCallbacked(request, function (err, state) {
+        if (err) {
+          callback(err, null)
+          return
         }
-        reply.setCookie(VERIFIER_COOKIE_NAME, verifier, cookieOpts)
-      }
 
-      const urlOptions = Object.assign({}, generateCallbackUriParams(callbackUriParams, request, scope, state), {
-        redirect_uri: callbackUri,
-        scope,
-        state
-      }, pkceParams)
+        reply.setCookie('oauth2-redirect-state', state, cookieOpts)
+
+        // when PKCE extension is used
+        let pkceParams = {}
+        if (configured.pkce) {
+          const verifier = codeVerifier()
+          const challenge = configured.pkce === 'S256' ? codeChallenge(verifier) : verifier
+          pkceParams = {
+            code_challenge: challenge,
+            code_challenge_method: configured.pkce
+          }
+          reply.setCookie(VERIFIER_COOKIE_NAME, verifier, cookieOpts)
+        }
+
+        const urlOptions = Object.assign({}, generateCallbackUriParams(callbackUriParams, request, scope, state), {
+          redirect_uri: callbackUri,
+          scope,
+          state
+        }, pkceParams)
+
+        callback(null, oauth2.authorizeURL(urlOptions))
+      })
+    }
 
-      return oauth2.authorizeURL(urlOptions)
+    const generateAuthorizationUriPromisified = promisify(generateAuthorizationUriCallbacked)
+
+    function generateAuthorizationUri (request, reply, callback) {
+      if (!callback) {
+        return generateAuthorizationUriPromisified(request, reply)
+      }
+
+      generateAuthorizationUriCallbacked(request, reply, callback)
     }
 
     function startRedirectHandler (request, reply) {
-      const authorizationUri = generateAuthorizationUri(request, reply)
+      generateAuthorizationUriCallbacked(request, reply, function (err, authorizationUri) {
+        if (err) {
+          reply.code(500).send(err.message)
+          return
+        }
 
-      reply.redirect(authorizationUri)
+        reply.redirect(authorizationUri)
+      })
     }
 
     const cbk = function (o, code, pkceParams, callback) {
@@ -172,6 +204,24 @@ function fastifyOauth2 (fastify, options, next) {
       return callbackify(o.oauth2.getToken.bind(o.oauth2, body))(callback)
     }
 
+    function checkStateFunctionCallbacked (request, callback) {
+      const boundCheckStateFunction = checkStateFunction.bind(fastify)
+
+      if (checkStateFunction.length <= 1) {
+        Promise.resolve(boundCheckStateFunction(request))
+          .then(function (result) {
+            if (result) {
+              callback()
+            } else {
+              callback(new Error('Invalid state'))
+            }
+          })
+          .catch(function (err) { callback(err) })
+      } else {
+        boundCheckStateFunction(request, callback)
+      }
+    }
+
     function getAccessTokenFromAuthorizationCodeFlowCallbacked (request, reply, callback) {
       const code = request.query.code
       const pkceParams = configured.pkce ? { code_verifier: request.cookies['oauth2-code-verifier'] } : {}
@@ -183,7 +233,7 @@ function fastifyOauth2 (fastify, options, next) {
         clearCodeVerifierCookie(reply)
       }
 
-      checkStateFunction(request, function (err) {
+      checkStateFunctionCallbacked(request, function (err) {
         if (err) {
           callback(err)
           return