in src/yaral.ts [176:418]
async register(server, options) {
options = {
cache: '_default',
enabled: true,
includeHeaders: true,
default: [],
exclude: () => false,
onLimit: () => {
/*do nothing*/
},
onPass: () => {
/*do nothing*/
},
event: 'onPreAuth',
//Timeout enabled by default with a value of 1000 ms. On timeout, by default we continue.
timeout: {
enabled: true,
timeout: 1000,
ontimeout: (_request, reply) => reply.continue,
},
...options,
};
Joi.assert(options, schema);
// If we aren't enabled, don't bother doing anything.
if (!options.enabled) {
return;
}
const limitus = options.limitus || buildLimitus(server, options.cache);
const buckets: {
[key: string]: Bucket;
} = {};
options.buckets.forEach(bucket => {
const b = new Bucket(bucket, limitus);
buckets[b.name()] = b;
});
/**
* Returns a configuration object for the route based on its specific
* rules.
*/
function resolveRouteOpts(
req: Request,
): {
enabled: boolean;
buckets: string[];
exclude: (req: Request) => boolean;
} {
const routeOpts = req.route.settings.plugins.yaral;
const opts = {
enabled: true,
buckets: <string[]>[],
exclude: (_req: Request) => false,
};
if (!routeOpts) {
// do nothing
} else if (typeof routeOpts === 'string') {
// specifying bucket as string
opts.buckets = [routeOpts];
} else if (Array.isArray(routeOpts)) {
// specifying array of buckets
opts.buckets = routeOpts;
} else {
// providing a literal object
Object.assign(opts, routeOpts);
if (typeof opts.buckets === 'string') {
opts.buckets = [opts.buckets];
}
}
if (opts.enabled) {
opts.buckets = opts.buckets.concat(options.default);
opts.enabled =
!(options.exclude(req) || opts.exclude(req)) &&
opts.buckets.length + options.default.length > 0;
}
return opts;
}
/**
* Returns the bucket used to rate limit the specified response.
*/
const matchBucket = (
info: IYaralInternalData,
res: ResponseObject,
): {
bucket: Bucket;
id: string | number | object;
} => {
for (let i = 0; i < info.buckets.length; i++) {
const bucket = buckets[info.buckets[i]];
if (bucket.matches(res.statusCode)) {
return { bucket: bucket, id: info.ids[i] };
}
}
return undefined;
};
/**
* Adds rate limit headers to the response if they're set.
*/
const addHeaders = (
res: Boom.Output | ResponseObject,
headers: { [key: string]: string | number },
) => {
if (options.includeHeaders) {
Object.assign(res.headers, headers);
}
};
const getRequestLogDetails = (
err: Error,
req: Request,
isTimedout: boolean,
duration: number,
) => {
return {
name: 'yaral-timeout',
url: req.url.href || '',
duration: duration,
success: !err,
properties: {
error: err ? err.stack : '',
isTimedout: isTimedout,
},
};
};
async function createTimeout<T>(call: () => Promise<T>, req: Request): Promise<T> {
const startTime = Date.now();
const p = call();
if (!options.timeout.enabled) {
return p;
}
try {
return await new Promise<T>((resolve, reject) => {
const t = setTimeout(
() => reject(new TimeoutError('Call Timed Out')),
options.timeout.timeout,
);
p.then(innerV => {
clearTimeout(t);
resolve(innerV);
}, reject);
});
} catch (e) {
server.log(
['ratelimit', 'timeout'],
getRequestLogDetails(e, req, true, Date.now() - startTime),
);
throw e;
}
}
//Appropriately handles different types of Errors
const handleError = (err: Error, req: Request, reply: ResponseToolkit) => {
//In case there is a redis timeout continue executing
//did not put it in the same block as Limitus.Rejected for the sake of future logging
if (err instanceof TimeoutError) {
server.log(
['ratelimit', 'timeout'],
getRequestLogDetails(err, req, true, options.timeout.timeout),
);
return options.timeout.ontimeout(req, reply);
}
if (!(err instanceof Limitus.Rejected)) {
server.log(['error', 'ratelimit'], err);
}
// Internal errors should not halt everything.
return reply.continue;
};
server.ext(options.event, async (req, tk) => {
const opts = resolveRouteOpts(req);
if (opts.enabled === false) {
return tk.continue;
}
const info = {
buckets: opts.buckets,
ids: opts.buckets.map(b => buckets[b].identify(req)),
limited: false,
};
req.plugins.yaral = info;
try {
await Promise.all(
info.buckets.map((name, i) =>
createTimeout(() => limitus.checkLimited(name, info.ids[i]), req),
),
);
options.onPass(req);
return tk.continue;
} catch (err) {
//Internal Error or Timeout Error
if (!(err instanceof Limitus.Rejected)) {
return handleError(err, req, tk);
}
// FIXME: limitus: this is actually a number??
const reset: number = <any>err.bucketName;
const headers = {
'X-RateLimit-Remaining': '0',
'X-RateLimit-Reset': reset.toString(),
};
// Continue the request if onLimit dictates that we cancel limiting.
const lResponse = options.onLimit(req, tk, err.info, reset, headers);
if (lResponse) {
if (lResponse === cancel) {
return tk.continue;
}
return lResponse;
}
info.limited = true;
const res = Boom.tooManyRequests();
addHeaders(res.output, headers);
throw res;
}
});
server.ext('onPreResponse', async (req, reply) => {
const res = (<Boom<any>>req.response).output || <ResponseObject>req.response;
const opts = req.plugins.yaral && matchBucket(req.plugins.yaral, <ResponseObject>res);
if (!opts || req.plugins.yaral.limited) {
return reply.continue;
}
try {
const data = await createTimeout(() => limitus.drop(opts.bucket.name(), opts.id), req);
addHeaders(res, opts.bucket.headers(data));
return reply.continue;
} catch (err) {
return handleError(err, req, reply);
}
});
},