diff --git a/src/limits/MemoryLimitListener.cc b/src/limits/MemoryLimitListener.cc index feb5bcc..9f018cb 100644 --- a/src/limits/MemoryLimitListener.cc +++ b/src/limits/MemoryLimitListener.cc @@ -40,33 +40,52 @@ 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(); + if (!vmPeakValid_) { + return tracer::TraceAction::CONTINUE; + } + bool doesUnmap = (tracee.getSyscallArgument(3)&MREMAP_DONTUNMAP)!=0; + // Allow user to shrink its memory + if(!doesUnmap&&tracee.getSyscallArgument(1)>=tracee.getSyscallArgument(2)){ + return tracer::TraceAction::CONTINUE; + } + uint64_t newMemoryAllocated=tracee.getSyscallArgument(2); + // Do not count already allocated memory + if(!doesUnmap) newMemoryAllocated-=tracee.getSyscallArgument(1); + 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..9fc92a5 100644 --- a/src/limits/MemoryLimitListener.h +++ b/src/limits/MemoryLimitListener.h @@ -5,6 +5,13 @@ #include "seccomp/policy/SyscallPolicy.h" #include "tracer/TraceEventListener.h" +#include + +#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 namespace s2j { @@ -40,6 +47,7 @@ class MemoryLimitListener pid_t childPid_; std::vector syscallRules_; + tracer::TraceAction handleMemoryAllocation(uint64_t allocatedMemoryKb); }; } // namespace limits