From 0f457bb625be6b337225f933eece54ae9c243bb4 Mon Sep 17 00:00:00 2001
From: pythontest <mmm05892@gmail.com>
Date: Tue, 6 Feb 2024 21:37:05 +0100
Subject: [PATCH 1/2] Added "Memory limit exceeded" error checking when
 allocatiing memory using mremap

---
 src/limits/MemoryLimitListener.cc | 75 ++++++++++++++++++++++---------
 src/limits/MemoryLimitListener.h  |  1 +
 2 files changed, 54 insertions(+), 22 deletions(-)

diff --git a/src/limits/MemoryLimitListener.cc b/src/limits/MemoryLimitListener.cc
index feb5bcc..8f99b92 100644
--- a/src/limits/MemoryLimitListener.cc
+++ b/src/limits/MemoryLimitListener.cc
@@ -8,9 +8,16 @@
 #include "seccomp/action/ActionTrace.h"
 #include "seccomp/filter/LibSeccompFilter.h"
 
+#include <sys/mman.h>
 #include <sys/resource.h>
 #include <sys/time.h>
 
+#ifndef MREMAP_DONTUNMAP
+// Due to this flag being introduced in Linux 5.7 there are many system which do
+// not define it
+#define MREMAP_DONTUNMAP 4
+#endif
+
 #include <csignal>
 #include <cstdint>
 #include <fstream>
@@ -40,33 +47,57 @@ MemoryLimitListener::MemoryLimitListener(uint64_t memoryLimitKb)
                     if (!vmPeakValid_) {
                         return tracer::TraceAction::CONTINUE;
                     }
-
-                    uint64_t memoryUsage = getMemoryUsageKb() +
-                                           tracee.getSyscallArgument(1) / 1024;
-                    memoryPeakKb_ = std::max(memoryPeakKb_, memoryUsage);
-                    outputBuilder_->setMemoryPeak(memoryPeakKb_);
-                    logger::debug(
-                            "Memory usage after mmap ",
-                            VAR(memoryUsage),
-                            ", ",
-                            VAR(memoryPeakKb_));
-
-                    if (memoryUsage > memoryLimitKb_) {
-                        outputBuilder_->setKillReason(
-                                printer::OutputBuilder::KillReason::MLE,
-                                "memory limit exceeded");
-                        logger::debug(
-                                "Limit ",
-                                VAR(memoryLimitKb_),
-                                " exceeded, killing tracee");
-                        return tracer::TraceAction::KILL;
-                    }
-                    return tracer::TraceAction::CONTINUE;
+                    return handleMemoryAllocation(
+                            tracee.getSyscallArgument(1) / 1024);
                 }),
                 Arg(0) == 0 && Arg(1) > MEMORY_LIMIT_MARGIN / 2));
     }
+    syscallRules_.emplace_back(seccomp::SeccompRule(
+            "mremap",
+            seccomp::action::ActionTrace([this](tracer::Tracee& tracee) {
+                TRACE();
+                uint64_t old_size = tracee.getSyscallArgument(1);
+                uint64_t new_size = tracee.getSyscallArgument(2);
+                uint64_t flags = tracee.getSyscallArgument(3);
+                if (!vmPeakValid_) {
+                    return tracer::TraceAction::CONTINUE;
+                }
+                bool doesntUnmap =
+                        (flags & MREMAP_DONTUNMAP) == MREMAP_DONTUNMAP;
+                // Allow user to shrink its memory
+                if (!doesntUnmap && old_size >= new_size) {
+                    return tracer::TraceAction::CONTINUE;
+                }
+                uint64_t newMemoryAllocated = new_size;
+                // Do not count already allocated memory
+                if (!doesntUnmap)
+                    newMemoryAllocated -= old_size;
+                newMemoryAllocated /= 1024;
+                return handleMemoryAllocation(newMemoryAllocated);
+            }),
+            Arg(2) > MEMORY_LIMIT_MARGIN / 2));
 }
+tracer::TraceAction MemoryLimitListener::handleMemoryAllocation(
+        uint64_t allocatedMemoryKb) {
+    uint64_t memoryUsage = getMemoryUsageKb() + allocatedMemoryKb;
+    memoryPeakKb_ = std::max(memoryPeakKb_, memoryUsage);
+    outputBuilder_->setMemoryPeak(memoryPeakKb_);
+    logger::debug(
+            "Memory usage after allocation ",
+            VAR(memoryUsage),
+            ", ",
+            VAR(memoryPeakKb_));
 
+    if (memoryUsage > memoryLimitKb_) {
+        outputBuilder_->setKillReason(
+                printer::OutputBuilder::KillReason::MLE,
+                "memory limit exceeded");
+        logger::debug(
+                "Limit ", VAR(memoryLimitKb_), " exceeded, killing tracee");
+        return tracer::TraceAction::KILL;
+    }
+    return tracer::TraceAction::CONTINUE;
+}
 void MemoryLimitListener::onPostForkChild() {
     TRACE();
 
diff --git a/src/limits/MemoryLimitListener.h b/src/limits/MemoryLimitListener.h
index 7430fa1..fb7bf50 100644
--- a/src/limits/MemoryLimitListener.h
+++ b/src/limits/MemoryLimitListener.h
@@ -40,6 +40,7 @@ class MemoryLimitListener
     pid_t childPid_;
 
     std::vector<seccomp::SeccompRule> syscallRules_;
+    tracer::TraceAction handleMemoryAllocation(uint64_t allocatedMemoryKb);
 };
 
 } // namespace limits

From 65c5a25ff5fc9e5f95d5f956caf4cfc5e8d4560a Mon Sep 17 00:00:00 2001
From: pythontest <mmm05892@gmail.com>
Date: Wed, 7 Feb 2024 10:55:57 +0100
Subject: [PATCH 2/2] Fixed a bug that caused a Memory Limit Exceeded error
 when trying to allocate more than MEMORY_LIMIT_MARGIN/2 when there was no
 memory limit

---
 src/limits/MemoryLimitListener.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/limits/MemoryLimitListener.cc b/src/limits/MemoryLimitListener.cc
index 8f99b92..06f3b0a 100644
--- a/src/limits/MemoryLimitListener.cc
+++ b/src/limits/MemoryLimitListener.cc
@@ -88,7 +88,7 @@ tracer::TraceAction MemoryLimitListener::handleMemoryAllocation(
             ", ",
             VAR(memoryPeakKb_));
 
-    if (memoryUsage > memoryLimitKb_) {
+    if (memoryLimitKb_ > 0 && memoryUsage > memoryLimitKb_) {
         outputBuilder_->setKillReason(
                 printer::OutputBuilder::KillReason::MLE,
                 "memory limit exceeded");