[xen master] tools/pygrub: Restrict depriv operation with RLIMIT_AS

[patchbot@xxxxxxx]

commit 03d6720a4c62c283f9a9f09858eeccd24299b312
Author:     Alejandro Vallejo <alejandro.vallejo@xxxxxxxxx>
AuthorDate: Mon Nov 6 15:05:05 2023 +0000
Commit:     Andrew Cooper <andrew.cooper3@xxxxxxxxxx>
CommitDate: Thu Nov 23 17:21:01 2023 +0000

    tools/pygrub: Restrict depriv operation with RLIMIT_AS
    
    Prevents the depriv pygrub from consuming more than a fixed amount of
    memory.
    
    Signed-off-by: Alejandro Vallejo <alejandro.vallejo@xxxxxxxxx>
    Acked-by: Andrew Cooper <andrew.cooper3@xxxxxxxxxx>
---
 tools/pygrub/src/pygrub | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tools/pygrub/src/pygrub b/tools/pygrub/src/pygrub
index 2c06684d65..9d51f96070 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:
+        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
--
generated by git-patchbot for /home/xen/git/xen.git#master