qps-path.js 2.72 KB
/**
 * 控制路由请求次数
 * @date: 2018/03/05
 */
'use strict';

const _ = require('lodash');
const cache = global.yoho.cache.master;
const logger = global.yoho.logger;
const md5 = require('md5');
const pathToRegexp = require('path-to-regexp');
const Promise = require('bluebird');

const zk = require('./zk');
const {limitKey} = require('./vars');

const INVALIDTIME = 3600 * 24; // 24h
const IP_WHITE_LIST = [
    '106.38.38.146',
    '106.38.38.147',
    '106.39.86.227',
    '218.94.75.58',
    '218.94.75.50',
    '218.94.77.166'
];

module.exports = async({user}, next) => {
    if (!user.app || !user.path || !user.ip) {
        return next();
    }

    const app = user.app;

    if (_.get(zk, `${app}.close.risk`, false)) {
        return next();
    }

    const ip = user.ip;
    const path = user.path;
    // const risks = _.get(zk, `${app}.json.risk`, [{route: '/product/(.*).html', interval: 5000, requests: 10}]);
    const risks = _.get(zk, `${app}.json.risk`, []);
    let router = {};

    if (_.isEmpty(path) || _.isEmpty(risks) || IP_WHITE_LIST.indexOf(ip) > -1) {
        return next();
    }

    _.isArray(risks) && risks.some(item => {
        if (item.state === 'off') {
            return false;
        }

        if (!item.regRoute) {
            item.regRoute = pathToRegexp(item.route);
            item.interval = parseInt(item.interval, 10);
            item.requests = parseInt(item.requests, 10);
        }

        if (item.regRoute.test(path)) {
            router = item;
            return true;
        }

        return false;
    });

    logger.error(`risk==> router: ${JSON.stringify(router)}, path: ${path}`); // eslint-disable-line

    if (_.isEmpty(router)) {
        return next();
    }

    let keyPath = md5(`${router.regRoute}`);
    let limitEnable = `${app}:risk:${limitKey}:${keyPath}:${ip}`; // 查询这个key是否生效
    let configKey = `${app}:risk:count:${keyPath}:${ip}`;

    const inters = await Promise.all([
        cache.getAsync(limitEnable),
        cache.getAsync(configKey),
    ]);

    logger.debug(`risk==> cache: %s %d`, limitEnable, inters[0], configKey, inters[1]); // eslint-disable-line

    if (inters[0]) {
        logger.debug('[qps:route] this user[%j] has rejected', user);
        return;
    }

    if (!inters[1]) {
        cache.setAsync(configKey, 1, router.interval || 300);
        return next();
    }

    inters[1] = parseInt(`0${inters[1]}`, 10);
    if (inters[1] <= router.requests) {
        router = [];
        cache.incrAsync(configKey, 1);
        return next();
    }

    logger.info('[qps:route] this user[%j] is being marked as rejected', user);
    await Promise.all([
        cache.setAsync(limitEnable, 1, INVALIDTIME),
        cache.delAsync(configKey)
    ]);
};