77 lines
1.7 KiB
JavaScript
77 lines
1.7 KiB
JavaScript
var url = require('url');
|
|
var urlJoin = require('url-join');
|
|
var http = require('http');
|
|
var ServerResponse = http.ServerResponse;
|
|
var WebSocketServer = require('ws').Server;
|
|
|
|
var wsServers = {};
|
|
|
|
/**
|
|
* @param {express.Application} app
|
|
* @param {http.Server} [server]
|
|
*/
|
|
module.exports = function (app, server) {
|
|
if(!server) {
|
|
server = http.createServer(app);
|
|
|
|
app.listen = function()
|
|
{
|
|
return server.listen.apply(server, arguments)
|
|
}
|
|
}
|
|
|
|
function addSocketRoute(route, middleware, callback) {
|
|
var args = [].splice.call(arguments, 0);
|
|
var wsPath = urlJoin(app.mountpath, route);
|
|
|
|
if (args.length < 2)
|
|
throw new SyntaxError('Invalid number of arguments');
|
|
|
|
if (args.length === 2) {
|
|
middleware = [middleware];
|
|
} else if (typeof middleware === 'object') {
|
|
middleware.push(callback);
|
|
} else {
|
|
middleware = args.slice(1);
|
|
}
|
|
|
|
var wss = new WebSocketServer({
|
|
server: server,
|
|
path: wsPath
|
|
});
|
|
|
|
wsServers[wsPath] = wss;
|
|
|
|
wss.on('connection', function(ws) {
|
|
var response = new ServerResponse(ws.upgradeReq);
|
|
response.writeHead = function (statusCode) {
|
|
if (statusCode > 200) ws.close();
|
|
};
|
|
ws.upgradeReq.method = 'ws';
|
|
|
|
app.handle(ws.upgradeReq, response, function(err) {
|
|
var idx = 0;
|
|
(function next (err) {
|
|
if (err) return;
|
|
var cur = middleware[idx++];
|
|
if (!middleware[idx]) {
|
|
cur(ws, ws.upgradeReq);
|
|
} else {
|
|
cur(ws.upgradeReq, response, next);
|
|
}
|
|
}(err));
|
|
});
|
|
});
|
|
|
|
return app;
|
|
};
|
|
|
|
app.ws = addSocketRoute;
|
|
|
|
return {
|
|
app: app,
|
|
getWss: function (route) {
|
|
return wsServers[route];
|
|
}
|
|
};
|
|
};
|