finish gateway websocket example

client/components/home/ping.jsx

@@ -1,12 +1,74 @@
-import { useState } from 'react';
+import { useState, useEffect, useRef, useContext } from 'react';
 import { Button } from '../common/button';
+import { io } from 'socket.io-client';
+import { AuthContext } from '../../utils/auth_context';
 
 export const Ping = () => {
   const [pings, setPings] = useState([]);
   const [key, setKey] = useState('defaultkey');
+  const [currentRoom, setCurrentRoom] = useState(null);
+  const [authToken] = useContext(AuthContext);
+  const [socket, setSocket] = useState(null);
+
+  useEffect(() => {
+    // instantiates a socket object and initiates the connection...
+    // you probably want to make sure you are only doing this in one component at a time.
+    const socket = io({
+      auth: { token: authToken },
+      query: { message: 'I am the query ' },
+    });
+
+    // adds an event listener to the connection event
+    socket.on('connect', () => {
+      setSocket(socket);
+    });
+
+    // adds event listener to the disconnection event
+    socket.on('disconnect', () => {
+      console.log('Disconnected');
+    });
+
+    // recieved a pong event from the server
+    socket.on('pong', (data) => {
+      console.log('Recieved pong', data);
+    });
+
+    // IMPORTANT! Unregister from all events when the component unmounts and disconnect.
+    return () => {
+      socket.off('connect');
+      socket.off('disconnect');
+      socket.off('pong');
+      socket.disconnect();
+    };
+  }, []);
+
+  useEffect(() => {
+    // if our token changes we need to tell the socket also
+    if (socket) {
+      // this is a little weird because we are modifying this object in memory
+      // i dunno a better way to do this though...
+      socket.auth.token = authToken;
+    }
+  }, [authToken]);
+
+  if (!socket) return 'Loading...';
+
+  const sendPing = () => {
+    // sends a ping to the server to be broadcast to everybody in the room
+    currentRoom && socket.emit('ping', { currentRoom });
+  };
+
+  const joinRoom = () => {
+    // tells the server to remove the current client from the current room and add them to the new room
+    socket.emit('join-room', { currentRoom, newRoom: key }, (response) => {
+      console.log(response);
+      setCurrentRoom(response.room);
+    });
+  };
+
   return (
     <>
-      <header>Ping</header>
+      <header>Ping: {currentRoom || '(No room joined)'}</header>
       <section>
         <input
           type="text"
@@ -14,8 +76,8 @@ export const Ping = () => {
           value={key}
           onChange={(e) => setKey(e.target.value)}
         />
-        <Button>Ping Connect</Button>
+        <Button onClick={joinRoom}>Connect To Room</Button>
-        <Button>Send Ping</Button>
+        <Button onClick={sendPing}>Send Ping</Button>
       </section>
     </>
   );

client/utils/use_jwt_refresh.js

@@ -12,7 +12,7 @@ export const useJwtRefresh = (authToken, setAuthToken) => {
         } else {
           setAuthToken(null);
         }
-      }, 60000 * 0.5); // 10 minutes
+      }, 60000 * 10); // 10 minutes
     }
     return () => clearTimeout(refreshTimer.current);
   }, [authToken]);

server/app.module.ts

@@ -4,6 +4,7 @@ import { TypeOrmModule } from '@nestjs/typeorm';
 import { AppController } from './app.controller';
 import { config } from './database/config';
 import { UsersModule } from './modules/users.module';
+import { PingGateway } from './providers/gateways/ping.gateway';
 import { AuthGuard } from './providers/guards/auth.guard';
 import { RolesGuard } from './providers/guards/roles.guard';
 import { JwtService } from './providers/services/jwt.service';
@@ -15,6 +16,7 @@ import { GuardUtil } from './providers/util/guard.util';
   imports: [TypeOrmModule.forRoot(config), UsersModule],
   controllers: [AppController],
   providers: [
+    PingGateway,
     UsersService,
     RolesService,
     JwtService,

server/decorators/gateway_jwt_body.decorator.ts

@@ -0,0 +1,6 @@
+import { createParamDecorator, ExecutionContext } from '@nestjs/common';
+import { Socket } from 'socket.io';
+export const GatewayJwtBody = createParamDecorator((data: unknown, ctx: ExecutionContext) => {
+  const req = ctx.switchToHttp().getRequest() as Socket;
+  return req.handshake.auth.jwtBody;
+});

server/providers/gateways/ping.gateway.ts

@@ -1,11 +1,67 @@
-import { SubscribeMessage, WebSocketGateway } from '@nestjs/websockets';
+import { UseGuards } from '@nestjs/common';
+import { ConnectedSocket, MessageBody, OnGatewayConnection, OnGatewayDisconnect, OnGatewayInit, SubscribeMessage, WebSocketGateway, WebSocketServer, WsException } from '@nestjs/websockets';
+import { GatewayJwtBody } from 'server/decorators/gateway_jwt_body.decorator';
+import { JwtBodyDto } from 'server/dto/jwt_body.dto';
+import { Server, Socket } from 'socket.io';
+import { GatewayAuthGuard } from '../guards/gatewayauth.guard';
+import { JwtService } from '../services/jwt.service';
+
+class JoinPayload {
+  currentRoom?: string;
+  newRoom: string;
+}
+
+class PingPayload {
+  currentRoom: string;
+}
 
 @WebSocketGateway()
-export class PingGateway {
+@UseGuards(GatewayAuthGuard)
-  constructor() {}
+export class PingGateway implements OnGatewayInit, OnGatewayConnection, OnGatewayDisconnect {
+  @WebSocketServer()
+  server: Server;
 
-  @SubscribeMessage('ping/:id')
+  constructor(private jwtService: JwtService) {}
-  public handlePing() {
+
-    return { message: 'recieved ping' };
+  afterInit(server: Server) {
+    console.log('Sockets initialized');
+  }
+
+  handleConnection(client: Socket) {
+    // you can do things like add users to rooms
+    // or emit events here.
+    // IMPORTANT! The GatewayAuthGuard doesn't trigger on these handlers
+    // if you need to do anything in this method you need to authenticate the JWT
+    // manually.
+    try {
+      const jwt = client.handshake.auth.token;
+      const jwtBody = this.jwtService.parseToken(jwt);
+      console.log(client.handshake.query);
+      console.log('Client Connected: ', jwtBody.userId);
+    } catch (e) {
+      throw new WsException('Invalid token');
+    }
+  }
+
+  handleDisconnect(client: Socket) {
+    console.log('Client Disconnected');
+  }
+
+  @SubscribeMessage('ping')
+  public handlePing(
+    @ConnectedSocket() client: Socket,
+    @MessageBody() payload: PingPayload,
+    @GatewayJwtBody() jwtBody: JwtBodyDto,
+  ) {
+    this.server.to(payload.currentRoom).emit('pong', { message: { userId: jwtBody.userId } });
+    console.log(client.rooms);
+  }
+
+  @SubscribeMessage('join-room')
+  public async joinRoom(client: Socket, payload: JoinPayload) {
+    console.log(payload);
+    payload.currentRoom && (await client.leave(payload.currentRoom));
+    await client.join(payload.newRoom);
+    return { msg: 'Joined room', room: payload.newRoom };
   }
 }

server/providers/guards/gatewayauth.guard.ts

@@ -0,0 +1,27 @@
+import { Injectable, CanActivate, ExecutionContext } from '@nestjs/common';
+import { JwtService } from '../services/jwt.service';
+import { GuardUtil } from '../util/guard.util';
+import { Socket } from 'socket.io';
+import { WsException } from '@nestjs/websockets';
+
+@Injectable()
+export class GatewayAuthGuard implements CanActivate {
+  constructor(private guardUtil: GuardUtil, private jwtService: JwtService) {}
+
+  canActivate(context: ExecutionContext) {
+    // Handlers and Controllers can both skip this guard in the event that
+    if (this.guardUtil.shouldSkip(this, context)) {
+      return true;
+    }
+
+    const req = context.switchToHttp().getRequest() as Socket;
+    const jwt = req.handshake.auth.token;
+    if (!jwt) throw new WsException('Invalid auth token');
+    try {
+      req.handshake.auth.jwtBody = this.jwtService.parseToken(jwt);
+    } catch (e) {
+      throw new WsException('Invalid auth token');
+    }
+    return true;
+  }
+}