universo-platform-3d

1
import { Injectable, Logger } from '@nestjs/common'
2
import { Server, WebSocket } from 'ws'
3
import { FirebaseAuthenticationService } from '../firebase/firebase-authentication.service'
4
import { RedisPubSubService } from '../redis/redis-pub-sub.service'
5
import { CHANNELS } from '../redis/redis.channels'
6
import { v4 as uuidv4 } from 'uuid'
7

8
/**
9
 * @description This is used with WsAuthGuard to ensure that the user is authed before processing a request
10
 * See: https://github.com/nestjs/nest/issues/882#issuecomment-1493106283
11
 */
12
@Injectable()
13
export class WsAuthHelperService {
14
  constructor(
15
    private readonly logger: Logger,
16
    private readonly redisPubSubService: RedisPubSubService,
17
    private readonly firebaseAuthService: FirebaseAuthenticationService
18
  ) {}
19
  public initializationSuccess: { [key: string]: boolean } = {}
20
  private initializationMap = new Map<string, Promise<any>>()
21
  /**
22
   * Primary property for holding subscriptions to channels
23
   */
24
  private channelSubs: Record<string, WebSocket[]> = {}
25

26
  handleConnectionHelper(client: WebSocket, args: any) {
27
    this.initializationMap.set(client['id'], this.initialize(client, args))
28
  }
29

30
  async finishInitialization(client: WebSocket): Promise<any> {
31
    return await this.initializationMap.get(client['id'])
32
  }
33

34
  private async initialize(client: WebSocket, args: any): Promise<any> {
35
    // ensure max listeners is set high enough - nestjs bug (older version causes error)
36
    client.setMaxListeners(20)
37

38
    /** Crude authentication check until further defined */ const [
39
      { headers }
40
    ] = args
41
    let spaceId = headers?.space
42
    let token = headers?.authorization
43

44
    // attach the token to the client
45
    client['token'] = token
46
    // assign the socket a uuid with the uuid library
47
    client['id'] = uuidv4()
48

49
    let isFirebaseToken = false
50

51
    // if no token, check for Sec-WebSocket-Protocol header (this is how browsers have to send headers: see #5 here https://stackoverflow.com/a/77060459/3777933)
52
    if (!token) {
53
      const secWebSocketProtocol = headers['sec-websocket-protocol']
54
      if (secWebSocketProtocol) {
55
        // Important: if this order is changed, it must be changed in ws-auth-helper.service.ts on the react app too. it expects ordered array
56
        const result = secWebSocketProtocol.split(',')
57
        token = result[0]
58
        client['token'] = token
59
        if (result[1]) {
60
          spaceId = result[1]
61
        }
62
      }
63
    }
64

65
    /*
66
      check if the token is a valid firebase token
67
      if it is, set the user object on the client
68
      this data will be used to check the roles 
69
      we can catch headers when using WS only during the handleConnection event,
70
      so here we check the user's token and embed the decrypted token data in the user's connection object 
71
      when using the Firebase token, methods with role verification will be executed only
72
     */
73

74
    /* 
75
    if the token is the WSS_SECRET, set the role to admin
76
        and allows to use methods with admin role
77
        when using the WSS_SECRET token, methods will be executed only without role verification
78
      */
79
    if (token === process.env.WSS_SECRET) {
80
      client['role'] = 'admin'
81
    } else {
82
      try {
83
        const decodedJwt = await this.decodeJwt(token)
84
        if (decodedJwt) {
85
          isFirebaseToken = true
86
          client['user'] = decodedJwt
87
          console.log('isFirebaseToken', isFirebaseToken)
88
        }
89
      } catch (error) {
90
        this.logger.log(
91
          `Invalid tokens (WSS_SECRET or Firebase JWT) for GodotGateway handleConnection: Disconnecting`,
92
          WsAuthHelperService.name
93
        )
94
        return client.close(1014, 'Invalid token')
95
      }
96
    }
97

98
    // if the token is neither the WSS_SECRET nor a valid firebase token, close the connection
99
    if (!token || (token !== process.env.WSS_SECRET && !isFirebaseToken)) {
100
      const msg =
101
        'Invalid tokens (WSS_SECRET or Firebase JWT) for GodotGateway handleConnection: Disconnecting'
102
      this.logger.error(msg, WsAuthHelperService.name)
103
      return client.close(1014, 'Invalid token')
104
    }
105

106
    this.logger.log(
107
      `handleConnection: ${JSON.stringify(
108
        {
109
          spaceId
110
        },
111
        null,
112
        2
113
      )}`,
114
      WsAuthHelperService.name
115
    )
116
    this.setupSubscriber(client, spaceId)
117

118
    this.initializationMap.delete(client['id'])
119
    this.initializationSuccess[client['id']] = true
120
    return Promise.resolve()
121
  }
122

123
  setupSubscriber(client: WebSocket, spaceId: string) {
124
    if (!spaceId) {
125
      this.logger.log(
126
        `setupSubscriber: attempted but spaceId was falsey: ${JSON.stringify(
127
          {
128
            spaceId
129
          },
130
          null,
131
          2
132
        )}`,
133
        WsAuthHelperService.name
134
      )
135
      return
136
    }
137
    const subChannel = `${CHANNELS.SPACE}:${spaceId}`
138
    client['subscriberChannel'] = subChannel
139
    // Add to tracked channels
140
    if (!this.channelSubs[subChannel]) {
141
      this.channelSubs[subChannel] = new Array<WebSocket>()
142
    }
143
    this.channelSubs[subChannel].push(client)
144
    // setup the subscription listener if this is the first subscriber to the channel.
145
    if (this.channelSubs[subChannel].length == 1) {
146
      this.redisPubSubService.subscriber.subscribe(subChannel, (message) => {
147
        this.handleSubscribedChannelReceivedMessage(subChannel, message)
148
      })
149
    }
150
  }
151

152
  /**
153
   *
154
   * @description Called when a Redis pubsub message is received from a channel that was subscribed to
155
   */
156
  handleSubscribedChannelReceivedMessage(subChannel: string, message: string) {
157
    this.logger.log(
158
      `handleSubscribedChannelReceivedMessage: ${JSON.stringify(
159
        {
160
          subChannel,
161
          message
162
        },
163
        null,
164
        2
165
      )}\n
166
      Sending message to subscribers: ${this.channelSubs[subChannel].length}`,
167
      WsAuthHelperService.name
168
    )
169
    this.channelSubs[subChannel].forEach((client) => {
170
      client.send(message)
171
    })
172
  }
173

174
  removeSubscriber(client: WebSocket) {
175
    // remove the client from the channel subscriptions.
176
    const subchannel = client['subscriberChannel']
177
    if (!subchannel || !this.channelSubs[subchannel]) {
178
      this.logger.log(
179
        `setupSubscriber: attempted but subchannel was falsey: ${JSON.stringify(
180
          {
181
            client,
182
            subchannel
183
          },
184
          null,
185
          2
186
        )}`,
187
        WsAuthHelperService.name
188
      )
189
      return
190
    }
191
    this.channelSubs[subchannel] = this.channelSubs[subchannel].filter(
192
      (c1) => c1 !== client
193
    )
194
    // unsubscribe from the subscription listener if this is the last subscriber to the channel.
195
    if (this.channelSubs[subchannel].length == 0) {
196
      this.redisPubSubService.subscriber.unsubscribe(subchannel)
197
    }
198
  }
199

200
  async decodeJwt(token) {
201
    return await this.firebaseAuthService.verifyIdToken(
202
      token ? token.replace('Bearer ', '') : '',
203
      true
204
    )
205
  }
206
}
207