[RFC] Return 401 for an authentication error on WebSockets (#3411)
* Return 401 for an authentication error on WebSocket * Use upgradeReq instead of a custom object
This commit is contained in:
		
							parent
							
								
									5e2c5e95b6
								
							
						
					
					
						commit
						9a81be0d37
					
				
					 1 changed files with 48 additions and 39 deletions
				
			
		|  | @ -95,7 +95,6 @@ const startWorker = (workerId) => { | |||
|   const app    = express(); | ||||
|   const pgPool = new pg.Pool(Object.assign(pgConfigs[env], dbUrlToConfig(process.env.DATABASE_URL))); | ||||
|   const server = http.createServer(app); | ||||
|   const wss    = new WebSocket.Server({ server }); | ||||
|   const redisNamespace = process.env.REDIS_NAMESPACE || null; | ||||
| 
 | ||||
|   const redisParams = { | ||||
|  | @ -186,14 +185,10 @@ const startWorker = (workerId) => { | |||
|     }); | ||||
|   }; | ||||
| 
 | ||||
|   const authenticationMiddleware = (req, res, next) => { | ||||
|     if (req.method === 'OPTIONS') { | ||||
|       next(); | ||||
|       return; | ||||
|     } | ||||
| 
 | ||||
|     const authorization = req.get('Authorization'); | ||||
|     const accessToken = req.query.access_token; | ||||
|   const accountFromRequest = (req, next) => { | ||||
|     const authorization = req.headers.authorization; | ||||
|     const location = url.parse(req.url, true); | ||||
|     const accessToken = location.query.access_token; | ||||
| 
 | ||||
|     if (!authorization && !accessToken) { | ||||
|       const err = new Error('Missing access token'); | ||||
|  | @ -208,6 +203,26 @@ const startWorker = (workerId) => { | |||
|     accountFromToken(token, req, next); | ||||
|   }; | ||||
| 
 | ||||
|   const wsVerifyClient = (info, cb) => { | ||||
|     accountFromRequest(info.req, err => { | ||||
|       if (!err) { | ||||
|         cb(true, undefined, undefined); | ||||
|       } else { | ||||
|         log.error(info.req.requestId, err.toString()); | ||||
|         cb(false, 401, 'Unauthorized'); | ||||
|       } | ||||
|     }); | ||||
|   }; | ||||
| 
 | ||||
|   const authenticationMiddleware = (req, res, next) => { | ||||
|     if (req.method === 'OPTIONS') { | ||||
|       next(); | ||||
|       return; | ||||
|     } | ||||
| 
 | ||||
|     accountFromRequest(req, next); | ||||
|   }; | ||||
| 
 | ||||
|   const errorMiddleware = (err, req, res, next) => { | ||||
|     log.error(req.requestId, err.toString()); | ||||
|     res.writeHead(err.statusCode || 500, { 'Content-Type': 'application/json' }); | ||||
|  | @ -352,10 +367,12 @@ const startWorker = (workerId) => { | |||
|     streamFrom(`timeline:hashtag:${req.query.tag}:local`, req, streamToHttp(req, res), streamHttpEnd(req), true); | ||||
|   }); | ||||
| 
 | ||||
|   const wss    = new WebSocket.Server({ server, verifyClient: wsVerifyClient }); | ||||
| 
 | ||||
|   wss.on('connection', ws => { | ||||
|     const location = url.parse(ws.upgradeReq.url, true); | ||||
|     const token    = location.query.access_token; | ||||
|     const req      = { requestId: uuid.v4() }; | ||||
|     const req      = ws.upgradeReq; | ||||
|     const location = url.parse(req.url, true); | ||||
|     req.requestId  = uuid.v4(); | ||||
| 
 | ||||
|     ws.isAlive = true; | ||||
| 
 | ||||
|  | @ -363,33 +380,25 @@ const startWorker = (workerId) => { | |||
|       ws.isAlive = true; | ||||
|     }); | ||||
| 
 | ||||
|     accountFromToken(token, req, err => { | ||||
|       if (err) { | ||||
|         log.error(req.requestId, err); | ||||
|         ws.close(); | ||||
|         return; | ||||
|       } | ||||
| 
 | ||||
|       switch(location.query.stream) { | ||||
|       case 'user': | ||||
|         streamFrom(`timeline:${req.accountId}`, req, streamToWs(req, ws), streamWsEnd(req, ws)); | ||||
|         break; | ||||
|       case 'public': | ||||
|         streamFrom('timeline:public', req, streamToWs(req, ws), streamWsEnd(req, ws), true); | ||||
|         break; | ||||
|       case 'public:local': | ||||
|         streamFrom('timeline:public:local', req, streamToWs(req, ws), streamWsEnd(req, ws), true); | ||||
|         break; | ||||
|       case 'hashtag': | ||||
|         streamFrom(`timeline:hashtag:${location.query.tag}`, req, streamToWs(req, ws), streamWsEnd(req, ws), true); | ||||
|         break; | ||||
|       case 'hashtag:local': | ||||
|         streamFrom(`timeline:hashtag:${location.query.tag}:local`, req, streamToWs(req, ws), streamWsEnd(req, ws), true); | ||||
|         break; | ||||
|       default: | ||||
|         ws.close(); | ||||
|       } | ||||
|     }); | ||||
|     switch(location.query.stream) { | ||||
|     case 'user': | ||||
|       streamFrom(`timeline:${req.accountId}`, req, streamToWs(req, ws), streamWsEnd(req, ws)); | ||||
|       break; | ||||
|     case 'public': | ||||
|       streamFrom('timeline:public', req, streamToWs(req, ws), streamWsEnd(req, ws), true); | ||||
|       break; | ||||
|     case 'public:local': | ||||
|       streamFrom('timeline:public:local', req, streamToWs(req, ws), streamWsEnd(req, ws), true); | ||||
|       break; | ||||
|     case 'hashtag': | ||||
|       streamFrom(`timeline:hashtag:${location.query.tag}`, req, streamToWs(req, ws), streamWsEnd(req, ws), true); | ||||
|       break; | ||||
|     case 'hashtag:local': | ||||
|       streamFrom(`timeline:hashtag:${location.query.tag}:local`, req, streamToWs(req, ws), streamWsEnd(req, ws), true); | ||||
|       break; | ||||
|     default: | ||||
|       ws.close(); | ||||
|     } | ||||
|   }); | ||||
| 
 | ||||
|   const wsInterval = setInterval(() => { | ||||
|  |  | |||
		Loading…
	
	Add table
		
		Reference in a new issue