feat(socket): authenticate in handshake (#254)

* use wss

* feat(socket): authenticate in handshake

* timeout

* bump versions

* don't send back token

packages/board/src/linker.ts

@@ -128,7 +128,7 @@ export class BoardLinker {
         if (!e.choice) {
           e.choice = <any>{}
         }
-        e.choice!.userId = this.inputUserToken.value
+        e.choice!.userId = this.inputUserId.value
       })
     }
     if (this.inputUserToken.value.length) {

packages/chat/package.json

@@ -15,6 +15,6 @@
   ],
   "devDependencies": {},
   "dependencies": {
-    "@botpress/messaging-socket": "0.0.1"
+    "@botpress/messaging-socket": "0.0.2"
   }
 }

packages/chat/src/chat.ts

@@ -29,7 +29,6 @@ export class Webchat {
     await this.storage.setup()
     await this.locale.setup()
     await this.lang.setup()
-    await this.socket.connect({ autoLogin: false })
     await this.user.setup()
     await this.conversation.setup()
     await this.messages.setup()

packages/chat/src/user/system.ts

@@ -22,7 +22,7 @@ export class WebchatUser extends WebchatSystem {
     const event = { choice: saved }
     await this.emitter.emit(UserEvents.Choose, event)
 
-    const user = await this.socket.login(event.choice)
+    const user = await this.socket.connect(event.choice)
     this.storage.set(STORAGE_ID, user)
 
     await this.set(user)

packages/server/src/socket.ts

@@ -12,8 +12,8 @@ export class Socket {
   private messages: MessageSocket
 
   constructor({ clients, users, userTokens, conversations, messages, sockets }: App) {
-    this.manager = new SocketManager(sockets)
-    this.users = new UserSocket(this.manager, clients, users, userTokens, sockets)
+    this.manager = new SocketManager(clients, users, userTokens, sockets)
+    this.users = new UserSocket(this.manager, users)
     this.conversations = new ConversationSocket(this.manager, users, conversations)
     this.messages = new MessageSocket(this.manager, conversations, messages)
   }

packages/server/src/socket/manager.ts

@@ -4,19 +4,30 @@ import clc from 'cli-color'
 import { Server } from 'http'
 import Joi from 'joi'
 import Socket from 'socket.io'
+import { validate as validateUuid } from 'uuid'
 import yn from 'yn'
+import { ClientService } from '../clients/service'
+import { UserTokenService } from '../user-tokens/service'
+import { UserService } from '../users/service'
+import { Schema } from './schema'
 import { SocketService } from './service'
 
 export class SocketManager {
   private logger = new Logger('Socket')
   private ws: Socket.Server | undefined
   private handlers: { [type: string]: SocketHandler } = {}
 
-  constructor(private sockets: SocketService) {}
+  constructor(
+    private clients: ClientService,
+    private users: UserService,
+    private userTokens: UserTokenService,
+    private sockets: SocketService
+  ) {}
 
   async setup(server: Server) {
     if (yn(process.env.ENABLE_EXPERIMENTAL_SOCKETS)) {
       this.ws = new Socket.Server(server, { serveClient: false, cors: { origin: '*' } })
+      this.ws.use(this.handleSocketAuthentication.bind(this))
       this.ws.on('connection', this.handleSocketConnection.bind(this))
     }
   }
@@ -41,24 +52,16 @@ export class SocketManager {
     })
   }
 
-  public handle(
-    type: string,
-    schema: Joi.ObjectSchema<any>,
-    callback: (socket: SocketRequest) => Promise<void>,
-    checkUserId?: boolean
-  ) {
+  public handle(type: string, schema: Joi.ObjectSchema<any>, callback: (socket: SocketRequest) => Promise<void>) {
     this.handlers[type] = async (socket: Socket.Socket, message: SocketMessage) => {
-      // TODO: remove this
-      if (checkUserId !== false) {
-        const userId = this.sockets.getUserId(socket)
-        if (!userId) {
-          return this.reply(socket, message, {
-            error: true,
-            message: 'socket does not have user rights'
-          })
-        }
-        message.userId = userId
+      const userId = this.sockets.getUserId(socket)
+      if (!userId) {
+        return this.reply(socket, message, {
+          error: true,
+          message: 'socket does not have user rights'
+        })
       }
+      message.userId = userId
 
       const { error } = schema.validate(message.data)
       if (error) {
@@ -76,12 +79,64 @@ export class SocketManager {
     })
   }
 
+  private async handleSocketAuthentication(socket: Socket.Socket, next: (err?: Error) => void) {
+    try {
+      const { error } = Schema.Socket.Auth.validate(socket.handshake.auth)
+      if (error) {
+        return next(new Error(error.message))
+      }
+
+      const { clientId, creds } = socket.handshake.auth as {
+        clientId: uuid
+        creds?: { userId: uuid; userToken: string }
+      }
+
+      const client = await this.clients.getById(clientId)
+      if (!client) {
+        return next(new Error('Client not found'))
+      }
+
+      if (creds) {
+        const user = await this.users.get(creds.userId)
+        if (user?.clientId === clientId) {
+          const [userTokenId, userTokenToken] = creds.userToken.split('.')
+          if (
+            validateUuid(userTokenId) &&
+            userTokenToken?.length &&
+            (await this.userTokens.verifyToken(userTokenId, userTokenToken))
+          ) {
+            socket.data.creds = creds
+            // we don't need to send it back if it was already sent to us
+            delete socket.data.creds.userToken
+            return next()
+          }
+        }
+      }
+
+      const user = await this.users.create(clientId)
+      const tokenRaw = await this.userTokens.generateToken()
+      const userToken = await this.userTokens.create(user.id, tokenRaw, undefined)
+      socket.data.creds = { userId: user.id, userToken: `${userToken.id}.${tokenRaw}` }
+
+      next()
+    } catch (e) {
+      this.logger.error(e, 'An error occurred when authenticating a socket connection')
+
+      next(new Error('an error occurred authenticating socket'))
+    }
+  }
+
   private async handleSocketConnection(socket: Socket.Socket) {
     try {
       this.logger.debug(`${clc.blackBright(`[${socket.id}]`)} ${clc.bgBlue(clc.magentaBright('connection'))}`)
 
+      const { creds } = socket.data
+      delete socket.data
+
       await this.setupSocket(socket)
-      this.sockets.create(socket)
+      await this.sockets.create(socket, creds.userId)
+
+      socket.emit('login', creds)
     } catch (e) {
       this.logger.error(e, 'An error occurred during a socket connection')
     }

packages/server/src/socket/schema.ts

@@ -0,0 +1,13 @@
+import Joi from 'joi'
+
+const Socket = {
+  Auth: Joi.object({
+    clientId: Joi.string().guid().required(),
+    creds: Joi.object({
+      userId: Joi.string().guid().required(),
+      userToken: Joi.string().required()
+    }).optional()
+  }).required()
+}
+
+export const Schema = { Socket }

packages/server/src/socket/service.ts

@@ -26,8 +26,26 @@ export class SocketService extends Service {
     this.cacheByUserId = await this.cachingService.newServerCache('cache_sockets_by_user_id')
   }
 
-  public create(socket: Socket) {
-    this.sockets[socket.id] = {}
+  public async create(socket: Socket, userId: uuid) {
+    const current = this.socketsByUserId[userId]
+
+    if (!current || !current.find((x) => x.id === socket.id)) {
+      const state = {
+        socket,
+        userId
+      }
+      this.sockets[socket.id] = state
+      this.cache.set(socket.id, state)
+
+      const list = [...(current || []), socket]
+      this.socketsByUserId[userId] = list
+      this.cacheByUserId.set(userId, list)
+
+      // this is the first socket connection this user has on this server
+      if (list.length === 1) {
+        await this.emitter.emit(SocketEvents.UserConnected, { userId: state.userId })
+      }
+    }
   }
 
   public async delete(socket: Socket) {
@@ -67,28 +85,6 @@ export class SocketService extends Service {
     }
   }
 
-  public async registerForUser(socket: Socket, userId: uuid) {
-    const current = this.socketsByUserId[userId]
-
-    if (!current || !current.find((x) => x.id === socket.id)) {
-      const state = {
-        socket,
-        userId
-      }
-      this.sockets[socket.id] = state
-      this.cache.set(socket.id, state)
-
-      const list = [...(current || []), socket]
-      this.socketsByUserId[userId] = list
-      this.cacheByUserId.set(userId, list)
-
-      // this is the first socket connection this user has on this server
-      if (list.length === 1) {
-        await this.emitter.emit(SocketEvents.UserConnected, { userId: state.userId })
-      }
-    }
-  }
-
   public listByUser(userId: string) {
     const cached = this.cacheByUserId.get(userId)
     if (cached) {
@@ -102,5 +98,6 @@ export class SocketService extends Service {
 }
 
 export interface SocketState {
-  userId?: uuid
+  socket: Socket
+  userId: uuid
 }

packages/server/src/users/schema.ts

@@ -10,13 +10,7 @@ const Api = {
 }
 
 const Socket = {
-  Get: Joi.object({}).required(),
-
-  Auth: Joi.object({
-    clientId: Joi.string().guid().required(),
-    userId: Joi.string().guid().optional(),
-    userToken: Joi.string().optional()
-  })
+  Get: Joi.object({}).required()
 }
 
 export const Schema = { Api, Socket }

packages/server/src/users/socket.ts

@@ -1,64 +1,14 @@
-import { uuid } from '@botpress/messaging-base'
-import { ClientService } from '../clients/service'
 import { SocketManager, SocketRequest } from '../socket/manager'
-import { SocketService } from '../socket/service'
-import { UserTokenService } from '../user-tokens/service'
 import { Schema } from './schema'
 import { UserService } from './service'
 
 export class UserSocket {
-  constructor(
-    private sockets: SocketManager,
-    private clients: ClientService,
-    private users: UserService,
-    private userTokens: UserTokenService,
-    private socketService: SocketService
-  ) {}
+  constructor(private sockets: SocketManager, private users: UserService) {}
 
   setup() {
-    // TODO: this should be done when establishing the socket connection
-    this.sockets.handle('users.auth', Schema.Socket.Auth, this.auth.bind(this), false)
     this.sockets.handle('users.get', Schema.Socket.Get, this.get.bind(this))
   }
 
-  async auth(socket: SocketRequest) {
-    const {
-      clientId,
-      userId,
-      userToken: userTokenRaw
-    }: { clientId: uuid; userId: uuid; userToken: string } = socket.data
-
-    const client = await this.clients.getById(clientId)
-    if (!client) {
-      return socket.notFound('Client not found')
-    }
-
-    // TODO: refactor here
-
-    let success = true
-    let user = userId ? await this.users.get(userId) : undefined
-    let token = undefined
-
-    if (!user || user.clientId !== client.id) {
-      success = false
-    } else {
-      const [userTokenId, userTokenToken] = userTokenRaw.split('.')
-      if (!(await this.userTokens.verifyToken(userTokenId, userTokenToken))) {
-        success = false
-      }
-    }
-
-    if (!success) {
-      user = await this.users.create(clientId)
-      const tokenRaw = await this.userTokens.generateToken()
-      const userToken = await this.userTokens.create(user.id, tokenRaw, undefined)
-      token = `${userToken.id}.${tokenRaw}`
-    }
-
-    await this.socketService.registerForUser(socket.socket, user!.id)
-    socket.reply({ userId: user!.id, userToken: token })
-  }
-
   async get(socket: SocketRequest) {
     socket.reply(await this.users.get(socket.userId))
   }

packages/socket/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@botpress/messaging-socket",
-  "version": "0.0.1",
+  "version": "0.0.2",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "source": "src/index.ts",