koaでwebsocketサーバーを構築する

koa

メジャーな Express.js を開発したチームによって作成されたNode.jsのフレームワークで、
Express.jsの後継にあたるそうです。今回はこのkoaを使ってwebsoketサーバーを実装してみたいと思います。

koa-websocket

koaを使用してwebsocketサーバーを構築する際のmiddleware部分のパッケージになります。
調べた結果あまりkoa + websocket構成でのmiddlewareパッケージがなく、こちらが一番使われていそうだったので試してみる事にしました。

サンプルプロジェクトの作成

今回は以下の構成でプロジェクト作成しました。

.
├── Dockerfile
├── docker-compose.yml
├── index.ts
├── node_modules
├── nodemon.json
├── package.json
├── tsconfig.json
└── yarn.lock

主要なファイルの詳細は以下になります。

Dockerfile

FROM node:16-alpine
WORKDIR /usr/src/app

RUN apk update && \
    apk add git vim bash

COPY package*.json yarn.lock ./
RUN yarn install
COPY . ./
CMD [ "yarn", "start" ]

docker-compose.yml

version: '3.3'
volumes:
  modules_data:
    driver: local

services:
  wss:
    build: .
    image: koa/websocket
    container_name: "koa-websocket"
    command: sh -c "yarn install && yarn dev"
    tty: true
    volumes:
      - .:/usr/src/app
      - modules_data:/usr/src/app/node_modules
    ports:
      - '8080:8080'
    working_dir: /usr/src/app

package.json

{
  "name": "koa-websocket-sample",
  "version": "1.0.0",
  "main": "index.js",
  "license": "MIT",
  "scripts": {
    "dev": "nodemon",
    "start": "ts-node index.ts"
  },
  "dependencies": {},
  "devDependencies": {
    "@types/node": "^18.0.0",
    "nodemon": "^2.0.16",
    "ts-node": "^10.8.1",
    "typescript": "^4.7.3"
  }
}

nodemon.json

{
  "watch": ["."],
   "ext": "ts",
  "exec": "ts-node index.ts",
  "env": {
    "__LOCAL__": "true"
  }
}

パッケージインストール

$ yarn add koa koa-websocket
$ yarn add -D @types/koa @types/koa-websocket

今回のテーマとは関係無いですが、utility用のパッケージとして radash も追加してます。
radashに関してはこちらの記事に詳細があります。

$ yarn add radash

koa-compose

KoaのMiddlewareを構築する際のutilityパッケージです。koaをインストールすると使えます。

fn = compose([a, b, c, ...])

複数のMiddlewareを合成に1つのMiddlewareを返す compose を提供します。

簡単なサンプル実装

index.ts

import Koa from 'koa';
import { websockify, WebSocketServer } from 'koa-websocket';
import { toInt } from 'radash';
import { IncomingMessage } from 'http';

const port = toInt(process.env.PORT, 8080);

const app = websockify(new Koa());
const server: WebSocketServer = app.ws;

server.use((ctx: Koa.Context) => {
  ctx.websocket.on('open', (request: IncomingMessage) => {
    console.log('ctx.websocket connection open');
  });

  ctx.websocket.on('close', () => {
    console.log('ctx.websocket connection close');
  });

  ctx.websocket.on('upgrade', () => {
    console.log('ctx.websocket connection upgrade');
  });

  ctx.websocket.on('ping', () => {
    console.log('ctx.websocket connection ping');
  });

  ctx.websocket.on('message', (message: any) => {
    console.log('ctx.websocket message', message);
  });
});

app.listen(port);

この状態でコンテナ起動し、クライアントから接続すると message 受信時にログが出力されているかと思います。

問題点

ここで、koa-websocket の実装を見ていると接続時のイベントで何か処理を実施したくても、そのイベントが取れなさそうです。。

  ctx.websocket.on('open', (request: IncomingMessage) => {
    console.log('ctx.websocket connection open');
  });

↑こちらのイベントも一切呼ばれないので、このままだと接続時の処理が出来なさそうです。折角なのでTypescriptで書き直しつつカスタムしてみようかと思います。

Typescriptで書き直しつつカスタム

lib/ws/server.ts に以下の内容で実装します。
使う場合は koa-websocket のインポート部分を以下に修正します。

import { websockify, WebSocketServer } from './lib/ws/server';

lib/ws/server.ts

import ws from 'ws';
import Koa from 'koa';
import url from 'url';
import { IncomingMessage, ServerResponse } from 'http';
import compose from 'koa-compose';
import * as https from 'https';

declare module 'koa' {
  interface Context {
    websocket: ws;
    path: string;
  }
}

const wsServer = ws.Server;

export class WebSocketServer {
  private _app: Koa;
  private _server: ws.Server<ws.WebSocket> | undefined = undefined;
  private _middlewares: any[] = [];

  constructor(app: Koa) {
    this._app = app;
  }

  listen(options?: ws.ServerOptions | undefined) {
    this._server = new wsServer(options);
    this._server.on('connection', this.onConnection.bind(this));
  }

  onConnection(socket: ws.WebSocket, request: IncomingMessage) {
    const fn = compose(this._middlewares);

    const context = this._app.createContext(
      request,
      new ServerResponse(request)
    );
    context.websocket = socket;
    context.path = (request.url && url.parse(request.url).pathname) ?? '';

    fn(context)
      .then(() => socket.emit('open', request))
      .catch((e) => console.error(e));
  }

  use(fn: any) {
    this._middlewares.push(fn);
  }
}

export const websockify = (
  app: any,
  wsOptions?: ws.ServerOptions,
  httpsOptions?: https.ServerOptions
) => {
  const oldListen = app.listen;
  app.listen = function listen(...args: any[]) {
    if (typeof httpsOptions === 'object') {
      const httpsServer = https.createServer(httpsOptions, app.callback());
      app.server = httpsServer.listen(...(args as any));
    } else {
      app.server = oldListen.apply(app, args as any);
    }
    const options: { [key: string]: any } = { server: app.server };
    if (wsOptions) {
      Object.keys(wsOptions).forEach((key) => {
        if (Object.prototype.hasOwnProperty.call(wsOptions, key)) {
          options[key] = (wsOptions as any)[key];
        }
      });
    }
    app.ws.listen(options);
    return app.server;
  };
  app.ws = new WebSocketServer(app);
  return app;
};

以下の部分で接続時に open を呼ぶようにしています。

    fn(context)
      .then(() => socket.emit('open', request))
      .catch((e) => console.error(e));

これで

  ctx.websocket.on('open', (request: IncomingMessage) => {
    console.log('ctx.websocket connection open');
  });

↑が接続時に呼ばれる様になったかと思います。

まとめ

こうみると実際にkoa + websocket構成でwebsocketサーバー実装する場合は、外部パッケージを使うより、Middleware部分は自前で実装した方が、実装量や後々のメンテしやすさを考えると良さそうかなと思いました。