diff --git a/index.js b/index.js index bcf94a4..3061ad5 100644 --- a/index.js +++ b/index.js @@ -7,8 +7,6 @@ process.on('uncaughtException', err => { const httpWrapper = require('./libs/http-wrapper'); const Server = httpWrapper.Server; -const TokenBucket = require('./libs/TokenBucket'); -const START_TOKENS = 6000000; const CHECK_FOR_RESTART_COMMAND_FILE_INTERVAL = 500; (function loadDefaultComponents(){ @@ -41,8 +39,6 @@ function HttpServer({ listeningPort, rootFolder, sslConfig, dynamicPort, restart } let port = listeningPort || 8080; - const tokenBucket = new TokenBucket(START_TOKENS, 1, 10); - const conf = require('./config').getConfig(); const server = new Server(sslConfig); server.config = conf; @@ -127,31 +123,6 @@ function HttpServer({ listeningPort, rootFolder, sslConfig, dynamicPort, restart next(); }); - if (conf.preventRateLimit !== true) { - server.use(function (req, res, next) { - const ip = res.socket.remoteAddress; - tokenBucket.takeToken(ip, tokenBucket.COST_MEDIUM, function (err, remainedTokens) { - res.setHeader('X-RateLimit-Limit', tokenBucket.getLimitByCost(tokenBucket.COST_MEDIUM)); - res.setHeader('X-RateLimit-Remaining', tokenBucket.getRemainingTokenByCost(remainedTokens, tokenBucket.COST_MEDIUM)); - - if (err) { - if (err === TokenBucket.ERROR_LIMIT_EXCEEDED) { - res.statusCode = 429; - } else { - res.statusCode = 500; - } - - res.end(); - return; - } - - next(); - }); - }); - } else { - logger.trace(`Rate limit mechanism disabled!`); - } - server.options('/*', function (req, res) { const headers = {}; // IE8 does not allow domains to be specified, just the * @@ -176,6 +147,7 @@ function HttpServer({ listeningPort, rootFolder, sslConfig, dynamicPort, restart function addRootMiddlewares() { const LoggerMiddleware = require('./middlewares/logger'); const AuthorisationMiddleware = require('./middlewares/authorisation'); + const Throttler = require('./middlewares/throttler'); const OAuth = require('./middlewares/oauth'); const ResponseHeaderMiddleware = require('./middlewares/responseHeader'); const genericErrorMiddleware = require('./middlewares/genericErrorMiddleware'); @@ -187,6 +159,7 @@ function HttpServer({ listeningPort, rootFolder, sslConfig, dynamicPort, restart genericErrorMiddleware(server); requestEnhancements(server); + Throttler(server); if(conf.enableJWTAuthorisation) { new AuthorisationMiddleware(server); diff --git a/middlewares/throttler/index.js b/middlewares/throttler/index.js new file mode 100644 index 0000000..305f30a --- /dev/null +++ b/middlewares/throttler/index.js @@ -0,0 +1,55 @@ +const TokenBucket = require("../../libs/TokenBucket"); + +function Throttler(server){ + const logger = $$.getLogger("apihub", "Throttler"); + const START_TOKENS = 6000000; + const tokenBucket = new TokenBucket(START_TOKENS, 1, 10); + let remainingTokens = START_TOKENS; + const conf = require('../../config').getConfig(); + function throttlerHandler(req, res, next) { + const ip = res.socket.remoteAddress; + tokenBucket.takeToken(ip, tokenBucket.COST_MEDIUM, function (err, remainedTokens) { + res.setHeader('X-RateLimit-Limit', tokenBucket.getLimitByCost(tokenBucket.COST_MEDIUM)); + res.setHeader('X-RateLimit-Remaining', tokenBucket.getRemainingTokenByCost(remainedTokens, tokenBucket.COST_MEDIUM)); + remainingTokens = remainedTokens; + if (err) { + if (err === TokenBucket.ERROR_LIMIT_EXCEEDED) { + res.statusCode = 429; + } else { + res.statusCode = 500; + } + + res.end(); + return; + } + + next(); + }); + } + + function readyProbeHandler(req, res) { + const stats = { + remainingTokens: tokenBucket.getRemainingTokenByCost(remainingTokens, tokenBucket.COST_MEDIUM), + tokenLimit: tokenBucket.getLimitByCost(tokenBucket.COST_MEDIUM) + } + + res.setHeader('Content-Type', 'application/json'); + res.statusCode = 200; + res.write(JSON.stringify(stats)); + res.end(); + } + + if (conf.preventRateLimit !== true) { + server.use(throttlerHandler); + server.get("/ready-probe", readyProbeHandler); + } else { + logger.trace(`Rate limit mechanism disabled!`); + server.get("/ready-probe", function (req, res) { + res.statusCode = 200; + res.write("Server ready"); + res.end(); + }); + } +} + +module.exports = Throttler; \ No newline at end of file