NestJS-Web Socket(deep)

🖊️Validation pipe

이번에는 Pipe를 사용해보도록 하겠습니다.

REST API에서의 pipe와 Gateway에서의 API는 서로 다르지 않습니다. 공식문에서도 별로 다르지 않다고 합니다.

포스트맨으로 1번 사용자를 연결 후 JSON으로 아무런 값을 주지않고 create_chat을 해보도록 하겠습니다.

먼저 리스닝으로 exception과 receive_message를 활성화 해줍니다.

이렇게 나오면 안됩니다. 왜냐하면 어떠한 값이 뭐가 잘못 되었다고 나와야합니다. 그런데 InternalServerError가 나오게 됩니다.

우리는 이미 main.ts에서 Validation pipe를 전역적으로 작용하도록 만들었습니다. 하지만 작동을 하지 않습니다.

• main.ts

app.useGlobalPipes(new ValidationPipe({
    transform: true, // 변화는 해도 된다.
    transformOptions: {
      	enableImplicitConversion: true // 임의로 변환하는 것을 허가한다.
    },
    whitelist: true,
    forbidNonWhitelisted: true,
}));

아쉽게도 글로벌 파이프를 전역적으로 작용하는 것은 오직 REST API 컨트롤러에만 적용이 됩니다. 그래서 우리가 Gateway를 사용할 때는 따로 Validation을 gateway에 추가해줘야 합니다.

따라서 @SubscribeMessage별로 @UsePipes를 적용해야합니다.

• chats.gateway.ts

@UsePipes(new ValidationPipe({
    transform: true, // 변화는 해도 된다.
    transformOptions: {
      	enableImplicitConversion: true // 임의로 변환하는 것을 허가한다.
    },
    whitelist: true,
    forbidNonWhitelisted: true,
}))
@SubscribeMessage('create_chat')
async createChat(
    @MessageBody() data: CreateChatDto,
    @ConnectedSocket() socket: Socket,
) {
    const chat = await this.chatsService.createChat(
      	data, 
    );
}

포스트맨으로 다시 보내도 에러가 발생합니다.

현재 에러의 위치를 보면 validation.pipe로 나옵니다. 즉, dto에서는 통과가 되었다는 것입니다. 하지만 Bad Request Exception이 메세지로 전달이 되지않고 터져버렸습니다.

왜냐하면 우리는 에러를 던질 때 WsException으로 던져야합니다. 하지만 class-validator는 기본적으로 REST API를 위해서 설계가 된것이기 때문에 모든 Expection들이 HTTP Excetpion을 extends하고 있습니다.

하지만 WsException은 HTTP Exception을 extends하고 있지 않습니다. 그래서 다음과 같은 에러가 발생한 것입니다.

그러면 우리는 HTTP Exception들이 발생했을 때, WsException으로 변환해주는 코드를 작성하면 되는 것입니다.

---

🖊️Exception Filter 적용

이어서 HTTP Exception들이 발생했을 때, WsException으로 변환해주는 코드를 작성하겠습니다.

Exception Filter에서 HTTP Exception을 모두 잡아서 WsException으로 바꿔주면 에러들을 전부 잡을 수 있습니다.

common 폴더에서 작업을 진행하겠습니다.

• common/exception-filter/socket-catch-http.exception-filter.ts

@Catch(HttpException)
export class SocketCatchHttpExceptionFilter extends BaseWsExceptionFilter<HttpException> {
    // BaseWsExceptionFilter를 상속하면 Ws관련 Exception을 만들 수 있다.

    catch(exception: HttpException, host: ArgumentsHost): void {
        super.catch(exception, host);

        const socket = host.switchToWs().getClient(); // 소켓 가져오기
        socket.emit( // 현재 소켓에다가만 emit하기
            'exception', // 이벤트 이름
            {	
              	// 실제 응답에서 받는 {메세지 형태들}
              	data: exception.getResponse(),
            }
        )
    }
}

Exception Filter를 적용하겠습니다.

• chats.gateway.ts

@UsePipes(new ValidationPipe({
    transform: true,
    transformOptions: {
      	enableImplicitConversion: true
    },
    whitelist: true,
    forbidNonWhitelisted: true,
}))
@UseFilters(SocketCatchHttpExceptionFilter) // 적용
@SubscribeMessage('create_chat')
async createChat(
    @MessageBody() data: CreateChatDto,
    @ConnectedSocket() socket: Socket,
) {
    const chat = await this.chatsService.createChat(
      	data, 
    );
}

포스트맨으로 테스트를 하겠습니다.

User 1과 User 2모두 리스너를 exception, receive_message를 열겠습니다. 그리고 create_chat을 실행하면 에러가 2개가 나옵니다.

이유는 super.catch() 때문입니다. 따라서 해당 코드를 삭제합니다. 이후에 다시 똑같은 방법으로 테스트를 하게되면

다음 에러 메세지는 User 1번에게만 가게 됩니다. User 2번은 어떠한 에러 메세지를 받지 못합니다. 왜냐하면 연결된 사용자한테만 가는 emit 때문입니다.

즉, Gateway에서는 Pipe를 적용하기 위해서는 각각의 메소드 위에다가 적용을 해줘야합니다.

---

🖊️Guard 적용

이번에는 Guard를 적용해보겠습니다. REST API에서 적용한 방법과 동일합니다.

• auth/guard/socket/socket-bearer-token.guard.ts

@Injectable()
export class SocketBearerTokenGuard implements CanActivate {

    constructor(
    	private readonly authService: AuthService,
     	private readonly userService: UsersService,
    ) {}

    async canActivate(context: ExecutionContext): Promise<boolean> {
        // 지금 연결해서 사용하고 있는 소켓
        const socket = context.switchToWs().getClient();

        // 헤더 가져오기
        const headers = socket.handshake.headers;

        // Bearer xxx
        const rawToken = headers['authorization'];
        if (!rawToken) throw new WsException('토큰이 없습니다. ');

        const token = this.authService.extractTokenFromHeader(
            rawToken,
            true
        );

        const payload = this.authService.verifyToken(token);
        const user = await this.userService.getUserByEmail(payload.email);

        socket.user = user;
        socket.token = token;
        socket.tokeType = payload.tokenType;
    }
}

이제 try catch로 묶겠습니다. 왜냐하면 기본적으로 HTTP Exception을 주기 때문에 WsException으로 변경을 해줘야 합니다.

@Injectable()
export class SocketBearerTokenGuard implements CanActivate {

    constructor(
    	private readonly authService: AuthService,
     	private readonly userService: UsersService,
    ) {}

    async canActivate(context: ExecutionContext): Promise<boolean> {
        const socket = context.switchToWs().getClient();
        const headers = socket.handshake.headers;
        const rawToken = headers['authorization'];
        if (!rawToken) throw new WsException('토큰이 없습니다. ');

        try {
            const token = this.authService.extractTokenFromHeader(
                rawToken,
                true
            );

            const payload = this.authService.verifyToken(token);
            const user = await this.userService.getUserByEmail(payload.email);

            socket.user = user;
            socket.token = token;
            socket.tokeType = payload.tokenType;
            return true;
        } catch (error) {
          	throw new WsException('토큰이 유효하지 않습니다.')
        }
    }
}

chats.module.ts에서 provider로 등록을 하겠습니다.

• chats.module.ts

@Module({
    imports: [
        TypeOrmModule.forFeature([
          	ChatsModel,
          	MessagesModel,
        ]),
        CommonModule,
        AuthModule,
      	UsersModule,
    ],
    controllers: [
        ChatsController,
        MessagesController,
    ],
    providers: [
        ChatsGateway,
        ChatsService,
        ChatsMessagesService,
    ],
})
export class ChatsModule {}

이제 chats.gateway.ts에 적용을 하겠습니다.

• chats.gateway.ts

@UsePipes(new ValidationPipe({
    transform: true,
    transformOptions: {
      	enableImplicitConversion: true
    },
    whitelist: true,
    forbidNonWhitelisted: true,
}))
@UseFilters(SocketCatchHttpExceptionFilter)
@UseGuards(SocketBearerTokenGuard) // 추가
@SubscribeMessage('create_chat')
async createChat(
    @MessageBody() data: CreateChatDto,
  	// 인터섹션: user가 UsersModel이라고 존재한다.
  	// 토큰이 통과되면 user가 있다는 것을 알 수 있다.
    @ConnectedSocket() socket: Socket & {user: UsersModel}, 
) {
    const chat = await this.chatsService.createChat(
      	data, 
    );
}

포스트맨으로 테스트를 하겠습니다. 로그인을 하지 않고 create_chat을 누르면 예상한대로 나옵니다.

로그인을 하고 응답받은 accessToken을 Headers에 넣어주겠습니다.

따라서 REST API에서 적용한 것처럼 똑같이 Gateway에서도 Guard를 적용할 수 있는 것을 알 수 있습니다.

---

🖊️decorator 기반 로직 변경

나머지 기능을 완성하겠습니다. enter_chat에도 Guard, Filter, Pipe를 적용하겠습니다. 로그인을 해야만 접근이 가능하도록 만들겠습니다.

또한 send_message에도 동일하게 적용을 하겠습니다.

• chats.gateway.ts

@UsePipes(new ValidationPipe({
    transform: true,
    transformOptions: {
      	enableImplicitConversion: true
    },
    whitelist: true,
    forbidNonWhitelisted: true,
}))
@UseFilters(SocketCatchHttpExceptionFilter)
@UseGuards(SocketBearerTokenGuard)
@SubscribeMessage('send_message')
async sendMessage(
.
.
@UsePipes(new ValidationPipe({
    transform: true,
    transformOptions: {
    enableImplicitConversion: true
    },
    whitelist: true,
    forbidNonWhitelisted: true,
}))
@UseFilters(SocketCatchHttpExceptionFilter)
@UseGuards(SocketBearerTokenGuard)
@SubscribeMessage('enter_chat')
async enterChat(

추가적으로 CreateMessageDto에서 authorId를 number로 받고 있었습니다. 이 또한 제거를 하고 직접 받도록 하겠습니다.

• chats/messages/dto/create-message.dto.ts

import { PickType } from "@nestjs/mapped-types";
import { MessagesModel } from "../entities/messages.entity";
import { IsNumber } from "class-validator";

export class CreateMessagesDto extends PickType(MessagesModel, [
  	'message',
]) {
    @IsNumber()
    chatId: number;
  
  	// 제거
}

• messages.service.ts

async createMessage(
    dto: CreateMessagesDto,
    authorId: number
) {
    const message = await this.messagesRepository.save({
        chat: {
          	id: dto.chatId,
        },
        author: {
          	id: authorId, // 변경
        },
        message: dto.message,
    });
    return this.messagesRepository.findOne({
        where: {
          	id: message.id,
        },
        relations: {
          	chat: true,
        }
    });
}

• chats.gateway.ts

@UsePipes(new ValidationPipe({
    transform: true, // 변화는 해도 된다.
    transformOptions: {
      	enableImplicitConversion: true // 임의로 변환하는 것을 허가한다.
    },
    whitelist: true,
    forbidNonWhitelisted: true,
}))
@UseFilters(SocketCatchHttpExceptionFilter)
@UseGuards(SocketBearerTokenGuard)
@SubscribeMessage('send_message')
async sendMessage(
    @MessageBody() dto: CreateMessagesDto,
    @ConnectedSocket() socket: Socket & {user: UsersModel}, // 변경
) {
    const chatExists = await this.chatsService.checkIfChatExists(
      	dto.chatId,
    );

    if (!chatExists) {
        throw new WsException(
          	`존재하지 않는 채팅방입니다. Chat ID : ${dto.chatId}`,
        );
    }

    const message = await this.messagesService.createMessage(
        dto,
        socket.user.id
    );
    socket.to(message.chat.id.toString()).emit('receive_message', message.message);
}

다음과 같이 바꾸는 이유는 엑세스 토큰으로 부터 사용자 정보를 받아오기 위해서 입니다.

---

🖊️AccessToken을 매번 검증할때의 문제

지금부터는 accessToken을 적용하는 문제에 대해서 알아보겠습니다.

로그인을 하게 되면 accessToken과 refreshToken을 제공합니다. accessToken의 경우 5분입니다. 그 이후에는 만료가 되어서 다시 refreshToken으로 재발급을 받아야합니다.

포스트맨의 경우 connection을 하게되면 바꿀 수 없게 되어버립니다.

그럼 5분 안으로는 소통이 가능하지만, 그 이후 만료가 되면 토큰이 유효하지 않습니다. 에러를 던지게 됩니다. 그래서 매번 무엇인가를 검증을 한다면 Guard를 사용하는 것이 맞습니다.

하지만 Bearer Token 검증을 통해 사용자와의 socket과 연결에는 매번 Guard를 이용한 검증을 할 필요가 없습니다.

이제부터 사용자를 어떻게 각각의 소켓과 연결하는지를 알아보겠습니다.

---

🖊️Socket에서 사용자 정보 저장

먼저 chats.gateway.ts에 있는 Guard를 모두 지워줍니다.

• chats.gateway.ts

@WebSocketGateway({
    // ws:localhost:3000/chats
    namespace: '/chats',
})
export class ChatsGateway implements OnGatewayConnection{

    constructor(
    	private readonly chatsService: ChatsService,
     	private readonly messagesService: ChatsMessagesService
    ){}

    @WebSocketServer()
    server: Server;

    handleConnection(socket: Socket) {
      	console.log(`on connect called : ${socket.id}`);
    }

    @UsePipes(new ValidationPipe({
        transform: true,
        transformOptions: {
          	enableImplicitConversion: true
        },
        whitelist: true,
        forbidNonWhitelisted: true,
    }))
    @UseFilters(SocketCatchHttpExceptionFilter)
    @SubscribeMessage('enter_chat')
    async enterChat(
    	@MessageBody() data: EnterChatDto,
     	@ConnectedSocket() socket: Socket & {user: UsersModel},
    ) {  
        for (const chatId of data.chatIds) {
            const exists = await this.chatsService.checkIfChatExists(
              	chatId,
            );

            if (!exists) {
                throw new WsException({
                    code: 100,
                    message: `존재하지 않는 chat 입니다. chatId: ${chatId}`,
                });
            }
        }
        socket.join(data.chatIds.map((x) => x.toString()));
    }

    @UsePipes(new ValidationPipe({
        transform: true,
        transformOptions: {
          	enableImplicitConversion: true
        },
        whitelist: true,
        forbidNonWhitelisted: true,
    }))
    @UseFilters(SocketCatchHttpExceptionFilter)
    @SubscribeMessage('create_chat')
    async createChat(
    	@MessageBody() data: CreateChatDto,
     	@ConnectedSocket() socket: Socket & {user: UsersModel},
    ) {
        const chat = await this.chatsService.createChat(
          	data, 
        );
    }

    @UsePipes(new ValidationPipe({
        transform: true,
        transformOptions: {
          	enableImplicitConversion: true
        },
        whitelist: true,
        forbidNonWhitelisted: true,
    }))
    @UseFilters(SocketCatchHttpExceptionFilter)
    @SubscribeMessage('send_message')
    async sendMessage(
      	@MessageBody() dto: CreateMessagesDto,
       	@ConnectedSocket() socket: Socket & {user: UsersModel},
    ) {
        const chatExists = await this.chatsService.checkIfChatExists(
          	dto.chatId,
        );

        if (!chatExists) {
            throw new WsException(
             	 `존재하지 않는 채팅방입니다. Chat ID : ${dto.chatId}`,
            );
        }
        const message = await this.messagesService.createMessage(
          	dto,
          	socket.user.id
        );
        socket.to(message.chat.id.toString()).emit('receive_message', message.message);
    }
}

Guard를 지우고 어떻게 사용자 정보를 가지고 올 수 있을까? 

Socket이라는 것은 연결이 되면 서로 Pipe같은 것이 생깁니다. 그래서 한번 연결되어 있으면 그 연결은 지속이 됩니다.

async handleConnection(socket: Socket & {user: UsersModel}) {
    console.log(`on connect called : ${socket.id}`); // 어떤 id 소켓이 연결됨?
  
  	// 임시 테스트
    const user = await this.usersService.getUserByEmail('codefactory@codefactory.ai');
    socket.user = user; // 소켓에 사용자 정보 넣기
}

따라서 1번 연결되고 나면 socket에는 메세지를 보낼 때 계속해서 유지가 됩니다. 이것은 디버깅을 통해서 socket에 들어있는 값을 확인해보면 알 수 있습니다.

따라서 REST API처럼 모든 곳에 Guard를 붙일 필요 없이, handleConnection에서 핸들링만 해주고 socket.user = user;로 넣어주면, 나머지 메세지를 보낼 때 소켓에 들어있는 값을 신뢰할 수 있게 됩니다.

또한 여기서는 중간에 에러가 발생하면 에러를 throw 하는 것보다는 연결을 끊어주겠습니다.

async handleConnection(socket: Socket & {user: UsersModel}) {
    console.log(`on connect called : ${socket.id}`); // 어떤 id 소켓이 연결됨?
    const headers = socket.handshake.headers; // 헤더 가져오기
    const rawToken = headers['authorization']; // Bearer xxx
    if (!rawToken) socket.disconnect(); // 토큰이 없으면 연결 끊기

    try {
        const token = this.authService.extractTokenFromHeader(
            rawToken,
            true
        );

        const payload = this.authService.verifyToken(token);
        const user = await this.usersService.getUserByEmail(payload.email);
        socket.user = user;
        return true;
    } catch (error) {
      	socket.disconnect(); // 에러가 나면 연결 끊기
    }
}

포스트맨으로 로그인을 하고 토큰없이 connect를 하면 바로 연결이 끊기는 것을 알 수 있습니다.

---

🖊️Gateway Lifecycle Hooks

연결되었을 때 handleconnection을 통해서 특정 기능이 발생하도록 만들었습니다. 이것을 Lifecycle Hooks이라고 합니다. Lifecycle Hooks은 2개가 더 있습니다.

OnGatewayInit 입니다. OnGatewayInit는 afterInit 함수를 사용할 수 있게 만들어 줍니다. 이 함수는 실제 서버를 inject 받을 때 사용되는 함수이기도 하고, 그리고 Gateway가 초기화되었을 때 실행할 수 있는 함수 입니다.

• chats.gateway.ts

export class ChatsGateway implements OnGatewayConnection, OnGatewayInit{ // 추가

    constructor(
    	private readonly chatsService: ChatsService,
     	private readonly messagesService: ChatsMessagesService,
     	private readonly usersService: UsersService,
     	private readonly authService: AuthService,
    ){}

    @WebSocketServer()
    server: Server;

    afterInit(server: any) { // server: Server; 이것과 동일한 값을 받는다.
      	console.log(`after gateway init`); // 로그 확인
    }

after gateway init 이후로 소켓 메세지들이 구독된것을 알 수 있습니다. 그래서 Gateway가 시작했을 때 특정 함수 또는 로직을 실행하고 싶으면 afterInit이라는 Lifecycle Hooks을 사용하면 됩니다.

다음은 연결이 끊긴 이후를 관리하는 OnGatewayDisconnect입니다.

export class ChatsGateway implements OnGatewayConnection, OnGatewayInit, OnGatewayDisconnect{
.
.
handleDisconnect(socket: Socket) {
  	console.log(`on disconnect called : ${socket.id}`);
}

포스트맨으로 연결 후 disconnect를 하겠습니다.