qps-limit.js 3.47 KB
'use strict';

const logger = global.yoho.logger;
const cache = global.yoho.cache.master;
const config = global.yoho.config;
const ONE_DAY = 60 * 60 * 24;
const MAX_QPS = config.maxQps;
const MAX_QPS_10m = config.maxQps10m; // eslint-disable-line
const _ = require('lodash');

const PAGES = {
    '/product/^\\/([\\d]+)(.*)/': 5,
    '/product/list/index': 5,
    '/product/search/index': 5
};

function urlJoin(a, b) {
    if (_.endsWith(a, '/') && _.startsWith(b, '/')) {
        return a + b.substring(1, b.length);
    } else if (!_.endsWith(a, '/') && !_.startsWith(b, '/')) {
        return a + '/' + b;
    } else {
        return a + b;
    }
}

module.exports = (limiter, policy) => {
    const req = limiter.req,
        res = limiter.res,
        next = limiter.next; // eslint-disable-line

    const key = `pc:limiter:${limiter.remoteIp}`;
    const keyMax = `pc:limiter:max:${limiter.remoteIp}`;
    const key10m = `pc:limiter:10m:${limiter.remoteIp}`;
    const key10mMax = `pc:limiter:10m:max:${limiter.remoteIp}`;

    res.on('render', function() {
        let route = req.route ? req.route.path : '';
        let appPath = req.app.mountpath;

        if (_.isArray(route) && route.length > 0) {
            route = route[0];
        }

        let pageKey = urlJoin(appPath, route.toString()); // route may be a regexp
        let pageIncr = PAGES[pageKey] || 0;

        if (pageIncr > 0) {
            cache.incr(key, pageIncr, (err) => {}); // eslint-disable-line
            cache.incr(key10m, pageIncr, (err) => {}); // eslint-disable-line
        }
    });

    return cache.getMultiAsync([key, key10m, keyMax, key10mMax]).then((results) => {
        let result = results[key];
        let result10m = results[key10m];

        logger.debug('qps limiter: ' + key + '@' + result + ' max: ' + MAX_QPS);
        logger.debug('qps limiter:10m ' + key10m + '@' + result10m + ' max: ' + MAX_QPS_10m); // eslint-disable-line

        // 达到1分钟或是10分钟的访问限制,禁止访问
        if (results[keyMax] === 1 || results[key10mMax] === 1) {
            return Promise.resolve(policy);
        }

        // 默认数据设置
        if (!result && !_.isNumber(result)) {
            cache.setAsync(key, 1, 60);  // 设置key,1m失效
        }

        if (!result10m && !_.isNumber(result10m)) {
            cache.setAsync(key10m, 1, 600);  // 设置key,10m失效
        }

        // 第一次访问,都没计数,直接过
        if (!result && !_.isNumber(result) && !result10m && !_.isNumber(result10m)) {
            return Promise.resolve(true);
        }

        if (result === -1 || result10m === -1) {
            return Promise.resolve(true);
        }

        // 判断 qps 10分钟
        if (result10m === 9999) {
            res.statusCode = 403;
            return Promise.resolve(policy);
        } else if (result10m > MAX_QPS_10m) { // eslint-disable-line
            cache.setAsync(key10mMax, 1, ONE_DAY);
            logger.debug('req limit', key10m);

            return Promise.resolve(policy);
        }

        // 判断 qps 1分钟
        if (result === 9999) {
            res.statusCode = 403;
            return Promise.resolve(policy);
        } else if (result > MAX_QPS) { // 判断 qps
            cache.setAsync(keyMax, 1, ONE_DAY);
            logger.debug('req limit', key);

            return Promise.resolve(policy);
        }

        cache.incrAsync(key, 1); // qps + 1
        cache.incrAsync(key10m, 1); // qps + 1
        return Promise.resolve(true);
    });
};