[PATCH 3/6] tools/pygrub: Restrict depriv operation with RLIMIT_AS

[Alejandro Vallejo <alejandro.vallejo@xxxxxxxxx>]

Prevents the depriv pygrub from consuming more than a fixed amount of
memory.

Signed-off-by: Alejandro Vallejo <alejandro.vallejo@xxxxxxxxx>
---
 tools/pygrub/src/pygrub | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tools/pygrub/src/pygrub b/tools/pygrub/src/pygrub
index 327cf51774..b96bdfd849 100755
--- a/tools/pygrub/src/pygrub
+++ b/tools/pygrub/src/pygrub
@@ -39,6 +39,11 @@ SECTOR_SIZE = 512
 # pygrub
 LIMIT_FSIZE = 128 << 20
 
+# Unless provided through the env variable PYGRUB_MAX_RAM_USAGE_MB, then
+# this is the maximum amount of memory allowed to be used by the depriv
+# pygrub.
+LIMIT_AS = 2 * LIMIT_FSIZE
+
 CLONE_NEWNS = 0x00020000 # mount namespace
 CLONE_NEWNET = 0x40000000 # network namespace
 CLONE_NEWIPC = 0x08000000 # IPC namespace
@@ -75,6 +80,11 @@ def downgrade_rlimits():
     resource.setrlimit(resource.RLIMIT_CORE,     (0, 0))
     resource.setrlimit(resource.RLIMIT_MEMLOCK,  (0, 0))
 
+    max_ram_usage = LIMIT_AS
+    if "PYGRUB_MAX_RAM_USAGE_MB" in os.environ.keys():
+        max_ram_usage = int(os.environ["PYGRUB_MAX_RAM_USAGE_MB"]) << 20
+    resource.setrlimit(resource.RLIMIT_AS,  (max_ram_usage, max_ram_usage))
+
     # py2's resource module doesn't know about resource.RLIMIT_MSGQUEUE
     #
     # TODO: Use resource.RLIMIT_MSGQUEUE after python2 is deprecated
-- 
2.34.1