Fix remove mappings (#93)

Author: tyroguru

src/OID.cpp

@@ -323,16 +323,20 @@ static ExitStatus::ExitStatus runScript(const std::string &fileName,
    */
   if (oidConfig.attachToProcess) {
     if (oidConfig.removeMappings) {
+      ExitStatus::ExitStatus ret = ExitStatus::Success;
+
       if (!oid->segConfigExists()) {
         LOG(INFO) << "No config exists for pid " << oidConfig.pid
                   << " : cannot remove mappings";
+        ret = ExitStatus::UsageError;
       } else if (!oid->unmapSegments(true)) {
         LOG(ERROR) << "Failed to remove segments in target process with PID "
                    << oidConfig.pid;
-        return ExitStatus::SegmentRemovalError;
+        ret = ExitStatus::SegmentRemovalError;
       }
 
-      return ExitStatus::Success;
+      oid->contTargetThread();
+      return ret;
     }
 
     if (oidConfig.dataSegSize > 0) {
@@ -679,6 +683,14 @@ int main(int argc, char *argv[]) {
     return ExitStatus::UsageError;
   }
 
+  /*
+   * This is unfortunately necessary to stop users having to specify a script
+   * just to remove mappings (which doesn't make sense).
+   */
+  if (oidConfig.removeMappings && scriptFile.empty() && scriptSource.empty()) {
+    scriptSource = "entry:unknown_function:arg0";
+  }
+
   OICodeGen::Config codeGenConfig{
       .useDataSegment = true,
       .chaseRawPointers = chaseRawPointers,

test/integration.py

@@ -1,9 +1,11 @@
 import glob
 import json
+import os
 import os.path
 import shutil
 import subprocess
 import tempfile
+import time
 import unittest
 from contextlib import contextmanager
 from enum import Enum
@@ -237,6 +239,63 @@ def test_no_args_shows_help(self):
         self.expectReturncode(proc, ExitStatus.USAGE_ERROR)
         self.assertIn(b"usage: ", proc.stdout)
 
+    def test_remove_mappings(self):
+        with subprocess.Popen(
+            f"{self.binary_path} 100",
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            shell=True,
+        ) as target_proc:
+            pid = target_proc.pid
+
+            # The sleep is unfortunate but it allows the segments to get
+            # established in the target process.
+            time.sleep(2)
+
+            numsegs = subprocess.run(
+                f"cat /proc/{pid}/maps | wc -l",
+                shell=True,
+                capture_output=True,
+                check=True,
+            )
+            beforeSegs = int(numsegs.stdout.decode("ascii"))
+
+            proc = subprocess.run(
+                f"{self.oid} --script {self.script()} -t 1 --pid {pid}",
+                shell=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            self.expectReturncode(proc, ExitStatus.SUCCESS)
+
+            numsegs = subprocess.run(
+                f"cat /proc/{pid}/maps | wc -l",
+                shell=True,
+                capture_output=True,
+                check=True,
+            )
+            afterSegs = int(numsegs.stdout.decode("ascii"))
+            self.assertEqual(beforeSegs, afterSegs - 2)
+
+            # remove both the text and data segments
+            proc = subprocess.run(
+                f"{self.oid} -r --pid {pid}",
+                shell=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            self.expectReturncode(proc, ExitStatus.SUCCESS)
+
+            numsegs = subprocess.run(
+                f"cat /proc/{pid}/maps | wc -l",
+                shell=True,
+                capture_output=True,
+                check=True,
+            )
+            afterRemSegs = int(numsegs.stdout.decode("ascii"))
+            self.assertEqual(beforeSegs, afterRemSegs)
+
     def test_metrics_data_is_generated(self):
         with self.spawn_oid(self.script()) as proc:
             self.expectReturncode(proc, ExitStatus.SUCCESS)