diff --git a/cvmts/src/IPData.ts b/cvmts/src/IPData.ts index ce53011..676b9e8 100644 --- a/cvmts/src/IPData.ts +++ b/cvmts/src/IPData.ts @@ -40,6 +40,19 @@ export class IPDataManager { return data; } + static GetIPDataMaybe(address: string) { + if (IPDataManager.ipDatas.has(address)) { + // Note: We already check for if it exists, so we use ! here + // because TypeScript can't exactly tell that in this case, + // only in explicit null or undefined checks + let ref = IPDataManager.ipDatas.get(address)!; + ref.refCount++; + return ref; + } + + return null; + } + static ForEachIPData(callback: (d: IPData) => void) { for (let tuple of IPDataManager.ipDatas) callback(tuple[1]); } diff --git a/cvmts/src/WebSocket/WSServer.ts b/cvmts/src/WebSocket/WSServer.ts index 71f4fdd..918cfe6 100644 --- a/cvmts/src/WebSocket/WSServer.ts +++ b/cvmts/src/WebSocket/WSServer.ts @@ -119,15 +119,15 @@ export default class WSServer extends EventEmitter implements NetworkServer { ip = req.socket.remoteAddress; } - // TODO: Implement + let ipdata = IPDataManager.GetIPDataMaybe(ip); - // Get the amount of active connections coming from the requesting IP. - //let connections = this.clients.filter((client) => client.IP.address == ip); - // If it exceeds the limit set in the config, reject the connection with a 429. - //if (connections.length + 1 > this.Config.http.maxConnections) { - // socket.write('HTTP/1.1 429 Too Many Requests\n\n429 Too Many Requests'); - // socket.destroy(); - //} + if(ipdata != null) { + let connections = ipdata.refCount; + if (connections + 1 > this.Config.collabvm.maxConnections) { + socket.write('HTTP/1.1 429 Too Many Requests\n\n429 Too Many Requests'); + socket.destroy(); + } + } this.wsServer.handleUpgrade(req, socket, head, (ws: WebSocket) => { this.wsServer.emit('connection', ws, req); @@ -149,4 +149,4 @@ export default class WSServer extends EventEmitter implements NetworkServer { this.logger.Info(`New WebSocket connection from ${user.IP.address}`); } -} \ No newline at end of file +}