Exponential Backoff Guard in NestJS

exp-backoff-throttler.guard.ts

import { ExecutionContext, Injectable } from '@nestjs/common';
import { ThrottlerGuard, ThrottlerLimitDetail, ThrottlerRequest } from '@nestjs/throttler';

type BlockEntry = {
    minDuration: number;
    blockDuration: number;
    totalHits: number;
    blockedUntil: number;
};

@Injectable()
export class ExpBackoffThrottlerGuard extends ThrottlerGuard {
    private blockMap: Map<string, BlockEntry> = new Map();

    private readonly MAX_BLOCK_DURATION = 5 * 60 * 1000;

    protected getTracker(req: Record<string, any>): Promise<string> {
        return req.user?.id || req.ip;
    }

    protected async handleRequest(requestProps: ThrottlerRequest): Promise<boolean> {
        const getThrottlerSuffix = (name?: string) => (name === 'default' ? '' : `-${name}`);

        const {
            context,
            throttler,
            limit,
            ttl,
            blockDuration: dur,
            getTracker,
            generateKey,
        } = requestProps;
        const { req, res } = this.getRequestResponse(context);
        const tracker = await getTracker(req, context);
        const key = generateKey(context, tracker, throttler.name || 'default');
        if (!this.blockMap.has(key))
            this.blockMap.set(key, {
                minDuration: dur,
                blockDuration: dur,
                totalHits: 0,
                blockedUntil: 0,
            });
        const blockEntry = this.blockMap.get(key) as BlockEntry;
        const { blockDuration, blockedUntil, totalHits } = blockEntry;
        const now = Date.now();
        if (blockedUntil && now < blockedUntil) {
            // Any hit within the block period will throw
            const remTime = Math.ceil(blockDuration / 1e3);
            res.header(`Retry-After${getThrottlerSuffix(throttler.name)}`, remTime);
            await this.throwThrottlingException(context, {
                limit,
                ttl,
                key,
                tracker,
                totalHits: totalHits + 1,
                timeToExpire: remTime,
                isBlocked: true,
                timeToBlockExpire: remTime,
            });
        }
        const allowed = await super.handleRequest({
            ...requestProps,
            blockDuration,
        });
        // If we get here, the request is allowed and we lift the block
        if (blockedUntil) {
            this.blockMap.set(key, {
                ...blockEntry,
                blockDuration: blockEntry.minDuration,
                blockedUntil: 0,
            });
        }
        return allowed;
    }

    protected throwThrottlingException(
        context: ExecutionContext,
        throttlerLimitDetail: ThrottlerLimitDetail,
    ): Promise<void> {
        const { key, totalHits } = throttlerLimitDetail;
        const blockEntry = this.blockMap.get(key) as BlockEntry;
        const { blockDuration } = blockEntry;
        this.blockMap.set(key, {
            ...blockEntry,
            blockedUntil: Date.now() + blockDuration,
            // Double the block duration for every hit (up to the max)
            blockDuration: Math.min(this.MAX_BLOCK_DURATION, blockDuration * 2),
            totalHits,
        });
        return super.throwThrottlingException(context, throttlerLimitDetail);
    }
}