From cee765fecaf863dd5e0e1bfdc767ef38bb517fb0 Mon Sep 17 00:00:00 2001
From: Justin Langston <justin@bitpay.com>
Date: Sun, 23 Dec 2018 11:27:15 -0500
Subject: [PATCH] feat(api): Rate Limits

---
 packages/bitcore-node/src/config.ts           |  5 ++
 packages/bitcore-node/src/models/rateLimit.ts | 56 +++++++++++++++++++
 packages/bitcore-node/src/routes/index.ts     |  3 +-
 .../bitcore-node/src/routes/middleware.ts     | 23 ++++++++
 packages/bitcore-node/src/types/Config.ts     |  3 +
 5 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 packages/bitcore-node/src/models/rateLimit.ts

diff --git a/packages/bitcore-node/src/config.ts b/packages/bitcore-node/src/config.ts
index 74eefebaf72..deceb89c3dc 100644
--- a/packages/bitcore-node/src/config.ts
+++ b/packages/bitcore-node/src/config.ts
@@ -58,6 +58,11 @@ const Config = function(): ConfigType {
     dbPort: process.env.DB_PORT || '27017',
     numWorkers: cpus().length,
     api: {
+      rateLimiter: {
+        whitelist: [
+          '::ffff:127.0.0.1'
+        ]
+      },
       wallets: {
         allowCreationBeforeCompleteSync: false,
         allowUnauthenticatedCalls: false
diff --git a/packages/bitcore-node/src/models/rateLimit.ts b/packages/bitcore-node/src/models/rateLimit.ts
new file mode 100644
index 00000000000..9add3b812ed
--- /dev/null
+++ b/packages/bitcore-node/src/models/rateLimit.ts
@@ -0,0 +1,56 @@
+import { BaseModel } from './base';
+import { ObjectID } from 'mongodb';
+
+export type IRateLimit = {
+  _id?: ObjectID;
+  identifier: string;
+  method: string;
+  period: string;
+  count: number;
+  time?: Date;
+  expireAt?: Date;
+  value?: any;
+};
+
+export class RateLimit extends BaseModel<IRateLimit> {
+  constructor() {
+    super('ratelimits');
+  }
+  allowedPaging = [];
+
+  onConnect() {
+    this.collection.createIndex({ identifier: 1, time: 1, method: 1, count: 1 }, { background: true });
+    this.collection.createIndex({ expireAt: 1 }, { expireAfterSeconds: 0, background: true });
+  }
+
+  incrementAndCheck(identifier: string, method: string) {
+    return Promise.all([
+      this.collection.findOneAndUpdate(
+        { identifier, method, period: 'second', time: {$gt: new Date(Date.now() - 1000)} },
+        {
+          $setOnInsert: { time: new Date(), expireAt: new Date(Date.now() + 10 * 1000)  },
+          $inc: { count: 1 }
+        },
+        { upsert: true, returnOriginal: false }
+      ),
+      this.collection.findOneAndUpdate(
+        { identifier, method, period: 'minute', time: { $gt: new Date(Date.now() - 60 * 1000) } },
+        {
+          $setOnInsert: { time: new Date(), expireAt: new Date(Date.now() + 2 * 60 * 1000) },
+          $inc: { count: 1 }
+        },
+        { upsert: true, returnOriginal: false }
+      ),
+      this.collection.findOneAndUpdate(
+        { identifier, method, period: 'hour', time: { $gt: new Date(Date.now() - 60 * 60 * 1000) } },
+        {
+          $setOnInsert: { time: new Date(), expireAt: new Date(Date.now() + 2 * 60 * 1000) },
+          $inc: { count: 1 }
+        },
+        { upsert: true, returnOriginal: false }
+      ),
+    ]); 
+  }
+}
+
+export let RateLimitModel = new RateLimit();
diff --git a/packages/bitcore-node/src/routes/index.ts b/packages/bitcore-node/src/routes/index.ts
index 026196e28ef..25c3f491bbb 100644
--- a/packages/bitcore-node/src/routes/index.ts
+++ b/packages/bitcore-node/src/routes/index.ts
@@ -2,10 +2,11 @@ import config from '../config';
 import { Request, Response } from 'express';
 import express from 'express';
 import cors from 'cors';
-import { LogRequest } from "./middleware";
+import { LogRequest, RateLimiter } from "./middleware";
 
 const app = express();
 const bodyParser = require('body-parser');
+app.use(RateLimiter('GLOBAL', 10, 200, 4000));
 app.use(bodyParser.json());
 app.use(
   bodyParser.raw({
diff --git a/packages/bitcore-node/src/routes/middleware.ts b/packages/bitcore-node/src/routes/middleware.ts
index 8988c670eb0..27b979dc2bb 100644
--- a/packages/bitcore-node/src/routes/middleware.ts
+++ b/packages/bitcore-node/src/routes/middleware.ts
@@ -1,5 +1,7 @@
 import logger from '../logger';
 import * as express from 'express';
+import { RateLimitModel } from '../models/rateLimit';
+import config from '../config';
 
 type TimedRequest = {
   startTime?: Date;
@@ -42,3 +44,24 @@ export function LogRequest(req: TimedRequest, res: express.Response, next: expre
   res.on('close', LogPhase('CLOSED'));
   next();
 }
+
+export function RateLimiter(method: string, perSecond: number, perMinute: number, perHour: number) {
+  return async (req: express.Request, res: express.Response, next: express.NextFunction) => {
+    try {
+      const identifier = req.header('CF-Connecting-IP') || req.socket.remoteAddress || '';
+      if (config.api.rateLimiter.whitelist.includes(identifier)) {
+        return next();
+      }
+      let [perSecondResult, perMinuteResult, perHourResult] = await RateLimitModel.incrementAndCheck(identifier, method);
+      if (
+        (perSecondResult.value as any).count > perSecond ||
+        (perMinuteResult.value as any).count > perMinute ||
+        (perHourResult.value as any).count > perHour) {
+        return res.status(429).send('Rate Limited');
+      }
+    } catch (err) {
+      logger.error('Rate Limiter failed');
+    }
+    return next();
+  }
+}
diff --git a/packages/bitcore-node/src/types/Config.ts b/packages/bitcore-node/src/types/Config.ts
index c0d75da9d27..9d63ffdb079 100644
--- a/packages/bitcore-node/src/types/Config.ts
+++ b/packages/bitcore-node/src/types/Config.ts
@@ -10,6 +10,9 @@ export default interface Config {
     [currency: string]: any;
   };
   api: {
+    rateLimiter: {
+      whitelist: [string]; 
+    },
     wallets: {
       allowCreationBeforeCompleteSync?: boolean;
       allowUnauthenticatedCalls?: boolean;