1 files changed, 25 insertions, 2 deletions

diff --git a/server/middlewares/auth.ts b/server/middlewares/auth.ts
index e6025c8ce..0eefa2a8e 100644
--- a/server/middlewares/auth.ts
+++ b/server/middlewares/auth.ts
@@ -1,6 +1,7 @@
 import express from 'express'
 import { Socket } from 'socket.io'
 import { getAccessToken } from '@server/lib/auth/oauth-model'
+import { RunnerModel } from '@server/models/runner/runner'
 import { HttpStatusCode } from '../../shared/models/http/http-error-codes'
 import { logger } from '../helpers/logger'
 import { handleOAuthAuthenticate } from '../lib/auth/oauth'
@@ -27,7 +28,7 @@ function authenticate (req: express.Request, res: express.Response, next: expres
 function authenticateSocket (socket: Socket, next: (err?: any) => void) {
   const accessToken = socket.handshake.query['accessToken']
 
-  logger.debug('Checking socket access token %s.', accessToken)
+  logger.debug('Checking access token in runner.')
 
   if (!accessToken) return next(new Error('No access token provided'))
   if (typeof accessToken !== 'string') return next(new Error('Access token is invalid'))
@@ -73,9 +74,31 @@ function optionalAuthenticate (req: express.Request, res: express.Response, next
 
 // ---------------------------------------------------------------------------
 
+function authenticateRunnerSocket (socket: Socket, next: (err?: any) => void) {
+  const runnerToken = socket.handshake.auth['runnerToken']
+
+  logger.debug('Checking runner token in socket.')
+
+  if (!runnerToken) return next(new Error('No runner token provided'))
+  if (typeof runnerToken !== 'string') return next(new Error('Runner token is invalid'))
+
+  RunnerModel.loadByToken(runnerToken)
+    .then(runner => {
+      if (!runner) return next(new Error('Invalid runner token.'))
+
+      socket.handshake.auth.runner = runner
+
+      return next()
+    })
+    .catch(err => logger.error('Cannot get runner token.', { err }))
+}
+
+// ---------------------------------------------------------------------------
+
 export {
   authenticate,
   authenticateSocket,
   authenticatePromise,
-  optionalAuthenticate
+  optionalAuthenticate,
+  authenticateRunnerSocket
 }