Deflake test_linux_network_stacksmash_64

pyproject.toml

@@ -14,7 +14,6 @@ dependencies = [
     "povsim",
     "compilerex",
     "pwntools",
-    "flaky",
 ]
 
 [tool.setuptools.package-data]

rex/exploit/techniques/call_shellcode.py

@@ -4,10 +4,14 @@
 from .. import Exploit, CannotExploit
 from ..technique import Technique
 from ..nopsleds import NopSleds
-from ..actions import RexCommandAction
+from ..actions import RexCommandAction, RexWaitAction
 
 l = logging.getLogger("rex.exploit.techniques.call_shellcode")
 
+# Seconds to wait after sending the shellcode before sending any follow-up shell
+# commands, so the popped shell is reading stdin by the time the commands arrive.
+SHELL_SPAWN_DELAY = 2
+
 class CallShellcode(Technique):
 
     name = "call_shellcode"
@@ -64,6 +68,12 @@ def apply(self, cmd=None, use_nopsled=True, **kwargs): #pylint:disable=arguments
             if not cmd.endswith(b"\n"):
                 cmd += b"\n"
             channel_name = self.crash.input_type_to_channel(self.crash.input_type)
+            # Give the shellcode time to land and exec the shell before we send
+            # any commands. Otherwise the payload and the commands go out back to
+            # back, and a target that services them with a single read() can
+            # swallow both at once -- consuming the commands before the shell is
+            # up, which leaves the shell blocking forever on an empty stdin.
+            self.crash.actions.append(RexWaitAction(SHELL_SPAWN_DELAY))
             act = RexCommandAction(cmd, channel_name=channel_name)
             self.crash.actions.append(act)
             act = RexCommandAction(b"exit\n", channel_name=channel_name)

tests/test_rex.py

@@ -1,13 +1,14 @@
 # pylint: disable=line-too-long
 import os
-import random
+import socket
 import subprocess
 import sys
 import tempfile
 import time
 import struct
 import logging
 
+import psutil
 import pytest
 
 import archr
@@ -17,8 +18,6 @@
 from angr.state_plugins.trace_additions import FormatInfoStrToInt, FormatInfoDontConstrain
 from rex.exploit.cgc.type1.cgc_type1_shellcode_exploit import CGCType1ShellcodeExploit
 
-from flaky import flaky
-
 bin_location = str(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../../binaries'))
 cache_location = str(os.path.join(bin_location, 'tests_data/rop_gadgets_cache'))
 tests_dir = str(os.path.dirname(os.path.realpath(__file__)))
@@ -235,14 +234,42 @@ def test_linux_stacksmash_32():
         _check_arsenal_has_send(exploit.arsenal)
 
 
-@flaky(max_runs=3, min_passes=1)
+def _get_free_tcp_port():
+    """Return a TCP port that is currently free."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def _is_listening_on(port):
+    """Return True if any process is listening on TCP `port`."""
+    return any(
+        conn.status == psutil.CONN_LISTEN and conn.laddr and conn.laddr.port == port
+        for conn in psutil.net_connections(kind="tcp")
+    )
+
+
+def _wait_until_listening(port, timeout=30):
+    """Block until a process is listening on TCP `port`.
+
+    We check the socket state rather than connecting, because the target
+    accept()s exactly one connection -- a probe connection would be consumed
+    instead of the exploit's.
+    """
+    deadline = time.time() + timeout
+    while not _is_listening_on(port):
+        if time.time() > deadline:
+            raise TimeoutError(f"target never started listening on port {port}")
+        time.sleep(0.05)
+
+
 def test_linux_network_stacksmash_64():
     # Test exploiting a simple network server with a stack-based buffer overflow.
     inp = b'\x00' * 500
     lib_path = os.path.join(bin_location, "tests/x86_64")
     # ld_path = os.path.join(lib_path, "ld-linux-x86-64.so.2")
     path = os.path.join(lib_path, "network_overflow")
-    port = random.randint(8000, 9000)
+    port = _get_free_tcp_port()
     with archr.targets.LocalTarget([path, str(port)], path,
                                    target_arch='x86_64',
                                    ipv4_address="127.0.0.1",
@@ -258,7 +285,7 @@ def test_linux_network_stacksmash_64():
 
         # let's actually run the exploit
 
-    new_port = random.randint(9001, 10000)
+    new_port = _get_free_tcp_port()
     with archr.targets.LocalTarget([path, str(new_port)],
                                    path,
                                    target_arch='x86_64',
@@ -267,8 +294,9 @@ def test_linux_network_stacksmash_64():
         try:
             new_target.run_command("")
 
-            # wait for the target to load
-            time.sleep(.5)
+            # wait for the target to actually be listening before we launch the
+            # exploit (the generated script connects once, with no retry)
+            _wait_until_listening(new_port)
 
             temp_script = tempfile.NamedTemporaryFile(suffix=".py", delete=False)
             exploit_location = temp_script.name
@@ -278,7 +306,7 @@ def test_linux_network_stacksmash_64():
 
             exploit_result = subprocess.check_output(["python", exploit_location,
                                                       "127.0.0.1", str(new_port),
-                                                      ], timeout=3)
+                                                      ], timeout=30)
             assert b"hello" in exploit_result
         finally:
             os.unlink(exploit_location)