qps-limit.js 2.1 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 _ = require('lodash');

const PAGES = {
    '/product/\\/pro_([\\d]+)_([\\d]+)\\/(.*)/': 5,
    '/product/list/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;

    const key = `pc:limiter:${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 (/^\/p([\d]+)/.test(req.path)) {
            pageIncr = 5;
        }

        if (pageIncr > 0) {
            cache.incr(key, pageIncr, (err) => {});
        }
    });

    return cache.getAsync(key).then((result) => {
        logger.debug('qps limiter: ' + key + '@' + result + ' max: ' + MAX_QPS);
        if (result && _.isNumber(result)) {
            if (result === -1) {
                return Promise.resolve(true);
            }

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

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

            }
        } else {
            cache.setAsync(key, 1, 60); // 设置key,1m失效
            return Promise.resolve(true);
        }
    });
};