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

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

const zk = {};

if (config.zookeeperServer) {
    require('yoho-zookeeper')(config.zookeeperServer, 'pc', zk.pc = {}, global.yoho.cache);
    require('yoho-zookeeper')(config.zookeeperServer, 'wap', zk.wap = {}, global.yoho.cache);
}

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'
];

const key = 'risk1';

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`, []);
    let router = {};

    logger.debug(`risk => risks: ${JSON.stringify(risks)}, path: ${path}, ip: ${ip}`); // eslint-disable-line
    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.debug(`risk => router: ${JSON.stringify(router)}, path: ${path}`); // eslint-disable-line
    if (_.isEmpty(router)) {
        return next();
    }

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

    await Promise.all([
        cache.getAsync(limitKey),
        cache.getAsync(configKey),
    ]).then(inters => {
        logger.debug(`risk => getCache: ${JSON.stringify(inters)}, path: ${path}`); // eslint-disable-line
        if (inters[0]) {
            logger.info('[qps:route] this user[%o] has rejected', user);
            return;
        }

        if (typeof inters[1] === 'undefined') {
            cache.setAsync(configKey, 1, router.interval || 300);
            return;
        }

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

        logger.warn('[qps:route] this user[%o] is being marked as rejected', user);
        return Promise.all([
            cache.setAsync(limitKey, 1, INVALIDTIME),
            cache.delAsync(configKey)
        ]);
    }).then(result => {
        logger.debug('[qps:route] user[%o] result[%o]', user, result); // eslint-disable-line
    }).catch(e => {
        logger.error(`risk => path: ${path}, err: ${e.message}`);
    }).finally(() => {
        return next();
    });
};