chat-service/src/socket/ws.controller.ts

import { App, Inject, OnWSConnection, OnWSDisConnection, OnWSMessage, WSController } from '@midwayjs/decorator';
import { Context } from '@midwayjs/ws';
import * as http from 'http';
import get = require('lodash/get');
import group = require('lodash/group');
import { Application } from '@midwayjs/ws';
import last = require('lodash/last');
import { Types } from 'mongoose';
const ObjectId = Types.ObjectId;
import { NeedSend } from '../entity/chat/needSend.entity';
import { ReturnModelType } from '@typegoose/typegoose';
import { InjectEntityModel } from '@midwayjs/typegoose';

@WSController('*/ws/*')
export class WsSocketController {
  @Inject()
  ctx: Context;
  @App('webSocket')
  wsApp: Application;
  @InjectEntityModel(NeedSend)
  model: ReturnModelType<typeof NeedSend>;

  @OnWSConnection()
  async onConnectionMethod(socket: Context, request: http.IncomingMessage) {
    const url = get(request, 'url');
    const arr = url.split('/');
    const token = last(arr);
    socket.send('connect success');
    const acs = this.getClients();
    // 最后一个是该websocket实例,赋上token
    const client = last(acs);
    client.token = token;
    await this.checkNeedSend(client);
  }
  /**
   * 给对象补发 未发送出去的消息
   * @param client ws连接实例
   */
  async checkNeedSend(client) {
    const to = get(client, 'token');
    if (!to) return;
    let list = await this.model.find({ to }).lean();
    list = list.map(i => ({ ...i, type: get(i, 'msg.type') }));
    const groups = group(list, 'type');
    for (const key in groups) {
      const list = groups[key];
      const type = get(last(list), 'type');
      await this.toSend({ type }, to);
    }
    await this.model.deleteMany({ to });
  }

  // 获取信息
  @OnWSMessage('message')
  async gotMessage(data: Buffer) {
    // const msg = data.toString();
    return '(╯‵□′)╯︵┻━┻';
  }

  @OnWSDisConnection()
  async disconnect(id) {
    // 断开连接
  }
  /**
   * 给指定对象发送消息
   * @param data 要发送的数据
   * @param token 店铺id/用户id
   */
  async toSend(data: object, token) {
    const clients = this.getClient(token);
    if (!clients) {
      // 存储到待发送表中
      await this.model.create({ to: token, msg: data });
      return;
    }
    clients.send(JSON.stringify(data));
  }
  /**
   * 获取所有连接实例
   */
  getClients(): Array<any> {
    const acs = [];
    const clients = this.wsApp.clients;
    clients.forEach(e => {
      acs.push(e);
    });
    return acs;
  }
  /**
   * 获取指定实例连接
   * @param token 用户id/店铺id
   */
  getClient(token: string) {
    const clients = this.getClients();
    const client = clients.find(f => new ObjectId(f.token).equals(token));
    return client;
  }
}