@@ -10,9 +10,7 @@ import {
1010 MessageBody ,
1111} from '@nestjs/websockets' ;
1212import { Server , Socket } from 'socket.io' ;
13- import { Logger , Inject , UnauthorizedException } from '@nestjs/common' ;
14- import { AuthService } from '../auth/auth.service' ;
15- import { AuthUserPayload } from '../auth/types/auth-user.type' ;
13+ import { Logger } from '@nestjs/common' ;
1614
1715@WebSocketGateway ( {
1816 cors : {
@@ -27,73 +25,18 @@ export class NotificationsGateway implements OnGatewayConnection, OnGatewayDisco
2725 private logger : Logger = new Logger ( 'NotificationsGateway' ) ;
2826 private userSockets = new Map < string , string [ ] > ( ) ;
2927 private socketUsers = new Map < string , string > ( ) ;
30- private eventRateLimits = new Map < string , { count : number ; resetTime : number } > ( ) ;
31- private readonly MAX_CONNECTIONS_PER_USER = 5 ;
32- private readonly MAX_EVENTS_PER_MINUTE = 60 ;
33- private readonly RATE_LIMIT_WINDOW = 60 * 1000 ; // 1 minute in milliseconds
3428
35- constructor ( @Inject ( AuthService ) private readonly authService : AuthService ) { }
36-
37- private extractTokenFromHandshake ( client : Socket ) : string | null {
38- // Try to get token from authorization header
39- const authHeader = client . handshake . headers . authorization ;
40- if ( authHeader ) {
41- const [ scheme , token ] = authHeader . split ( ' ' ) ;
42- if ( scheme === 'Bearer' && token ) {
43- return token ;
44- }
45- }
46-
47- // Fallback to query parameter for websocket connections that can't send headers
48- const tokenFromQuery = client . handshake . query . token as string ;
49- if ( tokenFromQuery ) {
50- return tokenFromQuery ;
51- }
52-
53- return null ;
54- }
55-
56- async handleConnection ( client : Socket ) {
57- try {
58- // Extract and validate JWT token
59- const token = this . extractTokenFromHandshake ( client ) ;
60- if ( ! token ) {
61- this . logger . warn ( `Connection rejected: Missing authentication token (client: ${ client . id } )` ) ;
62- client . disconnect ( true ) ;
63- return ;
64- }
65-
66- let authUser : AuthUserPayload ;
67- try {
68- authUser = await this . authService . validateAccessToken ( token ) ;
69- } catch ( error ) {
70- this . logger . warn ( `Connection rejected: Invalid JWT token (client: ${ client . id } )` ) ;
71- client . disconnect ( true ) ;
72- return ;
73- }
74-
75- const userId = authUser . sub ;
76-
77- // Check max connections per user
78- const currentConnections = this . userSockets . get ( userId ) || [ ] ;
79- if ( currentConnections . length >= this . MAX_CONNECTIONS_PER_USER ) {
80- this . logger . warn ( `Connection rejected: User ${ userId } exceeded max connections (${ this . MAX_CONNECTIONS_PER_USER } )` ) ;
81- client . disconnect ( true ) ;
82- return ;
83- }
84-
85- // Add the new connection
86- currentConnections . push ( client . id ) ;
87- this . userSockets . set ( userId , currentConnections ) ;
29+ handleConnection ( client : Socket ) {
30+ const userId = client . handshake . query . userId as string ;
31+ if ( userId ) {
32+ const sockets = this . userSockets . get ( userId ) || [ ] ;
33+ sockets . push ( client . id ) ;
34+ this . userSockets . set ( userId , sockets ) ;
8835 this . socketUsers . set ( client . id , userId ) ;
8936
90- // Join the user's private room
9137 client . join ( `user:${ userId } ` ) ;
9238
93- this . logger . log ( `User ${ userId } connected (${ client . id } ). Active connections: ${ currentConnections . length } /${ this . MAX_CONNECTIONS_PER_USER } ` ) ;
94- } catch ( error ) {
95- this . logger . error ( `Error during connection handling: ${ error . message } ` , error . stack ) ;
96- client . disconnect ( true ) ;
39+ this . logger . log ( `User ${ userId } connected (${ client . id } )` ) ;
9740 }
9841 }
9942
@@ -109,47 +52,12 @@ export class NotificationsGateway implements OnGatewayConnection, OnGatewayDisco
10952 this . userSockets . delete ( userId ) ;
11053 }
11154 this . socketUsers . delete ( client . id ) ;
112- // Clean up rate limit entry for this socket
113- this . eventRateLimits . delete ( client . id ) ;
114- this . logger . log ( `User ${ userId } disconnected (${ client . id } ). Remaining connections: ${ sockets . length } ` ) ;
55+ this . logger . log ( `User ${ userId } disconnected (${ client . id } )` ) ;
11556 }
11657 }
11758
118- private checkRateLimit ( clientId : string ) : boolean {
119- const now = Date . now ( ) ;
120- const rateLimit = this . eventRateLimits . get ( clientId ) ;
121-
122- if ( ! rateLimit ) {
123- // First event, initialize rate limit
124- this . eventRateLimits . set ( clientId , { count : 1 , resetTime : now + this . RATE_LIMIT_WINDOW } ) ;
125- return true ;
126- }
127-
128- // Reset counter if window has expired
129- if ( now > rateLimit . resetTime ) {
130- rateLimit . count = 1 ;
131- rateLimit . resetTime = now + this . RATE_LIMIT_WINDOW ;
132- return true ;
133- }
134-
135- // Check if limit exceeded
136- if ( rateLimit . count >= this . MAX_EVENTS_PER_MINUTE ) {
137- this . logger . warn ( `Rate limit exceeded for client ${ clientId } ` ) ;
138- return false ;
139- }
140-
141- // Increment counter
142- rateLimit . count ++ ;
143- return true ;
144- }
145-
14659 @SubscribeMessage ( 'joinProperty' )
14760 handleJoinProperty ( @ConnectedSocket ( ) client : Socket , @MessageBody ( ) data : { propertyId : string } ) {
148- // Check rate limit
149- if ( ! this . checkRateLimit ( client . id ) ) {
150- return { event : 'error' , data : { message : 'Rate limit exceeded. Please try again later.' } } ;
151- }
152-
15361 if ( data ?. propertyId ) {
15462 client . join ( `property:${ data . propertyId } ` ) ;
15563 this . logger . log ( `Client ${ client . id } joined property room ${ data . propertyId } ` ) ;
@@ -160,29 +68,18 @@ export class NotificationsGateway implements OnGatewayConnection, OnGatewayDisco
16068
16169 @SubscribeMessage ( 'leaveProperty' )
16270 handleLeaveProperty ( @ConnectedSocket ( ) client : Socket , @MessageBody ( ) data : { propertyId : string } ) {
163- // Check rate limit
164- if ( ! this . checkRateLimit ( client . id ) ) {
165- return { event : 'error' , data : { message : 'Rate limit exceeded. Please try again later.' } } ;
166- }
167-
16871 if ( data ?. propertyId ) {
16972 client . leave ( `property:${ data . propertyId } ` ) ;
17073 this . logger . log ( `Client ${ client . id } left property room ${ data . propertyId } ` ) ;
17174 return { event : 'leftProperty' , data : { propertyId : data . propertyId } } ;
17275 }
173- return { event : 'error' , data : { message : 'propertyId is required' } } ;
17476 }
17577
17678 @SubscribeMessage ( 'joinTransaction' )
17779 handleJoinTransaction (
17880 @ConnectedSocket ( ) client : Socket ,
17981 @MessageBody ( ) data : { transactionId : string } ,
18082 ) {
181- // Check rate limit
182- if ( ! this . checkRateLimit ( client . id ) ) {
183- return { event : 'error' , data : { message : 'Rate limit exceeded. Please try again later.' } } ;
184- }
185-
18683 if ( data ?. transactionId ) {
18784 client . join ( `transaction:${ data . transactionId } ` ) ;
18885 this . logger . log ( `Client ${ client . id } joined transaction room ${ data . transactionId } ` ) ;
@@ -196,25 +93,14 @@ export class NotificationsGateway implements OnGatewayConnection, OnGatewayDisco
19693 @ConnectedSocket ( ) client : Socket ,
19794 @MessageBody ( ) data : { transactionId : string } ,
19895 ) {
199- // Check rate limit
200- if ( ! this . checkRateLimit ( client . id ) ) {
201- return { event : 'error' , data : { message : 'Rate limit exceeded. Please try again later.' } } ;
202- }
203-
20496 if ( data ?. transactionId ) {
20597 client . leave ( `transaction:${ data . transactionId } ` ) ;
20698 return { event : 'leftTransaction' , data : { transactionId : data . transactionId } } ;
20799 }
208- return { event : 'error' , data : { message : 'transactionId is required' } } ;
209100 }
210101
211102 @SubscribeMessage ( 'joinUser' )
212103 handleJoinUser ( @ConnectedSocket ( ) client : Socket , @MessageBody ( ) data : { userId : string } ) {
213- // Check rate limit
214- if ( ! this . checkRateLimit ( client . id ) ) {
215- return { event : 'error' , data : { message : 'Rate limit exceeded. Please try again later.' } } ;
216- }
217-
218104 const socketUserId = this . socketUsers . get ( client . id ) ;
219105 if ( socketUserId && socketUserId === data ?. userId ) {
220106 client . join ( `user:${ data . userId } ` ) ;
@@ -289,4 +175,4 @@ export class NotificationsGateway implements OnGatewayConnection, OnGatewayDisco
289175 sendToAll ( event : string , data : any ) {
290176 this . server . emit ( event , data ) ;
291177 }
292- }
178+ }
0 commit comments