From 2690d897f9b135763ad6f2a6132972b4a22f17c9 Mon Sep 17 00:00:00 2001
From: Defelo <elodef42@gmail.com>
Date: Wed, 28 Apr 2021 16:01:07 +0200
Subject: [PATCH] Added max_memory_usage parameter

---
 api/src/api/v2.js | 21 +++++++++++++++++++--
 api/src/config.js |  6 ++++++
 api/src/job.js    | 12 ++++++++----
 3 files changed, 33 insertions(+), 6 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;