Fix/timeout

[Kyle-Kyle]

[astewart-bah]

While testing the new timeout code, we ran into this issue:

angrop was attempting to write a log message to stdout and it was specifically in this log function:

https://github.com/python/cpython/blob/main/Lib/logging/__init__.py#L1154

The logging function has an except Exception block that catches the RopException and ignores it.

This caused the analysis to continue until the system ran out of RAM and the OOM Killer killed python.

This code block can demo the behavior:

import signal
import time


class RopException(Exception):
    pass


def timeout(seconds_before_timeout):
    def decorate(f):
        def handler(signum, frame):# pylint:disable=unused-argument
            print("[angrop] Timeout")
            raise RopException("[angrop] Timeout!")
        def new_f(*args, **kwargs):
            old = signal.signal(signal.SIGALRM, handler)
            old_time_left = signal.alarm(seconds_before_timeout)
            if 0 < old_time_left < seconds_before_timeout: # never lengthen existing timer
                signal.alarm(old_time_left)
            start_time = time.time()
            try:
                result = f(*args, **kwargs)
            finally:
                if old_time_left > 0: # deduct f's run time from the saved timer
                    old_time_left -= int(time.time() - start_time)
                signal.signal(signal.SIGALRM, old)
                signal.alarm(old_time_left)
            return result
        return new_f
    return decorate


@timeout(2)
def busy_function():
    run_count = 0
    while run_count < 5:
        run_count += 1
        print(f"Running {run_count}/5")
        try:
            time.sleep(1)
        except Exception as ex:
            print(f"Caught an exception of type {type(ex).__name__}")


if __name__ == "__main__":
    start_time = time.time()
    busy_function()
    end_time = time.time()
    print(f"The timeout-decorated function took {end_time - start_time} seconds to execute.")

[astewart-bah]

Another issue we found:

While testing the new timeout code, we continued seeing runaway timeout scenarios. This new case was due to weakref.

Here's an except from https://docs.python.org/3/library/weakref.html#weakref.ref

Exceptions raised by the callback will be noted on the standard error output, but cannot be propagated; they are handled in exactly the same way as exceptions raised from an object’s __del__() method.

I've included some example code below. It's not 100% reliable in reproducing the bug, maybe even as low as 10% reliable. Just keep running it until it print out a giant message to the screen letting you know you should check the stack trace.

We've also included a potential fix for this in the timeout() function included in the file. The code is commented out, but if you uncomment it and run the script until the error reproduces, you'll see that the timeout is able to handle the error this time.

import weakref
import signal
import time
import sys

class RopException(Exception):
    pass


def timeout(seconds_before_timeout):
    def decorate(f):
        def handler(signum, frame):# pylint:disable=unused-argument
            while frame.f_back:
                filename = frame.f_code.co_filename
                if frame.f_code.co_name == '__del__':
                    print("Delaying for 0.1s because of __del__")
                    signal.setitimer(signal.ITIMER_REAL, 0.1, 0)
                    return
                # if filename.startswith(sys.base_prefix) and filename.endswith("/weakref.py"):
                #     print("Delaying for 0.1s because of weakref.py")
                #     signal.setitimer(signal.ITIMER_REAL, 0.1, 0)
                #     return
                frame = frame.f_back
            print("[angrop] Timeout")
            raise RopException("[angrop] Timeout!")
        def new_f(*args, **kwargs):
            old = signal.signal(signal.SIGALRM, handler)
            old_time_left = signal.alarm(seconds_before_timeout)
            if 0 < old_time_left < seconds_before_timeout: # never lengthen existing timer
                signal.alarm(old_time_left)
            start_time = time.time()
            try:
                result = f(*args, **kwargs)
            finally:
                if old_time_left > 0: # deduct f's run time from the saved timer
                    old_time_left -= int(time.time() - start_time)
                signal.signal(signal.SIGALRM, old)
                signal.alarm(old_time_left)
            return result
        return new_f
    return decorate


iter_count = 1000000
_id2obj_dict = weakref.WeakValueDictionary()
def remember(obj):
    oid = id(obj)
    _id2obj_dict[oid] = obj
    return oid
def id2obj(oid):
    return _id2obj_dict[oid]


class UnimportantClass:
    def __init__(self, num):
        self.num = num


def profile():
    global _id2obj_dict
    global iter_count
    _id2obj_dict = weakref.WeakValueDictionary()
    ids = []
    print("Profiling...")
    start_time = time.time()
    for i in range(iter_count):
        uc = UnimportantClass(i)
        ids.append(remember(uc))
    for id in ids:
        try:
            ref = id2obj(id)
            del ref
        except KeyError:
            pass
    end_time = time.time()
    duration = end_time - start_time
    new_iter_count = int(iter_count * (0.75 / duration))
    print(f"Profiling results: {iter_count} iterations took {duration} seconds. Adjusting iter_count to {new_iter_count}")
    iter_count = new_iter_count


@timeout(1)
def busy_function_2():
    global _id2obj_dict
    run_count = 0
    while run_count < 5:
        run_count += 1
        print(f"Running {run_count}/5")
        _id2obj_dict = weakref.WeakValueDictionary()
        ids = []
        for i in range(iter_count):
            uc = UnimportantClass(i)
            ids.append(remember(uc))
        for id in ids:
            try:
                ref = id2obj(id)
                del ref
            except KeyError:
                pass
    # We only hit this is timeout failed
    print("*********************************************************")
    print("*********************************************************")
    print("!!!!!!!!!!!!!!! CHECK OUT THE STACK TRACE !!!!!!!!!!!!!!!")
    print("*********************************************************")
    print("*********************************************************")


if __name__ == "__main__":
    profile()
    start_time = time.time()
    busy_function_2()
    end_time = time.time()
    print(f"The timeout-decorated function took {end_time - start_time} seconds to execute.")