From a49acc7dbe4d63090250d6a1d1536542e362aef3 Mon Sep 17 00:00:00 2001
From: Felix <41747605+Defelo@users.noreply.github.com>
Date: Sat, 1 May 2021 11:10:53 +0200
Subject: [PATCH] Added max_memory_usage parameter (#238)

* Added max_memory_usage parameter

* Added max_memory_usage description to readme
---
 api/src/api/v2.js | 21 +++++++++++++++++++--
 api/src/config.js |  6 ++++++
 api/src/job.js    | 12 ++++++++----
 readme.md         |  4 +++-
 4 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/api/src/api/v2.js b/api/src/api/v2.js
index 72db74e..4d8713a 100644
--- a/api/src/api/v2.js
+++ b/api/src/api/v2.js
@@ -1,13 +1,14 @@
 const express = require('express');
 const router  = express.Router();
 
+const config = require('../config');
 const runtime = require('../runtime');
 const {Job} = require("../job");
 const package = require('../package')
 const logger = require('logplease').create('api/v1');
 
 router.post('/execute', async function(req, res){
-        const {language, version, files, stdin, args, run_timeout, compile_timeout} = req.body;
+        const {language, version, files, stdin, args, run_timeout, compile_timeout, max_memory_usage} = req.body;
 
         if(!language || typeof language !== "string")
         {
@@ -46,6 +47,21 @@ router.post('/execute', async function(req, res){
             }
         }
 
+        if (max_memory_usage) {
+            if (typeof max_memory_usage !== "number" || max_memory_usage < 0) {
+                return res
+                    .status(400)
+                    .send({
+                        message: "if specified, max_memory_usage must be a non-negative number"
+                    })
+            } else if (max_memory_usage > config.max_memory_usage) {
+                return res
+                    .status(400)
+                    .send({
+                        message: "max_memory_usage cannot exceed the configured limit of " + config.max_memory_usage
+                    })
+            }
+        }
 
 
     
@@ -68,7 +84,8 @@ router.post('/execute', async function(req, res){
             timeouts: {
                 run: run_timeout || 3000,
                 compile: compile_timeout || 10000
-            }
+            },
+            max_memory_usage: max_memory_usage || config.max_memory_usage
         });
 
         await job.prime();
diff --git a/api/src/config.js b/api/src/config.js
index c97b64c..12da1c6 100644
--- a/api/src/config.js
+++ b/api/src/config.js
@@ -108,6 +108,12 @@ const options = [
         default: 1000000, //1MB
         validators: []
     },
+    {
+        key: 'max_memory_usage',
+        desc: 'Max memory usage in bytes',
+        default: 256000000, //256MB
+        validators: []
+    },
     {
         key: 'repo_url',
         desc: 'URL of repo index',
diff --git a/api/src/job.js b/api/src/job.js
index b711ac1..258300a 100644
--- a/api/src/job.js
+++ b/api/src/job.js
@@ -19,7 +19,7 @@ let gid = 0;
 
 class Job {
 
-    constructor({ runtime, files, args, stdin, timeouts }) {
+    constructor({ runtime, files, args, stdin, timeouts, max_memory_usage }) {
         this.uuid =  uuidv4();
         this.runtime = runtime;
         this.files = files.map((file,i) => ({
@@ -30,6 +30,7 @@ class Job {
         this.args = args;
         this.stdin = stdin;
         this.timeouts = timeouts;
+        this.max_memory_usage = max_memory_usage;
 
         this.uid = config.runner_uid_min + uid;
         this.gid = config.runner_gid_min + gid;
@@ -75,7 +76,8 @@ class Job {
                 'prlimit',
                 '--nproc=' + config.max_process_count,
                 '--nofile=' + config.max_open_files,
-                '--fsize=' + config.max_file_size
+                '--fsize=' + config.max_file_size,
+                '--as=' + this.max_memory_usage
             ];
 
             const proc_call = [
@@ -161,7 +163,8 @@ class Job {
             compile = await this.safe_call(
                 path.join(this.runtime.pkgdir, 'compile'),
                 this.files.map(x => x.name),
-                this.timeouts.compile
+                this.timeouts.compile,
+                config.max_memory_usage
             );
         }
 
@@ -170,7 +173,8 @@ class Job {
         const run = await this.safe_call(
             path.join(this.runtime.pkgdir, 'run'),
             [this.files[0].name, ...this.args],
-            this.timeouts.run
+            this.timeouts.run,
+            this.max_memory_usage
         );
 
         this.state = job_states.EXECUTED;
diff --git a/readme.md b/readme.md
index f1d0f2e..21d152d 100644
--- a/readme.md
+++ b/readme.md
@@ -210,6 +210,7 @@ This endpoint requests execution of some arbitrary code.
 - `args` (*optional*) The arguments to pass to the program. Must be an array or left out. Defaults to `[]`.
 - `compile_timeout` (*optional*) The maximum time allowed for the compile stage to finish before bailing out in milliseconds. Must be a number or left out. Defaults to `10000` (10 seconds).
 - `run_timeout` (*optional*) The maximum time allowed for the run stage to finish before bailing out in milliseconds. Must be a number or left out. Defaults to `3000` (3 seconds).
+- `max_memory_usage` (*optional*) The maximum amount of memory the run stage is allowed to use. Must be a number or left out. Defaults to `256000000` (256 MB)
 
 ```json
 {
@@ -228,7 +229,8 @@ This endpoint requests execution of some arbitrary code.
         "3"
     ],
     "compile_timeout": 10000,
-    "run_timeout": 3000
+    "run_timeout": 3000,
+    "max_memory_usage": 256000000
 }
 ```
 A typical response upon successful execution will contain 1 or 2 keys `run` and `compile`.