import test from 'ava' import type { onConnectPayload } from '@hocuspocus/server' import * as MockDate from 'mockdate' import { Throttle } from '@hocuspocus/extension-throttle' const getOnConnectPayload = (ip: string) => { return { request: { headers: new Headers({ 'x-real-ip': ip, }), }, } as unknown as onConnectPayload } const generateRequests = async (instance: Throttle, ip: string, numberOfRequests: number) => { for (let i = 0; i < numberOfRequests; i += 1) { // eslint-disable-next-line no-await-in-loop await instance.onConnect(getOnConnectPayload(ip)) } } test('throttle extension bans properly', async t => { const throttle = new Throttle({ banTime: 5, throttle: 15 }) t.teardown(() => throttle.onDestroy()) const ip = '127.0.0.1' t.false(throttle.isBanned(ip)) await generateRequests(throttle, ip, 15) try { await throttle.onConnect(getOnConnectPayload(ip)) t.fail() } catch (e) { t.true(throttle.isBanned(ip)) } }) test('throttle extension unbans properly', async t => { const throttle = new Throttle({ banTime: 5, throttle: 15 }) t.teardown(() => throttle.onDestroy()) const ip = '127.0.0.1' t.false(throttle.isBanned(ip)) await generateRequests(throttle, ip, 15) try { await throttle.onConnect(getOnConnectPayload(ip)) t.fail() } catch (e) { t.true(throttle.isBanned(ip)) } MockDate.set(Date.now() + 1000 * (throttle.configuration.banTime * 60)) await throttle.onConnect(getOnConnectPayload(ip)) t.false(throttle.isBanned(ip)) MockDate.reset() }) test.serial('map cleanup works for connectionsByIp', async t => { const throttle = new Throttle({ consideredSeconds: 60 }) t.teardown(() => throttle.onDestroy()) const ip = '127.0.0.1' await generateRequests(throttle, ip, 10) t.is(throttle.connectionsByIp.get(ip)!.length, 10) MockDate.set(Date.now() + 1000 * throttle.configuration.consideredSeconds) await throttle.clearMaps() t.false(throttle.connectionsByIp.has(ip)) MockDate.reset() }) test.serial('map cleanup works for bannedIps', async t => { const throttle = new Throttle({ consideredSeconds: 60, throttle: 15 }) t.teardown(() => throttle.onDestroy()) const ip = '127.0.0.1' await generateRequests(throttle, ip, 15) try { await throttle.onConnect(getOnConnectPayload(ip)) // eslint-disable-next-line no-empty } catch (e) {} t.true(throttle.bannedIps.has(ip)) MockDate.set(Date.now() + 1000 * throttle.configuration.banTime * 60) await throttle.clearMaps() t.false(throttle.bannedIps.has(ip)) MockDate.reset() })