adding cls support for ws

src/middleware/cls.js

@@ -1,23 +1,26 @@
-
+var path = require('path');
 var continuationLocalStorage = require('continuation-local-storage');
+var APP_NAMESPACE = require(path.join(__dirname, '../../package.json')).name;
+var namespace = continuationLocalStorage.createNamespace(APP_NAMESPACE);
 
-var NAMESPACE = 'nodebb';
-var namespace = continuationLocalStorage.createNamespace(NAMESPACE);
+var cls = {};
 
-var cls = function (req, res, next) {
+cls.http = function (req, res, next) {
 	namespace.run(function() {
-		var routeData = {req: req, res: res};
-
-		if (process.env.NODE_ENV == 'development') {
-			routeData.audit = {created: process.hrtime()};
-		}
-		namespace.set('route', routeData);
+		namespace.set('http', {req: req, res: res});
 		next();
 	});
 };
 
-cls.storage = function () {
-	return cls.getNamespace(NAMESPACE);
+cls.socket = function (socket, payload, event, next) {
+	namespace.run(function() {
+		namespace.set('ws', {
+			socket: socket,
+			payload: payload,
+			// if it's a '*' event, then we grab it from the payload
+			event: event || ((payload || {}).data || [])[0]});
+		next();
+	});
 };
 
 cls.get = function (key) {

src/middleware/index.js

@@ -74,7 +74,7 @@ module.exports = function(app) {
 
 	app.use(middleware.addHeaders);
 	app.use(middleware.processRender);
-	app.use(cls);
+	app.use(cls.http);
 	auth.initialize(app, middleware);
 
 	return middleware;

src/socket.io/index.js

@@ -11,6 +11,7 @@ var	SocketIO = require('socket.io'),
 	user = require('../user'),
 	logger = require('../logger'),
 	ratelimit = require('../middleware/ratelimit'),
+	cls = require('../middleware/cls'),
 
 	Sockets = {},
 	Namespaces = {};
@@ -43,14 +44,20 @@ function onConnection(socket) {
 
 	logger.io_one(socket, socket.uid);
 
-	onConnect(socket);
+	cls.socket(socket, null, 'connection', function () {
+		onConnect(socket);
+	});
 
-	socket.on('disconnect', function(data) {
-		onDisconnect(socket, data);
+	socket.on('disconnect', function(payload) {
+		cls.socket(socket, payload, 'disconnect', function () {
+			onDisconnect(socket, payload);
+		});
 	});
 
 	socket.on('*', function(payload) {
-		onMessage(socket, payload);
+		cls.socket(socket, payload, null, function() {
+			onMessage(socket, payload);
+		});
 	});
 }