feat: add optional onAuthentication option to add user property in socket object (#62)

README.md

@@ -87,10 +87,38 @@ io.on('connection', async (socket) => {
 })
 ```
 
+### Server side with `onAuthentication` (example)
+
+```ts
+import { Server } from 'socket.io'
+import { authorize } from '@thream/socketio-jwt'
+
+const io = new Server(9000)
+io.use(
+  authorize({
+    secret: 'your secret or public key',
+    algorithms: ['RS256'],
+    onAuthentication: async decodedToken => {
+        // return the object that you want to add to the user property
+        // or throw an error if the token is unauthorized
+    }
+  })
+)
+
+io.on('connection', async (socket) => {
+  // jwt payload of the connected client
+  console.log(socket.decodedToken)
+  // You can do the same things of the previous example there...
+  // user object returned in onAuthentication
+  console.log(socket.user)
+})
+```
+
 ### `authorize` options
 
 - `secret` is a string containing the secret for HMAC algorithms, or a function that should fetch the secret or public key as shown in the example with `jwks-rsa`.
 - `algorithms` (default: `HS256`)
+- `onAuthentication` is a function that will be called with the decodedToken as a parameter after the token is authenticated. Return a value to add to the `user` property in the socket object.
 
 ### Client side
 

src/__test__/authorize.test.ts

@@ -1,7 +1,7 @@
 import axios from 'axios'
 import { io } from 'socket.io-client'
 
-import { fixtureStart, fixtureStop } from './fixture'
+import { fixtureStart, fixtureStop, getSocket, Profile } from './fixture'
 
 describe('authorize - with secret as string in options', () => {
   let token: string = ''
@@ -101,3 +101,76 @@ describe('authorize - with secret as callback in options', () => {
     })
   })
 })
+
+describe('authorize - with onAuthentication callback in options', () => {
+  let token: string = ''
+  let wrongToken: string = ''
+
+  beforeEach(async (done) => {
+    jest.setTimeout(15_000)
+    await fixtureStart(
+      async () => {
+        const response = await axios.post('http://localhost:9000/login')
+        token = response.data.token
+        const responseWrong = await axios.post('http://localhost:9000/login-wrong')
+        wrongToken = responseWrong.data.token
+        done()
+      },
+      {
+        secret: secretCallback,
+        onAuthentication: (decodedToken: Profile) => {
+          if (!decodedToken.checkField) {
+            throw new Error('Check Field validation failed')
+          }
+          return {
+            email: decodedToken.email
+          }
+        }
+      }
+    )
+  })
+
+  afterEach((done) => {
+    fixtureStop(done)
+  })
+
+  it('should connect the user', (done) => {
+    const socket = io('http://localhost:9000', {
+      auth: { token: `Bearer ${token}` }
+    })
+    socket.on('connect', () => {
+      socket.close()
+      done()
+    })
+  })
+
+  it('should contain user property', (done) => {
+    const socketServer = getSocket()
+    socketServer?.on('connection', (client: any) => {
+      expect(client.user.email).toEqual('john@doe.com')
+    })
+    const socket = io('http://localhost:9000', {
+      auth: { token: `Bearer ${token}` }
+    })
+    socket.on('connect', () => {
+      socket.close()
+      done()
+    })
+  })
+
+  it('should emit error when user validation fails', (done) => {
+    const socket = io('http://localhost:9000', {
+      auth: { token: `Bearer ${wrongToken}` }
+    })
+    socket.on('connect_error', (err: any) => {
+      try {
+        expect(err.message).toEqual('Check Field validation failed')
+      } catch (err) {
+        socket.close()
+        done(err)
+      }
+      socket.close()
+      done()
+    })
+  })
+})

src/__test__/fixture/index.ts

@@ -7,6 +7,12 @@ import enableDestroy from 'server-destroy'
 
 import { authorize, AuthorizeOptions } from '../../index'
 
+export interface Profile {
+  email: string
+  id: number
+  checkField: boolean
+}
+
 interface Socket {
   io: null | SocketIoServer
   init: (httpServer: HttpServer | HttpsServer) => void
@@ -34,9 +40,21 @@ export const fixtureStart = async (
     keySecret = await options.secret({ header: { alg: 'RS256' }, payload: '' })
   }
   app.post('/login', (_req, res) => {
-    const profile = {
+    const profile: Profile = {
       email: 'john@doe.com',
-      id: 123
+      id: 123,
+      checkField: true
+    }
+    const token = jwt.sign(profile, keySecret, {
+      expiresIn: 60 * 60 * 5
+    })
+    return res.json({ token })
+  })
+  app.post('/login-wrong', (_req, res) => {
+    const profile: Profile = {
+      email: 'john@doe.com',
+      id: 123,
+      checkField: false
     }
     const token = jwt.sign(profile, keySecret, {
       expiresIn: 60 * 60 * 5
@@ -56,3 +74,7 @@ export const fixtureStop = (callback: Function): void => {
   } catch (err) {}
   callback()
 }
+
+export const getSocket = (): SocketIoServer | null => {
+  return socket.io
+}

src/authorize.ts

@@ -14,6 +14,7 @@ interface ExtendedError extends Error {
 interface ExtendedSocket {
   encodedToken?: string
   decodedToken?: any
+  user?: any
 }
 
 type SocketIOMiddleware = (
@@ -34,10 +35,11 @@ type SecretCallback = (decodedToken: CompleteDecodedToken) => Promise<string>
 export interface AuthorizeOptions {
   secret: string | SecretCallback
   algorithms?: Algorithm[]
+  onAuthentication?: (decodedToken: any) => Promise<any> | any
 }
 
 export const authorize = (options: AuthorizeOptions): SocketIOMiddleware => {
-  const { secret, algorithms = ['HS256'] } = options
+  const { secret, algorithms = ['HS256'], onAuthentication } = options
   return async (socket, next) => {
     let encodedToken: string | null = null
     const { token } = socket.handshake.auth
@@ -78,6 +80,13 @@ export const authorize = (options: AuthorizeOptions): SocketIOMiddleware => {
       )
     }
     socket.decodedToken = decodedToken
+    if (onAuthentication != null) {
+      try {
+        socket.user = await onAuthentication(decodedToken)
+      } catch (err) {
+        return next(err)
+      }
+    }
     return next()
   }
 }