qps-limit.js
2.14 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
'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; // eslint-disable-line
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) => {}); // eslint-disable-line
}
});
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);
}
});
};