1 | 'use strict';
|
2 |
|
3 | const url = require('url'),
|
4 | compose = require('koa-compose'),
|
5 | co = require('co'),
|
6 | ws = require('ws');
|
7 | const WebSocketServer = ws.Server;
|
8 | const debug = require('debug')('koa:websockets');
|
9 |
|
10 | function KoaWebSocketServer (app) {
|
11 | this.app = app;
|
12 | this.middleware = [];
|
13 | }
|
14 |
|
15 | KoaWebSocketServer.prototype.listen = function (server) {
|
16 | this.server = new WebSocketServer({
|
17 | server: server
|
18 | });
|
19 | this.server.on('connection', this.onConnection.bind(this));
|
20 | };
|
21 |
|
22 | KoaWebSocketServer.prototype.onConnection = function(socket) {
|
23 | debug('Connection received');
|
24 | socket.on('error', function (err) {
|
25 | debug('Error occurred:', err);
|
26 | });
|
27 | const fn = co.wrap(compose(this.middleware));
|
28 |
|
29 | const context = this.app.createContext(socket.upgradeReq);
|
30 | context.websocket = socket;
|
31 | context.path = url.parse(socket.upgradeReq.url).pathname;
|
32 |
|
33 | fn.bind(context).call().catch(function(err) {
|
34 | debug(err);
|
35 | });
|
36 | };
|
37 |
|
38 | KoaWebSocketServer.prototype.use = function (fn) {
|
39 | this.middleware.push(fn);
|
40 | return this;
|
41 | };
|
42 |
|
43 | module.exports = function (app) {
|
44 | const oldListen = app.listen;
|
45 | app.listen = function () {
|
46 | debug('Attaching server...');
|
47 | app.server = oldListen.apply(app, arguments);
|
48 | app.ws.listen(app.server);
|
49 | return app.server;
|
50 | };
|
51 | app.ws = new KoaWebSocketServer(app);
|
52 | return app;
|
53 | };
|