NestJS WebSockets: Authentication, Authorization & Pipes

An example of how to use WebSockets with NestJS, and implement features like Authentication, Authorization & Pipes.

I’ve been wanting to try out and build a NestJS API that is able to transmit real-time data to clients for a few months now. The specific situation I had in mind was related to a project I’m working on, where clients have their own resources, and only they have access to them. I wanted to give them an option to view some of those resources in real-time, as and when they are updated.

Since I’ve written a number of articles around NestJS already, that’s what I’m continuing here with as well. NestJS, WebSockets and techniques to auth and validate requests.

As always, this article comes bundled with a full implementation example, which you can find and use over on my GitHub page. I won’t be going through every single detail over here, just the main sections we need to cover to make this work.

TLDR WebSockets

The WebSocket protocol enables full-duplex interaction between a client and a server, with lower overhead than alternatives such as HTTP polling, facilitating real-time data transfer to and from the server.

Prerequisites

To get started, we need to install a few extra packages that allow us to implement WebSockets and the surounding features:

yarn add @nestjs/websockets @nestjs/platform-socket.io socket.io zod

Building the Gateway

This is where we’ll construct our WebSocket server, which is our main entrypoint. In Nest, a gateway is simply a class annotated with @WebSocketGateway() decorator.

We’ll build a generic gateway, and name it events. We start the process easily by running:

nest g gateway events

We’ll get a new gateway ready to go that looks like this:

@WebSocketGateway()
export class EventsGateway {
  @SubscribeMessage('message')
  handleMessage(client: any, payload: any): string {
    return 'Hello world!';
  }
}

At this stage, you can use a tool like Postman to create a Socket.IO connection to your API (e.g. localhost:3000), and send anything to the message event name, as well as listen back to the same event name, and see if you get a response.

Authentication

Now that we have a simple server running, we’re going to immediately want to configure authorization, before anything else gets more complicated. We’ll start off with a new auth module:

nest g module auth

We’re going to want to create a Guard and a Socket.IO middleware. Normally, the guard should suffice, but they don’t work exactly like we’d need them to for WebSockets in NestJS, and as such, we’re doing both, to ensure we’re protected.

Auth Guard

For the sake of simplicity, in this article (and GitHub example) I’m not going to build something very extensive or even that safe for that matter. It’ll be up to you to add something more robust like JWT.

We can construct a new file in our auth module, call it socket.guard.ts. This is the code we need for it:

// Hardcoded user tokens for demonstration purposes
const VALID_USER_TOKENS = new Map<string, { id: string }>([
  ['valid-token-1', { id: 'user_1' }],
  ['valid-token-2', { id: 'user_2' }],
  ['valid-token-3', { id: 'user_3' }],
]);

export class SocketGuard implements CanActivate {
  canActivate(context: ExecutionContext): boolean | Promise<boolean> | Observable<boolean> {
    if (context.getType() !== 'ws') {
      return false;
    }

    const socket = context.switchToWs().getClient<Socket>();
    if (!SocketGuard.validateToken(socket)) {
      throw new WsException('Unauthorized');
    }

    return true;
  }

  static validateToken(socket: Socket): boolean {
    const headerToken = socket.handshake.headers.authorization;
    const authToken = socket.handshake.auth.token;

    const user = VALID_USER_TOKENS.get(headerToken || '') || VALID_USER_TOKENS.get(authToken);

    if (user !== undefined) {
      socket.data.userId = user.id;
      return true;
    }

    return false;
  }
}

Our main point of interest here is the validateToken function, where we take in data either from the request headers or the auth object of the socket connection. We then check that data (tokens) against a source of truth, and if they’re good, we return a positive response. Otherwise… bad luck, try again.

If it’s a success, we also add the user ID in the socket data, to be used later down the line.

Socket.IO Middleware

This middleware is what is actually going to be doing the auth work for us (as mentioned previously, Guards are unreliable here). The implementation is really simple, and we’re going to make use of the same validateToken function from above.

This is what the middleware needs to look like:

type SocketIoMiddleware = {
  (socket: Socket, next: (err?: any) => void): void;
};

export const SocketAuthMiddleware: SocketIoMiddleware = (socket, next) => {
  if (!SocketGuard.validateToken(socket)) {
    return next(new WsException('Unauthorized'));
  }
  return next();
};

Adding Auth to the Gateway

Now that we have the Guard and Middleware ready, we’re going to add them to our Gateway. Here’s what it needs to look like now:

@WebSocketGateway()
@UseGuards(SocketGuard)
export class EventsGateway implements OnGatewayInit<Server> {
  afterInit(server: Server): void {
    server.use(SocketAuthMiddleware);
  }

...
}

For the Guard, we make use of the UseGuards() decorator, and feed it the guard we built.

For the Middleware, we make sure that our EventsGateway class implements the OnGatewayInit interface, of type Server coming from Socket.IO. This forces us to implement the afterInit function, where we’ll tell our WebSocket server to make use of our middleware.

---

And that’s it for auth. If you now try to connect again to your sever using Postman, you’ll see that you get an immediate error response, telling the client they are not authenticated. To make it work again, Authorization header with a correct token in the value, and see it working like a charm. (Postman only allows for headers to be added, not actual WebSocket auth tokens, as of writing this)

Authorization

To showcase how this will work, we’re going to need an actual “subscribed” function where we can handle it. Keeping in tone with what I initially mentioned about my work, clients having their own resources, I’m going to create a function/event through which they can subscribe to see updates on a “project”. Again, to keep this example simple, we’re hardcoding some things, similarly to the authentication section.

Here’s the function (added in the gateway class), as well as the hardcoded permissions:

// Hardcoded user/project permissions for demonstration purposes
const USER_PROJECT_PERMISSIONS = new Map<string, string[]>([
  ['user_1', ['project_one']],
  ['user_2', ['project_one', 'project_two', 'project_three']],
  ['user_3', ['project_three']],
]);

@SubscribeMessage('join_project')
handleJoinProject(
  @MessageBody(new ZodValidationPipe(JoinProjectSchema)) payload: JoinProject,
  @ConnectedSocket() client: Socket,
): WsResponse<string> {
  // Get the user ID from the client data
  const userId: string | undefined = client.data.userId;
  if (userId === undefined) {
    throw new WsException('Unauthorized'); // This should never happen
  }

  // Check if the user has permissions to join the project
  if (!USER_PROJECT_PERMISSIONS.get(userId)?.some((id) => id === payload.projectId)) {
    throw new WsException('Project not found');
  }

  // Join the project room
  client.join(payload.projectId);

  return { event: 'join_project', data: `You have joined ${payload.projectId}!` };
}

A client can send a join_project message, with a payload we’ll discuss in the next topic, with the purpose of joining a room designated for that project.

Before that happens though, we do any checks we need to validate permissions. Depending on your implementation of the Authentication section, this part could be built in a smarter way, where you don’t need to go “outside” (e.g. a DB) to fetch permissions on each incoming message.

We make use of WsException to return valid errors that the client can interpret and handle/display.

Pipes

As you may have noticed in the code snippet above, we make use of Pipes to validate the incoming payload.

@MessageBody(new ZodValidationPipe(JoinProjectSchema)) payload: JoinProject

We’re able to do this with the help of zod. First step, is to create our custom pipe:

@Injectable()
export class ZodValidationPipe implements PipeTransform {
  constructor(private schema: Schema) {}

  transform(value: any) {
    try {
      this.schema.parse(value);
    } catch (error) {
      throw new WsException({
        message: 'Validation failed',
        errors: (error as ZodError).errors,
      });
    }
    return value;
  }
}

This is a reusable implementation that can be used in any message handler you need. Next up, we build our zod schema and TS type, to be used in the implementation:

const JoinProjectSchema = z.object({
  projectId: z.string(),
});

type JoinProject = z.infer<typeof JoinProjectSchema>;

With all these put together, if the client tries to send an invalid message payload to our server/handler, they will be met with an exception/error.

Wrapping up

To finish out this example, we’re going to need the server to actually send some updates to our clients, which have joined some project rooms. Keeping in tone with simplicity, we add a function in our gateway that sends a message to all our defined projects, with the exact date & time, on the chat message.

sendGlobalMessages(): void {
    this.server.to('project_one').emit('chat', `Hello, project_one! Time is -- ${new Date().toISOString()} --`);
    this.server.to('project_two').emit('chat', `Hello, project_two! Time is -- ${new Date().toISOString()} --`);
    this.server.to('project_three').emit('chat', `Hello, project_three! Time is -- ${new Date().toISOString()} --`);
  }

Similarly, you can build functions that are able to send specific pieces of data/types to the clients, when new data is available.

---

I hope you found this guide and example useful for your needs, and that you’re left with some valuable information on how to tackle WebSockets with NestJS now. Don’t forget to also check out the full code example on my GitHub page. Happy coding!