Nest JS Websockets - Rate Limiting and Acknowledgements

Welcome to Part 5: Rate Limiting, of building a realtime chat application with Nest JS and React.

In this part of the series, we’re going to implement rate limiting! With rate limiting, we will likely be denying users from sending messages - so to be respectful we should indicate to the user whether or not their message was delivered. To achieve this, we will also be introducing a new mechanism for our chat events called acknowledgements in socket.io parlance.

If you haven’t checkout out previous parts of this series, go back to follow along.

Part 1: Basics - laying the foundations of the application and implementing the most basic websockets features.
Part 2: Rooms - implementing websocket rooms and major UI improvements.
Part 3: Pipes and E2E Validation with Zod - validate mutations with Zod schemas shared between client and server. From form actions, to emitting client events, to handling events on the server - giving our application multiple levels of protection from malformed data.
Part 4: Guards and Authorization with CASL - attribute-based authorization with CASL and guards.

This series will attempt to touch on as many features of Nest's websockets integration as well as socket.io in general.

The entire series will be building on one repository that has both the server and the client for the application, so for each part you can download the repository from specific commits to follow along or you can clone the entire repository in it's finished state.

Download the repository at this commit to follow along - we’re not going to cover everything to keep the focus on rate limiting and acknowledgements.

Rate Limiting

Rate limiting is a process by which we can limit a user or agent from performing an excessive amount of actions in a given time frame. This is a critical part of developing any kind of application to prevent abusive use of your server or database and is one way to mitigate basic DDoS attacks.

Throttler Guard

Nest provides easy-to-use support for rate limiting with websockets. We will implement a WsThrottlerGuard by extending ThrottlerGuard from @nestjs/throttler.

In our guard’s handleRequest method, we have access to the execution context from which we can get a reference to the client’s IP address. We then use built-in methods from ThrottlerGuard to store the amount of times a client has sent a websocket event and we compare that with a ttl (time-to-live) that is passed to the guard.

src/server/chat/guards/throttler.guard.ts

import { Injectable, ExecutionContext } from '@nestjs/common'
import { ThrottlerGuard, ThrottlerException } from '@nestjs/throttler'

@Injectable()
export class WsThrottlerGuard extends ThrottlerGuard {
  async handleRequest(context: ExecutionContext, limit: number, ttl: number): Promise<boolean> {
    const client = context.switchToWs().getClient()
    const ip = client.conn.remoteAddress
    const key = this.generateKey(context, ip)
    const ttls = await this.storageService.getRecord(key)

    if (ttls.length >= limit) {
      throw new ThrottlerException()
    }

    await this.storageService.addRecord(key, ttl)
    return true
  }
}

We also need to include the ThrottlerModule from @nestjs/throttler in the module that will use it - in our case the chat module.

We can pass a default limit and ttl to it’s forRoot method. This means any gateway handler that implements our guard will by default allow 5 events in a 60 second time span.

import { Module } from '@nestjs/common'
import { ChatGateway } from './chat.gateway'
import { RoomModule } from '../room/room.module'
import { UserModule } from '../user/user.module'
import { CaslModule } from '../casl/casl.module'
import { ThrottlerModule } from '@nestjs/throttler'

@Module({
  imports: [RoomModule, UserModule, CaslModule, ThrottlerModule.forRoot({ limit: 5, ttl: 60 })],
  providers: [ChatGateway],
})
export class ChatModule {}

Throttler Guard Usage in Websocket Gateway

The only thing we now need to do is to use our guards in our chat websocket gateway via the @UseGuards decorator. If you’ve landed on this part of the series and want context for everything happening in the gateway, you can go back an skim through previous parts - but you don’t necessarily need to.

src/server/chat/chat.gateway.ts

import {
  MessageBody,
  SubscribeMessage,
  WebSocketGateway,
  WebSocketServer,
  OnGatewayConnection,
  OnGatewayDisconnect,
} from '@nestjs/websockets'
import { Logger, UseGuards, UsePipes } from '@nestjs/common'
import {
  ServerToClientEvents,
  ClientToServerEvents,
  Message,
  JoinRoom,
  KickUser,
} from '../../shared/interfaces/chat.interface'
import { Server, Socket } from 'socket.io'
import { RoomService } from '../room/room.service'
import { ZodValidationPipe } from '../pipes/zod.pipe'
import { ChatMessageSchema, JoinRoomSchema, KickUserSchema } from '../../shared/schemas/chat.schema'
import { UserService } from '../user/user.service'
import { ChatPoliciesGuard } from './guards/chat.guard'
import { WsThrottlerGuard } from './guards/throttler.guard'
import { Throttle } from '@nestjs/throttler'

@WebSocketGateway({})
export class ChatGateway implements OnGatewayConnection, OnGatewayDisconnect {
  constructor(private roomService: RoomService, private userService: UserService) {}

  @WebSocketServer() server: Server = new Server<ServerToClientEvents, ClientToServerEvents>()

  private logger = new Logger('ChatGateway')

  @Throttle(10, 30)
  @UseGuards(ChatPoliciesGuard<Message>, WsThrottlerGuard)
  @UsePipes(new ZodValidationPipe(ChatMessageSchema))
  @SubscribeMessage('chat')
  async handleChatEvent(
    @MessageBody()
    payload: Message
  ): Promise<boolean> {
    this.logger.log(payload)
    this.server.to(payload.roomName).emit('chat', payload)
    return true
  }

  @UseGuards(ChatPoliciesGuard<JoinRoom>, WsThrottlerGuard)
  @UsePipes(new ZodValidationPipe(JoinRoomSchema))
  @SubscribeMessage('join_room')
  async handleSetClientDataEvent(
    @MessageBody()
    payload: JoinRoom
  ): Promise<boolean> {
    this.logger.log(`${payload.user.socketId} is joining ${payload.roomName}`)
    await this.userService.addUser(payload.user)
    await this.server.in(payload.user.socketId).socketsJoin(payload.roomName)
    await this.roomService.addUserToRoom(payload.roomName, payload.user.userId)
    return true
  }

  @UseGuards(ChatPoliciesGuard<KickUser>)
  @UsePipes(new ZodValidationPipe(KickUserSchema))
  @SubscribeMessage('kick_user')
  async handleKickUserEvent(@MessageBody() payload: KickUser): Promise<boolean> {
    this.logger.log(`${payload.userToKick.userName} is getting kicked from ${payload.roomName}`)
    await this.server.to(payload.roomName).emit('kick_user', payload)
    await this.server.in(payload.userToKick.socketId).socketsLeave(payload.roomName)
    await this.server.to(payload.roomName).emit('chat', {
      user: {
        userId: 'serverId',
        userName: 'TheServer',
        socketId: 'ServerSocketId',
      },
      timeSent: new Date(Date.now()).toLocaleString('en-US'),
      message: `${payload.userToKick.userName} was kicked.`,
      roomName: payload.roomName,
    })
    return true
  }

  async handleConnection(socket: Socket): Promise<void> {
    this.logger.log(`Socket connected: ${socket.id}`)
  }

  async handleDisconnect(socket: Socket): Promise<void> {
    const user = await this.roomService.getFirstInstanceOfUser(socket.id)
    if (user !== 'Not Exists') {
      await this.userService.removeUserById(user.userId)
    }
    await this.roomService.removeUserFromAllRooms(socket.id)
    this.logger.log(`Socket disconnected: ${socket.id}`)
  }
}

Let’s drill into the important parts.

We have a handler that is subscribe to 'chat' events that handles messages. The default limit of 5 events per 60 seconds that we passed in the ThrottlerModule's forRoot method isn’t exactly what we want for chat messages. We want to increase this limit - which we can customize per handler by using the @Throttle(10, 30) decorator to increase the limit to 10 events per 30 seconds ttl.

@Throttle(10, 30)
@UseGuards(ChatPoliciesGuard<Message>, WsThrottlerGuard)
@UsePipes(new ZodValidationPipe(ChatMessageSchema))
@SubscribeMessage('chat')
async handleChatEvent(
  @MessageBody()
  payload: Message,
): Promise<boolean> {
  //...
}

The second handler we want to throttle is the one subscribed to 'join_room' events. Since we didn’t pass @Throttle, this will default to 5 actions per 60 seconds.

@UseGuards(ChatPoliciesGuard<JoinRoom>, WsThrottlerGuard)
@UsePipes(new ZodValidationPipe(JoinRoomSchema))
@SubscribeMessage('join_room')
async handleSetClientDataEvent(
  @MessageBody()
  payload: JoinRoom,
): Promise<boolean> {
  /...
}

In the event that one of our throttler guards detect a rate limit hit - an exception is thrown! The request will not go any further.

[1] [Nest] 53620  - 12/27/2022, 11:09:15 AM   ERROR [WsExceptionsHandler] ThrottlerException: Too Many Requests
[1] ThrottlerException: ThrottlerException: Too Many Requests
[1]     at WsThrottlerGuard.<anonymous> (/Users/austinhoward/code/nest-realtime/nest-react-websockets/src/server/chat/guards/throttler.guard.ts:17:13)
[1]     at Generator.next (<anonymous>)
[1]     at fulfilled (/Users/austinhoward/code/nest-realtime/nest-react-websockets/dist/server/server/chat/guards/throttler.guard.js:11:58)
[1]     at processTicksAndRejections (internal/process/task_queues.js:93:5)

Beyond brute force attacks, imagine we also made database calls or requests to external services in our gateway handler functions. We are now also protecting upstream services.

Acknowledgements

Now that we’ve implemented rate limiting, and also from previous parts of the series we’ve added data validation pipes and authorization guards - we need to start indicating back to the client whether or not their event was successful.

An important part of a good user experience is letting the user know what is happening - especially if there are failure cases!

For web socket based applications this is where acknowledgements come in handy.

Schemas

The first thing we need to do is update our zod schemas. Our web socket gateway event handlers should now accept an additional argument which is a callback function.

Here is the entire file with all of our schemas.

src/shared/schemas/chat.schema.ts

import { z } from 'zod'

export const UserIdSchema = z.string().min(1).max(24)

export const UserNameSchema = z
  .string()
  .min(1, { message: 'Must be at least 1 character.' })
  .max(16, { message: 'Must be at most 16 characters.' })

export const MessageSchema = z
  .string()
  .min(1, { message: 'Must be at least 1 character.' })
  .max(1000, { message: 'Must be at most 1000 characters.' })

export const TimeSentSchema = z.number()

export const RoomNameSchemaRegex = new RegExp('^\\S+\\w$')

export const RoomNameSchema = z
  .string()
  .min(2, { message: 'Must be at least 2 characters.' })
  .max(16, { message: 'Must be at most 16 characters.' })
  .regex(RoomNameSchemaRegex, {
    message: 'Must not contain spaces or special characters.',
  })

export const EventNameSchema = z.enum(['chat', 'kick_user', 'join_room'])

export const SocketIdSchema = z.string().length(20, { message: 'Must be 20 characters.' })

export const UserSchema = z.object({
  userId: UserIdSchema,
  userName: UserNameSchema,
  socketId: SocketIdSchema,
})

export const ChatMessageSchema = z.object({
  user: UserSchema,
  timeSent: TimeSentSchema,
  message: MessageSchema,
  roomName: RoomNameSchema,
  eventName: EventNameSchema,
})

export const RoomSchema = z.object({
  name: RoomNameSchema,
  host: UserSchema,
  users: UserSchema.array(),
})

export const JoinRoomSchema = z.object({
  user: UserSchema,
  roomName: RoomNameSchema,
  eventName: EventNameSchema,
})

export const KickUserSchema = z.object({
  user: UserSchema,
  userToKick: UserSchema,
  roomName: RoomNameSchema,
  eventName: EventNameSchema,
})

export const KickUserEventAckSchema = z.function().args(z.boolean()).returns(z.void())

export const KickUserEventSchema = z
  .function()
  .args(KickUserSchema, KickUserEventAckSchema)
  .returns(z.void())

export const ChatEventAckSchema = z.function().args(z.boolean()).returns(z.void())

export const ChatEventSchema = z
  .function()
  .args(ChatMessageSchema, ChatEventAckSchema)
  .returns(z.void())

export const JoinRoomEventAckSchema = z.function().args(z.string(), z.boolean()).returns(z.void())

export const JoinRoomEventSchema = z
  .function()
  .args(JoinRoomSchema, JoinRoomEventAckSchema)
  .returns(z.void())

export const ClientToServerEventsSchema = z.object({
  chat: ChatEventSchema,
  join_room: JoinRoomEventSchema,
  kick_user: KickUserEventSchema,
})

export const ServerToClientEventsSchema = z.object({
  chat: ChatEventSchema,
  kick_user: KickUserEventSchema,
})

Let’s focus in on the new acknowledgement callbacks.

With zod we can define schemas for everything including functions. Our event handler handles kicking users. We simply want the server to return a boolean indicating whether or not the handler successfully executed and did in fact kick the user.

// Kick user event callback function schema
export const KickUserEventAckSchema = z.function().args(z.boolean()).returns(z.void())

// Kick user event handler schema
export const KickUserEventSchema = z
  .function()
  .args(KickUserSchema, KickUserEventAckSchema)
  .returns(z.void())

Our chat event for handling chat messages we do the exact same schema for it’s callback.

// Chat event callback function schema
export const ChatEventAckSchema = z.function().args(z.boolean()).returns(z.void())

// Chat event handler schema
export const ChatEventSchema = z
  .function()
  .args(ChatMessageSchema, ChatEventAckSchema)
  .returns(z.void())

For the join room event, we’re going to do something similar but we’re going to use an additional argument in the callback function for a error message. The first argument in the JoinRoomEventAckSchema will be the error message (z.string) and the second one is the boolean returned from the handler function on the server just like the previous two events.

export const JoinRoomEventAckSchema = z.function().args(z.string(), z.boolean()).returns(z.void())

export const JoinRoomEventSchema = z
  .function()
  .args(JoinRoomSchema, JoinRoomEventAckSchema)
  .returns(z.void())

Acknowledgements Usage in Websocket Gateway

In the previous section where we added the ThrottlerGuard for rate limiting we saw the updated gateway. To support the new callback schemas we’re defined for our events, we simply need to return true from our handler functions.

src/server/chat/chat.gateway.ts

  @Throttle(10, 30)
  @UseGuards(ChatPoliciesGuard<Message>, WsThrottlerGuard)
  @UsePipes(new ZodValidationPipe(ChatMessageSchema))
  @SubscribeMessage('chat')
  async handleChatEvent(
    @MessageBody()
    payload: Message,
  ): Promise<boolean> {
    this.logger.log(payload);
    this.server.to(payload.roomName).emit('chat', payload);
    return true; // <---- return true when the handler function finishes
  }

  @UseGuards(ChatPoliciesGuard<JoinRoom>, WsThrottlerGuard)
  @UsePipes(new ZodValidationPipe(JoinRoomSchema))
  @SubscribeMessage('join_room')
  async handleSetClientDataEvent(
    @MessageBody()
    payload: JoinRoom,
  ): Promise<boolean> {
    this.logger.log(`${payload.user.socketId} is joining ${payload.roomName}`);
    await this.userService.addUser(payload.user);
    await this.server.in(payload.user.socketId).socketsJoin(payload.roomName);
    await this.roomService.addUserToRoom(payload.roomName, payload.user.userId);
    return true; // <---- return true when the handler function finishes
  }

  @UseGuards(ChatPoliciesGuard<KickUser>)
  @UsePipes(new ZodValidationPipe(KickUserSchema))
  @SubscribeMessage('kick_user')
  async handleKickUserEvent(
    @MessageBody() payload: KickUser,
  ): Promise<boolean> {
    this.logger.log(
      `${payload.userToKick.userName} is getting kicked from ${payload.roomName}`,
    );
    await this.server.to(payload.roomName).emit('kick_user', payload);
    await this.server
      .in(payload.userToKick.socketId)
      .socketsLeave(payload.roomName);
    await this.server.to(payload.roomName).emit('chat', {
      user: {
        userId: 'serverId',
        userName: 'TheServer',
        socketId: 'ServerSocketId',
      },
      timeSent: new Date(Date.now()).toLocaleString('en-US'),
      message: `${payload.userToKick.userName} was kicked.`,
      roomName: payload.roomName,
    });
    return true; // <---- return true when the handler function finishes
  }

Acknowledgements Usage in The Client

Now that we’ve added returns to our chat gateway handlers, we can now react to those on the client in a user friendly way. If you’ve been following along through the series, you’ll be familiar with how the client works up to this point. Here is our chat page that houses most of the application’s functionality.

src/client/pages/chat.tsx

import React, { useState, useEffect } from 'react'
import { MakeGenerics, useMatch, useNavigate } from '@tanstack/react-location'
import { io, Socket } from 'socket.io-client'
import {
  User,
  Message,
  ServerToClientEvents,
  ClientToServerEvents,
  KickUser,
  JoinRoom,
} from '../../shared/interfaces/chat.interface'
import { Header } from '../components/header'
import { UserList } from '../components/list'
import { MessageForm } from '../components/message.form'
import { Messages, ClientMessage } from '../components/messages'
import { ChatLayout } from '../layouts/chat.layout'
import { unsetRoom, useRoomQuery } from '../lib/room'
import { getUser } from '../lib/user'
import { ChatMessageSchema, JoinRoomSchema, KickUserSchema } from '../../shared/schemas/chat.schema'
import { LoadingLayout } from '../layouts/loading.layout'
import { Loading } from '../components/loading'

const socket: Socket<ServerToClientEvents, ClientToServerEvents> = io({
  autoConnect: false,
})

function Chat() {
  const {
    data: { user, roomName },
  } = useMatch<ChatLocationGenerics>()

  const [isConnected, setIsConnected] = useState(socket.connected)
  const [isJoinedRoom, setIsJoinedRoom] = useState(false)
  const [isJoiningDelay, setIsJoiningDelay] = useState(false)
  const [messages, setMessages] = useState<ClientMessage[]>([])
  const [toggleUserList, setToggleUserList] = useState<boolean>(false)

  const { data: room, refetch: roomRefetch } = useRoomQuery(roomName, isConnected)

  const navigate = useNavigate()

  useEffect(() => {
    if (!user || !roomName) {
      navigate({ to: '/', replace: true })
    } else {
      socket.on('connect', () => {
        setIsJoiningDelay(true)
        const joinRoom: JoinRoom = {
          roomName,
          user: { socketId: socket.id, ...user },
          eventName: 'join_room',
        }
        JoinRoomSchema.parse(joinRoom)
        setTimeout(() => {
          // default required 800 ms minimum join delay to prevent flickering
          setIsJoiningDelay(false)
        }, 800)
        socket.timeout(30000).emit('join_room', joinRoom, (err, response) => {
          if (err) {
            leaveRoom()
          }
          if (response) {
            setIsJoinedRoom(true)
          }
        })
        setIsConnected(true)
      })

      socket.on('disconnect', () => {
        setIsConnected(false)
      })

      socket.on('chat', (e) => {
        if (e.user.userId !== user.userId) {
          setMessages((messages) => [{ ...e, delivered: true }, ...messages])
        }
      })

      socket.on('kick_user', (e) => {
        if (e.userToKick.socketId === socket.id) {
          leaveRoom()
        }
      })

      socket.connect()
    }
    return () => {
      socket.off('connect')
      socket.off('disconnect')
      socket.off('chat')
      socket.off('kick_user')
    }
  }, [])

  const leaveRoom = () => {
    socket.disconnect()
    unsetRoom()
    navigate({ to: '/', replace: true })
  }

  const sendMessage = (message: string) => {
    if (user && socket && roomName) {
      const chatMessage: Message = {
        user: {
          userId: user.userId,
          userName: user.userName,
          socketId: socket.id,
        },
        timeSent: Date.now(),
        message,
        roomName: roomName,
        eventName: 'chat',
      }
      ChatMessageSchema.parse(chatMessage)
      setMessages((messages) => [{ ...chatMessage, delivered: false }, ...messages])
      socket.emit('chat', chatMessage, (response) => {
        if (response) {
          setMessages((messages) => {
            const previousMessageIndex = messages.findIndex((mes) => {
              if (mes.user.userId === user.userId && mes.timeSent === chatMessage.timeSent) {
                return mes
              }
            })
            if (previousMessageIndex === -1) {
              throw 'Previously sent message not found to update delivered status'
            }
            messages[previousMessageIndex] = {
              ...messages[previousMessageIndex],
              delivered: true,
            }
            return [...messages]
          })
        }
      })
    }
  }

  const kickUser = (userToKick: User) => {
    if (!room) {
      throw 'No room'
    }
    if (!user) {
      throw 'No current user'
    }
    const kickUserData: KickUser = {
      user: { ...user, socketId: socket.id },
      userToKick: userToKick,
      roomName: room.name,
      eventName: 'kick_user',
    }
    KickUserSchema.parse(kickUserData)
    socket.emit('kick_user', kickUserData, (response) => {
      if (response) {
        roomRefetch()
      }
    })
  }

  return (
    <>
      {user?.userId && roomName && room && isJoinedRoom && !isJoiningDelay ? (
        <ChatLayout>
          <Header
            isConnected={isConnected}
            users={room?.users ?? []}
            roomName={roomName}
            handleUsersClick={() => setToggleUserList((toggleUserList) => !toggleUserList)}
            handleLeaveRoom={() => leaveRoom()}
          ></Header>
          {toggleUserList ? (
            <UserList
              room={room}
              currentUser={{ socketId: socket.id, ...user }}
              kickHandler={kickUser}
            ></UserList>
          ) : (
            <Messages user={user} messages={messages}></Messages>
          )}
          <MessageForm sendMessage={sendMessage}></MessageForm>
        </ChatLayout>
      ) : (
        <LoadingLayout>
          <Loading message={`Joining ${roomName}`}></Loading>
        </LoadingLayout>
      )}
    </>
  )
}

export const loader = async () => {
  const user = getUser()
  return {
    user: user,
    roomName: sessionStorage.getItem('room'),
  }
}

type ChatLocationGenerics = MakeGenerics<{
  LoaderData: {
    user: Pick<User, 'userId' | 'userName'>
    roomName: string
  }
}>

export default Chat

When a client socket connects and is attempting to join a room, we want to improve the user experience by showing a loading screen until the server returns an acknowledgement from the 'join_room' event. To avoid flickering due to a fast connection, we also want to add a minimum delay of 800 milliseconds. To support this we need two useState hooks for when the user has successfully joined a room, and when the minimum 800 millisecond delay is finished.

const [isJoinedRoom, setIsJoinedRoom] = useState(false)
const [isJoiningDelay, setIsJoiningDelay] = useState(false)

We have our socket.on('connect' ...) listener on the client where we send the 'join_room' event to the server to join the client to a room. We fire off a useTimeout to do our minimum 800 millisecond delay that toggles isJoining to false.

As of socket.io version 4.4.0 there is now built in support for acknowledgement timeouts. We pass a timeout argument of 30000 milliseconds (30 seconds) which will throw an error if our server does not send back an acknowledgement within that time period. If the timeout it hit, the leaveRoom() function will disconnect the socket and redirect the user back to the login page where they can try to join again.

socket.on('connect', () => {
  setIsJoiningDelay(true)
  const joinRoom: JoinRoom = {
    roomName,
    user: { socketId: socket.id, ...user },
    eventName: 'join_room',
  }
  JoinRoomSchema.parse(joinRoom)
  setTimeout(() => {
    // default required 800 ms minimum join delay to prevent flickering
    setIsJoiningDelay(false)
  }, 800)
  socket.timeout(30000).emit('join_room', joinRoom, (err, response) => {
    if (err) {
      leaveRoom()
    }
    if (response) {
      setIsJoinedRoom(true)
    }
  })
  setIsConnected(true)
})

In the return of our chat page where we render UI elements, we have a top-level check if isJoinedRoom and !isJoiningDelay in which case we render the chat room components, and if not we reader a loading screen on the page.

return (
  <>
    {user?.userId && roomName && room && isJoinedRoom && !isJoiningDelay ? (
      <ChatLayout>
        <Header
          isConnected={isConnected}
          users={room?.users ?? []}
          roomName={roomName}
          handleUsersClick={() => setToggleUserList((toggleUserList) => !toggleUserList)}
          handleLeaveRoom={() => leaveRoom()}
        ></Header>
        {toggleUserList ? (
          <UserList
            room={room}
            currentUser={{ socketId: socket.id, ...user }}
            kickHandler={kickUser}
          ></UserList>
        ) : (
          <Messages user={user} messages={messages}></Messages>
        )}
        <MessageForm sendMessage={sendMessage}></MessageForm>
      </ChatLayout>
    ) : (
      <LoadingLayout>
        <Loading message={`Joining ${roomName}`}></Loading>
      </LoadingLayout>
    )}
  </>
)

Here is a snapshot of the loading screen.

So what about for our chat messages? How can we indicate to the user whether their messages are successfully being sent?

What we can do for messages is a delivered indicator alongside the messages in the UI.

To handle this, in our chat page we have a sendMessage() function that fires whenever a user sends a message.

Check if chat message passes schema check to ensure data is not malformed
Update state with the chat message but set it’s “delivered” status to false.
Emit the event and pass the acknowledgement callback function
When callback function is executed by the server, update the state of the previously pushed messages by setting the message’s “delivered” status to true.

const sendMessage = (message: string) => {
  if (user && socket && roomName) {
    const chatMessage: Message = {
      user: {
        userId: user.userId,
        userName: user.userName,
        socketId: socket.id,
      },
      timeSent: Date.now(),
      message,
      roomName: roomName,
      eventName: 'chat',
    }
    // Check if chat message passes schema check
    ChatMessageSchema.parse(chatMessage)
    // Update state with message "delivered" status to false
    setMessages((messages) => [{ ...chatMessage, delivered: false }, ...messages])
    // Emit 'chat' event with message and callback
    socket.emit('chat', chatMessage, (response) => {
      // If server response with response === true
      if (response) {
        // Update state by finding previously set message
        // and setting it's "delivered" status to true
        setMessages((messages) => {
          const previousMessageIndex = messages.findIndex((mes) => {
            if (mes.user.userId === user.userId && mes.timeSent === chatMessage.timeSent) {
              return mes
            }
          })
          if (previousMessageIndex === -1) {
            throw 'Previously sent message not found to update delivered status'
          }
          messages[previousMessageIndex] = {
            ...messages[previousMessageIndex],
            delivered: true,
          }
          return [...messages]
        })
      }
    })
  }
}

Now we can update our Messages component to display the “delivered” status of the messages.

We only want to display the status of the client’s own messages so we check if the client’s userId matches the message’s userId.

src/client/components/messages.tsx

import React from 'react'
import { Message, User } from '../../shared/interfaces/chat.interface'
export type ClientMessage = Message & { delivered: boolean }

const determineMessageStyle = (
  user: Pick<User, 'userId' | 'userName'>,
  messageUserId: Message['user']['userId']
) => {
  if (user && messageUserId === user.userId) {
    return {
      message: 'bg-slate-500 p-4 ml-24 rounded break-words',
      sender: 'ml-24 pl-4',
    }
  } else {
    return {
      message: 'bg-slate-800 p-4 mr-24 rounded break-words',
      sender: 'mr-24 pl-4',
    }
  }
}

export const Messages = ({
  user,
  messages,
}: {
  user: Pick<User, 'userId' | 'userName'>
  messages: ClientMessage[]
}) => {
  return (
    <div className="flex h-4/6 w-full flex-col-reverse overflow-y-scroll">
      {messages?.map((message, index) => {
        return (
          <div key={index + message.timeSent} className="mb-4">
            <div className={determineMessageStyle(user, message.user.userId).sender}>
              <span className="text-sm text-gray-400">{message.user.userName}</span>
              <span className="text-sm text-gray-400">{' ' + '•' + ' '}</span>
              <span className="text-sm text-gray-400">
                {new Date(message.timeSent).toLocaleString()}
              </span>
            </div>
            <div className={determineMessageStyle(user, message.user.userId).message}>
              <p className="text-white">{message.message}</p>
            </div>
            {user && message.user.userId === user.userId && (
              <p className="text-right text-xs text-gray-400">
                {message.delivered ? 'Delivered' : 'Not delivered'}
              </p>
            )}
          </div>
        )
      })}
    </div>
  )
}

Now if the user were to hit a rate limit on sending messages, they would know if the message was Not delivered or Delivered. Awesome!

Great! We added rate limiting to our websocket gateway, and introduced acknowledgements to send data back to the client to indicate what happened with their events. We could go much further with this but we’ll end it here for now!

Blog