[Rest Server] Add rate limit for RESTful API (#4418)

Add rate limit for RESTful API, defaults to:
* 600 requests per ip every 1 min
* 60 list job requests per user every 1 min
* 60 submit job requests per user every 1 hour

src/pylon/deploy/pylon-config/location.conf.template

@@ -29,7 +29,9 @@ location ~ ^/a/([^/]+):(\d*)(.*)$ {
 location ~ ^/rest-server/api(.*)$ {
   proxy_pass {{REST_SERVER_URI}}/api$1$is_args$args;
   proxy_redirect {{REST_SERVER_URI}}/api http://$http_host/rest-server/api;
-  #
+
+  proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+  proxy_set_header X-Real-IP $remote_addr;
   proxy_set_header Accept-Encoding "";
   subs_filter_types *;
   subs_filter

src/rest-server/config/rest-server.yaml

@@ -18,6 +18,9 @@
 service_type: "common"
 
 server-port: 9186
+rate-limit-api-per-min: 600
+rate-limit-list-job-per-min: 60
+rate-limit-submit-job-per-hour: 60
 jwt-secret: pai-secret
 jwt-expire-time: '7d'
 github-owner: Microsoft

src/rest-server/config/rest_server.py

@@ -49,16 +49,14 @@ def run(self):
         service_object_model = dict()
 
         service_object_model['uri'] = 'http://{0}:{1}'.format(master_ip, server_port)
-        service_object_model['server-port'] = server_port
-        service_object_model['jwt-secret'] = self.service_configuration['jwt-secret']
-        service_object_model['jwt-expire-time'] = self.service_configuration['jwt-expire-time']
-        service_object_model['default-pai-admin-username'] = self.service_configuration['default-pai-admin-username']
-        service_object_model['default-pai-admin-password'] = self.service_configuration['default-pai-admin-password']
-        service_object_model['github-owner'] = self.service_configuration['github-owner']
-        service_object_model['github-repository'] = self.service_configuration['github-repository']
-        service_object_model['github-path'] = self.service_configuration['github-path']
-        service_object_model['debugging-reservation-seconds'] = self.service_configuration['debugging-reservation-seconds']
-        service_object_model['enable-priority-class'] = self.service_configuration['enable-priority-class']
+        for k in [
+            'server-port', 'rate-limit-api-per-min', 'rate-limit-list-job-per-min',
+            'rate-limit-submit-job-per-hour', 'jwt-secret', 'jwt-expire-time',
+            'default-pai-admin-username', 'default-pai-admin-password',
+            'github-owner', 'github-repository', 'github-path',
+            'debugging-reservation-seconds', 'enable-priority-class'
+        ]:
+            service_object_model[k] = self.service_configuration[k]
         service_object_model['etcd-uris'] = ','.join('http://{0}:4001'.format(host['hostip'])
                                                      for host in machine_list
                                                      if host.get('k8s-role') == 'master')

src/rest-server/deploy/rest-server.yaml.template

@@ -87,6 +87,12 @@ spec:
         - name: YARN_URI
           value: http://{{ cluster_cfg['hadoop-resource-manager']['master-ip'] }}:8088
 {%- endif %}
+        - name: RATE_LIMIT_API_PER_MIN
+          value: "{{ cluster_cfg['rest-server']['rate-limit-api-per-min'] }}"
+        - name: RATE_LIMIT_LIST_JOB_PER_MIN
+          value: "{{ cluster_cfg['rest-server']['rate-limit-list-job-per-min'] }}"
+        - name: RATE_LIMIT_SUBMIT_JOB_PER_HOUR
+          value: "{{ cluster_cfg['rest-server']['rate-limit-submit-job-per-hour'] }}"
         - name: JWT_SECRET
           value: {{ cluster_cfg['rest-server']['jwt-secret'] }}
         - name: JWT_TOKEN_EXPIRE_TIME
@@ -139,7 +145,7 @@ spec:
         {%- if cluster_cfg['cluster']['common']['qos-switch'] == "true" %}
         resources:
           limits:
-            memory: "512Mi"
+            memory: "4Gi"
         {%- endif %}
       imagePullSecrets:
       - name: {{ cluster_cfg["cluster"]["docker-registry"]["secret-name"] }}

src/rest-server/package.json

@@ -40,6 +40,7 @@
     "eslint": "~4.18.2",
     "eslint-config-google": "~0.9.1",
     "express": "~4.16.2",
+    "express-rate-limit": "^5.1.1",
     "fs-extra": "~7.0.1",
     "http-errors": "~1.6.3",
     "joi": "~13.0.1",

src/rest-server/src/config/express.js

@@ -29,6 +29,7 @@ const swaggerUi = require('swagger-ui-express');
 const config = require('@pai/config');
 const logger = require('@pai/config/logger');
 const authnConfig = require('@pai/config/authn');
+const limiter = require('@pai/config/rate-limit');
 const querystring = require('querystring');
 const createError = require('@pai/utils/error');
 const routers = {
@@ -38,6 +39,7 @@ const routers = {
 
 const app = express();
 
+app.set('trust proxy', true);
 app.set('json spaces', config.env === 'development' ? 4 : 0);
 
 app.use(cors());
@@ -46,6 +48,7 @@ app.use(bodyParser.urlencoded({extended: true}));
 app.use(bodyParser.json());
 app.use(bodyParser.text({type: 'text/*'}));
 app.use(cookieParser());
+app.use(limiter.api);
 
 // setup the logger for requests
 app.use(morgan('dev', {'stream': logger.stream}));

src/rest-server/src/config/rate-limit.js

@@ -0,0 +1,63 @@
+// Copyright (c) Microsoft Corporation
+// All rights reserved.
+//
+// MIT License
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+// documentation files (the "Software"), to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
+// to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+// The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
+// BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+// module dependencies
+const Joi = require('joi');
+const rateLimit = require('express-rate-limit');
+
+
+let limiterConfig = {
+  apiPerMin: process.env.RATE_LIMIT_API_PER_MIN,
+  listJobPerMin: process.env.RATE_LIMIT_LIST_JOB_PER_MIN,
+  submitJobPerHour: process.env.RATE_LIMIT_SUBMIT_JOB_PER_HOUR,
+};
+
+const limiterConfigSchema = Joi.object().keys({
+  apiPerMin: Joi.number()
+    .integer()
+    .default(600),
+  listJobPerMin: Joi.number()
+    .integer()
+    .default(60),
+  submitJobPerHour: Joi.number()
+    .integer()
+    .default(60),
+}).required();
+
+const {error, value} = Joi.validate(limiterConfig, limiterConfigSchema);
+if (error) {
+  throw new Error(`rate limit config error\n${error}`);
+}
+limiterConfig = value;
+
+// module exports
+module.exports = {
+  api: rateLimit({
+    max: limiterConfig.apiPerMin,
+    windowMs: 1 * 60 * 1000,
+  }),
+  listJob: rateLimit({
+    max: limiterConfig.listJobPerMin,
+    windowMs: 1 * 60 * 1000,
+    keyGenerator: (req) => req.user.username,
+  }),
+  submitJob: rateLimit({
+    max: limiterConfig.submitJobPerHour,
+    windowMs: 1 * 60 * 60 * 1000,
+    keyGenerator: (req) => req.user.username,
+  }),
+};

src/rest-server/src/routes/v2/job.js

@@ -18,6 +18,7 @@
 
 // module dependencies
 const express = require('express');
+const limiter = require('@pai/config/rate-limit');
 const token = require('@pai/middlewares/token');
 const controller = require('@pai/controllers/v2/job');
 const protocol = require('@pai/middlewares/v2/protocol');
@@ -28,10 +29,11 @@ const router = new express.Router();
 
 router.route('/')
   /** GET /api/v2/jobs - List job */
-  .get(token.check, controller.list)
+  .get(token.check, limiter.listJob, controller.list)
   /** POST /api/v2/jobs - Update job */
   .post(
     token.check,
+    limiter.submitJob,
     protocol.submit,
     controller.update
   );

src/rest-server/yarn.lock

@@ -1087,6 +1087,11 @@ [email protected]:
     signal-exit "^3.0.0"
     strip-eof "^1.0.0"
 
+express-rate-limit@^5.1.1:
+  version "5.1.1"
+  resolved "https://registry.yarnpkg.com/express-rate-limit/-/express-rate-limit-5.1.1.tgz#572e75c47ef890a6c9a3347f27bf3557d571f9ed"
+  integrity sha512-puA1zcCx/quwWUOU6pT6daCt6t7SweD9wKChKhb+KSgFMKRwS81C224hiSAUANw/gnSHiwEhgozM/2ezEBZPeA==
+
 express@~4.16.2:
   version "4.16.3"
   resolved "https://registry.npmjs.org/express/-/express-4.16.3.tgz#6af8a502350db3246ecc4becf6b5a34d22f7ed53"