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:websocket');
|
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', (err) => {
|
25 | this.emit('error', err);
|
26 | debug('Error occurred:', err);
|
27 | });
|
28 | const fn = co.wrap(compose(this.middleware));
|
29 |
|
30 | const context = this.app.createContext(socket.upgradeReq);
|
31 | context.websocket = socket;
|
32 | context.path = url.parse(socket.upgradeReq.url).pathname;
|
33 |
|
34 | fn.bind(context).call().catch((err) => {
|
35 | debug(err);
|
36 | });
|
37 | };
|
38 |
|
39 | KoaWebSocketServer.prototype.use = function (fn) {
|
40 | this.middleware.push(fn);
|
41 | return this;
|
42 | };
|
43 |
|
44 | module.exports = function (app) {
|
45 | const oldListen = app.listen;
|
46 | app.listen = function () {
|
47 | debug('Attaching server...');
|
48 | app.server = oldListen.apply(app, arguments);
|
49 | app.ws.listen(app.server);
|
50 | return app.server;
|
51 | };
|
52 | app.ws = new KoaWebSocketServer(app);
|
53 | return app;
|
54 | };
|