From 76357216da6658842f952937798b0b71d360de40 Mon Sep 17 00:00:00 2001 From: feliam Date: Thu, 27 Jul 2017 19:41:08 -0300 Subject: [PATCH] Pub-sub like events (#371) * WIP New Policy class * WIP pubsub * Update Signal tests * small fixes from github comments * Fix event decode_instruction signature * Good merge * Good good merge * Eventful class commented * The million typos * Code revision * Fix tests for new mor strict Eventful --- manticore/core/cpu/abstractcpu.py | 52 +++---- manticore/core/executor.py | 147 ++++++++++++-------- manticore/core/state.py | 12 +- manticore/manticore.py | 37 +++-- manticore/platforms/decree.py | 15 ++- manticore/platforms/linux.py | 14 +- manticore/platforms/platform.py | 18 ++- manticore/platforms/windows.py | 11 +- manticore/utils/event.py | 216 +++++++++--------------------- tests/test_binaries.py | 1 + tests/test_events.py | 77 +++++++++++ tests/test_signal.py | 70 ---------- tests/test_state.py | 17 +-- tests/test_workspace.py | 15 +-- 14 files changed, 312 insertions(+), 390 deletions(-) create mode 100644 tests/test_events.py delete mode 100644 tests/test_signal.py diff --git a/manticore/core/cpu/abstractcpu.py b/manticore/core/cpu/abstractcpu.py index 3ff3ff3..6920c44 100644 --- a/manticore/core/cpu/abstractcpu.py +++ b/manticore/core/cpu/abstractcpu.py @@ -8,7 +8,7 @@ from ...utils.emulate import UnicornEmulator import sys from functools import wraps from itertools import islice, imap -from ...utils.event import Signal +from ...utils.event import Eventful import inspect import types import logging @@ -345,7 +345,7 @@ class SyscallAbi(Abi): ############################################################################ # Abstract cpu encapsulating common cpu methods used by platforms and executor. -class Cpu(object): +class Cpu(Eventful): ''' Base class for all Cpu architectures. Functionality common to all architectures (and expected from users of a Cpu) should be here. Commonly @@ -361,9 +361,9 @@ class Cpu(object): - stack_alias ''' - def __init__(self, regfile, memory): + def __init__(self, regfile, memory, **kwargs): assert isinstance(regfile, RegisterFile) - super(Cpu, self).__init__() + super(Cpu, self).__init__(**kwargs) self._regfile = regfile self._memory = memory self._instruction_cache = {} @@ -374,24 +374,6 @@ class Cpu(object): self._md.detail = True self._md.syntax = 0 - ##################################################### - # Signals - # signal handlers must have this signature: - # handler(cpu, *args, **kwargs) - self.will_decode_instruction = Signal() - self.will_execute_instruction = Signal() - self.did_execute_instruction = Signal() - self.will_emulate_instruction = Signal() - self.did_emulate_instruction = Signal() - self.will_read_register = Signal() - self.did_read_register = Signal() - self.will_write_register = Signal() - self.did_write_register = Signal() - self.will_read_memory = Signal() - self.will_write_memory = Signal() - self.did_read_memory = Signal() - self.did_write_memory = Signal() - # Ensure that regfile created STACK/PC aliases assert 'STACK' in self._regfile assert 'PC' in self._regfile @@ -452,9 +434,9 @@ class Cpu(object): :param value: register value :type value: int or long or Expression ''' - self.will_write_register(register, value) + self.publish('will_write_register', register, value) value = self._regfile.write(register, value) - self.did_write_register(register, value) + self.publish('did_write_register', register, value) return value def read_register(self, register): @@ -465,9 +447,9 @@ class Cpu(object): :return: register value :rtype: int or long or Expression ''' - self.will_read_register(register) + self.publish('will_read_register', register) value = self._regfile.read(register) - self.did_read_register(register, value) + self.publish('did_read_register', register, value) return value # Pythonic access to registers and aliases @@ -513,11 +495,11 @@ class Cpu(object): if size is None: size = self.address_bit_size assert size in SANE_SIZES - self.will_write_memory(where, expression, size) + self.publish('will_write_memory', where, expression, size) self.memory[where:where+size/8] = [Operators.CHR(Operators.EXTRACT(expression, offset, 8)) for offset in xrange(0, size, 8)] - self.did_write_memory(where, expression, size) + self.publish('did_write_memory', where, expression, size) def read_int(self, where, size=None): @@ -532,13 +514,13 @@ class Cpu(object): if size is None: size = self.address_bit_size assert size in SANE_SIZES - self.will_read_memory(where, size) + self.publish('will_read_memory', where, size) data = self.memory[where:where+size/8] assert (8 * len(data)) == size value = Operators.CONCAT(size, *map(Operators.ORD, reversed(data))) - self.did_read_memory(where, value, size) + self.publish('did_read_memory', where, value, size) return value @@ -727,12 +709,12 @@ class Cpu(object): if not self.memory.access_ok(self.PC,'x'): raise InvalidMemoryAccess(self.PC, 'x') - self.will_decode_instruction() + self.publish('will_decode_instruction', self.PC) instruction = self.decode_instruction(self.PC) self._last_pc=self.PC - self.will_execute_instruction(instruction) + self.publish('will_execute_instruction', instruction) if instruction.address != self.PC: return @@ -745,11 +727,11 @@ class Cpu(object): instruction.address, text_bytes, instruction.mnemonic, instruction.op_str) - self.will_emulate_instruction(instruction) + self.publish('will_emulate_instruction', instruction) self.emulate(instruction) - self.did_emulate_instruction(instruction) + self.publish('did_emulate_instruction', instruction) implementation = getattr(self, name, fallback_to_emulate) @@ -761,7 +743,7 @@ class Cpu(object): implementation(*instruction.operands) self._icount+=1 - self.did_execute_instruction(instruction) + self.publish('did_execute_instruction', instruction) def emulate(self, instruction): diff --git a/manticore/core/executor.py b/manticore/core/executor.py index 9ffb4d0..56dd335 100644 --- a/manticore/core/executor.py +++ b/manticore/core/executor.py @@ -10,7 +10,7 @@ except: import StringIO from ..utils.nointerrupt import WithKeyboardInterruptAs -from ..utils.event import Signal, forward_signals +from ..utils.event import Eventful from .smtlib import solver, Expression, SolverException from .state import Concretize, TerminateState from workspace import Workspace @@ -39,51 +39,88 @@ def sync(f): class Policy(object): ''' Base class for prioritization of state search ''' - def __init__(self): - pass + def __init__(self, executor, *args, **kwargs): + super(Policy, self).__init__(*args, **kwargs) + self._executor = executor + self._executor.subscribe('did_add_state', self._add_state_callback) + + @contextmanager + def locked_context(self): + ''' Policy shared context dictionary ''' + with self._executor.locked_context() as ctx: + policy_context = ctx.get('policy', None) + if policy_context is None: + policy_context = dict() + yield policy_context + ctx['policy'] = policy_context - def features(self, state): - ''' Save state features for prioritization before a state is stored ''' - pass + def _add_state_callback(self, state_id, state): + ''' Save prepare(state) on policy shared context before + the state is stored + ''' + with self.locked_context() as ctx: + metric = self.prepare(state) + if metric is not None: + ctx[state_id] = metric - def priority(self, state_id): - ''' A numeric value representing likelihood to reach the interesting program spot ''' - return 1.0 + def prepare(self, state): + ''' Process a state and keep enough data to later decide it's + priority #fixme rephrase + ''' + return None + + def choice(self, state_ids): + ''' Select a state id from states_id. + self.context has a dict mapping state_ids -> prepare(state)''' + raise NotImplementedError class Random(Policy): - def __init__(self): - super(Random, self).__init__() + def __init__(self, executor, *args, **kwargs): + super(Random, self).__init__(executor, *args, **kwargs) - def features(self, state): - ''' Save state features for prioritization before a state is stored ''' - pass + def choice(self, state_ids): + return random.choice(state_ids) - def priority(self, state_id): - ''' A numeric value representing likelihood to reach the interesting program spot ''' - return 1.0 +class Uncovered(Policy): + def __init__(self, executor, *args, **kwargs): + super(Uncovered, self).__init__(executor, *args, **kwargs) + #hook on the necesary executor signals + #on callbacks save data in executor.context['policy'] + + def prepare(self, state): + ''' this is what we need to save for choosing later ''' + return state.cpu.PC + + def choice(self, state_ids): + # Use executor.context['uncovered'] = state_id -> stats + # am + with self._executor.locked_context() as ctx: + lastpc = ctx['policy'] + visited = ctx.get('visited', ()) + interesting = set() + for _id in state_ids: + if lastpc.get(_id, None) not in visited: + interesting.add(_id) + if len(interesting) > 0: + return random.choice(tuple(interesting)) + else: + return random.choice(state_ids) -class Executor(object): +class Executor(Eventful): ''' The executor guides the execution of an initial state or a paused previous run. It handles all exceptional conditions (system calls, memory faults, concretization, etc.) ''' - def __init__(self, initial=None, workspace='', policy='random', context=None, **options): - # Signals / Callbacks handlers will be invoked potentially at different + def __init__(self, initial=None, workspace=None, policy='random', context=None, **kwargs): + super(Executor, self).__init__(**kwargs) + + + # Signals / Callbacks handlers will be invoked potentially at different # worker processes. State provides a local context to save data. - #Executor signals - self.will_start_run = Signal() - self.will_finish_run = Signal() - self.will_fork_state = Signal() - self.will_store_state = Signal() - self.will_load_state = Signal() - self.will_terminate_state = Signal() - self.will_generate_testcase = Signal() - - #Be sure every state will forward us their signals - self.will_load_state += self._register_state_callbacks + self.subscribe('will_load_state', self._register_state_callbacks) #The main executor lock. Acquire this for accessing shared objects self._lock = manager.Condition(manager.RLock()) @@ -105,21 +142,21 @@ class Executor(object): self._shared_context = manager.dict(context) #scheduling priority policy (wip) - self.policy = Random() + #Set policy + policies = {'random': Random, + 'uncovered': Uncovered + } + self._policy = policies[policy](self) + assert isinstance(self._policy, Policy) + if self.load_workspace(): if initial is not None: logger.error("Ignoring initial state") - # We loaded state ids, now load the actual state - - current_state_id = self.get() - initial = self._workspace.load_state(current_state_id) - self._register_state_callbacks(initial, current_state_id) - - self.add(initial) - ##FIXME PUBSUB We need to forward signals here so they get declared - ##forward signals from initial state so they are declared here - self._register_state_callbacks(initial, 0) # id param unused + else: + if initial is not None: + self.add(initial) + self.forward_events_from(initial, True) @contextmanager def locked_context(self): @@ -142,7 +179,7 @@ class Executor(object): Going up, we prepend state in the arguments. ''' #Forward all state signals - forward_signals(self, state, True) + self.forward_events_from(state, True) def add(self, state): ''' @@ -152,8 +189,8 @@ class Executor(object): ''' #save the state to secondary storage state_id = self._workspace.save_state(state) - self.will_store_state(state, state_id) self.put(state_id) + self.publish('did_add_state', state_id, state) return state_id def load_workspace(self): @@ -228,13 +265,10 @@ class Executor(object): logger.debug("Waiting for available states") self._lock.wait() - state_id = random.choice(self._states) + state_id = self._policy.choice(list(self._states)) del self._states[self._states.index(state_id)] return state_id - ############################################################### - # File Storage - def list(self): ''' Returns the list of states ids currently queued ''' return list(self._states) @@ -250,7 +284,8 @@ class Executor(object): #broadcast test generation. This is the time for other modules #to output whatever helps to understand this testcase - self.will_generate_testcase(state, message) + self.publish('will_generate_testcase', state) + def fork(self, state, expression, policy='ALL', setstate=None): @@ -282,7 +317,7 @@ class Executor(object): #We are about to fork current_state with self._lock: - self.will_fork_state(state, expression, solutions, policy) + self.publish('will_fork_state', state, expression, solutions, policy) #Build and enqueue a state for each solution children = [] @@ -329,7 +364,7 @@ class Executor(object): #load selected state from secondary storage if current_state_id is not None: current_state = self._workspace.load_state(current_state_id) - self.will_load_state(current_state, current_state_id) + self.publish('will_load_state', current_state, current_state_id) #notify siblings we have a state to play with self._start_run() @@ -348,7 +383,7 @@ class Executor(object): break else: #Notify this worker is done - self.will_terminate_state(current_state, current_state_id, 'Shutdown') + self.publish('will_terminate_state', current_state, current_state_id, 'Shutdown') current_state = None @@ -364,7 +399,7 @@ class Executor(object): except TerminateState as e: #Notify this worker is done - self.will_terminate_state(current_state, current_state_id, e) + self.publish('will_terminate_state', current_state, current_state_id, e) logger.debug("Generic terminate state") if e.testcase: @@ -377,7 +412,7 @@ class Executor(object): logger.error("Exception: %s\n%s", str(e), trace) #Notify this state is done - self.will_terminate_state(current_state, current_state_id, e) + self.publish('will_terminate_state', current_state, current_state_id, e) if solver.check(current_state.constraints): self.generate_testcase(current_state, "Solver failed" + str(e)) @@ -388,7 +423,7 @@ class Executor(object): trace = traceback.format_exc() logger.error("Exception: %s\n%s", str(e), trace) #Notify this worker is done - self.will_terminate_state(current_state, current_state_id, 'Exception') + self.publish('will_terminate_state', current_state, current_state_id, 'Exception') current_state = None logger.setState(None) @@ -398,6 +433,6 @@ class Executor(object): self._stop_run() #Notify this worker is done (not sure it's needed) - self.will_finish_run() + self.publish('will_finish_run') diff --git a/manticore/core/state.py b/manticore/core/state.py index faed797..c656c28 100644 --- a/manticore/core/state.py +++ b/manticore/core/state.py @@ -5,8 +5,7 @@ from collections import OrderedDict from .smtlib import solver, Bool from ..utils.helpers import issymbolic -from ..utils.event import Signal, forward_signals - +from ..utils.event import Eventful #import exceptions from .cpu.abstractcpu import ConcretizeRegister @@ -59,10 +58,8 @@ class ForkState(Concretize): super(ForkState, self).__init__(message, expression, policy='ALL', **kwargs) -from ..utils.event import Signal +class State(Eventful): - -class State(object): ''' Representation of a unique program state/path. @@ -71,7 +68,8 @@ class State(object): :ivar dict context: Local context for arbitrary data storage ''' - def __init__(self, constraints, platform): + def __init__(self, constraints, platform, **kwargs): + super(State, self).__init__(**kwargs) self.platform = platform self.forks = 0 self.constraints = constraints @@ -86,7 +84,7 @@ class State(object): #self.will_add_constraint = Signal() #Import all signals from platform - forward_signals(self, platform) + self.forward_events_from(platform) def __reduce__(self): return (self.__class__, (self.constraints, self.platform), diff --git a/manticore/manticore.py b/manticore/manticore.py index 026bc0a..68a9577 100644 --- a/manticore/manticore.py +++ b/manticore/manticore.py @@ -573,7 +573,7 @@ class Manticore(object): logger.debug("About to store state %r %r %r", state, expression, values, policy) - def _read_register_callback(self, state, reg_name, value): + def _read_register_callback(self, state, reg_name, value): logger.debug("Read Register %r %r", reg_name, value) def _write_register_callback(self, state, reg_name, value): @@ -585,7 +585,7 @@ class Manticore(object): def _write_memory_callback(self, state, address, value, size): logger.debug("Write Memory %r %r %r", address, value, size) - def _decode_instruction_callback(self, state): + def _decode_instruction_callback(self, state, pc): logger.debug("Decoding stuff instruction not available") @@ -707,36 +707,31 @@ class Manticore(object): self._executor = Executor(initial_state, workspace=ws_path, policy=self._policy, - dumpafter=self.dumpafter, - maxstates=self.maxstates, - maxstorage=self.maxstorage, - replay=replay, - dumpstats=self.should_profile, context=self.context) #Link Executor events to default callbacks in manticore object - self._executor.did_read_register += self._read_register_callback - self._executor.will_write_register += self._write_register_callback - self._executor.did_read_memory += self._read_memory_callback - self._executor.will_write_memory += self._write_memory_callback - self._executor.will_execute_instruction += self._execute_instruction_callback - self._executor.will_decode_instruction += self._decode_instruction_callback - self._executor.will_store_state += self._store_state_callback - self._executor.will_load_state += self._load_state_callback - self._executor.will_fork_state += self._fork_state_callback - self._executor.will_terminate_state += self._terminate_state_callback - self._executor.will_generate_testcase += self._generate_testcase_callback + self._executor.subscribe('did_read_register', self._read_register_callback) + self._executor.subscribe('will_write_register', self._write_register_callback) + self._executor.subscribe('did_read_memory', self._read_memory_callback) + self._executor.subscribe('will_write_memory', self._write_memory_callback) + self._executor.subscribe('will_execute_instruction', self._execute_instruction_callback) + self._executor.subscribe('will_decode_instruction', self._decode_instruction_callback) + self._executor.subscribe('will_store_state', self._store_state_callback) + self._executor.subscribe('will_load_state', self._load_state_callback) + self._executor.subscribe('will_fork_state', self._fork_state_callback) + self._executor.subscribe('will_terminate_state', self._terminate_state_callback) + self._executor.subscribe('will_generate_testcase', self._generate_testcase_callback) if self._hooks: - self._executor.will_execute_instruction += self._hook_callback + self._executor.subscribe('will_execute_instruction', self._hook_callback) if self._model_hooks: - self._executor.will_execute_instruction += self._model_hook_callback + self._executor.subscribe('will_execute_instruction', self._model_hook_callback) if self._assertions: - self._executor.will_execute_instruction += self._assertions_callback + self._executor.subscribe('will_execute_instruction', self._assertions_callback) self._time_started = time.time() diff --git a/manticore/platforms/decree.py b/manticore/platforms/decree.py index 410d519..74ededd 100644 --- a/manticore/platforms/decree.py +++ b/manticore/platforms/decree.py @@ -1,7 +1,6 @@ import cgcrandom import weakref import sys, os, struct -from ..utils.event import Signal, forward_signals # TODO use cpu factory from ..core.cpu.x86 import I386Cpu from ..core.cpu.abstractcpu import Interruption, Syscall, ConcretizeRegister @@ -10,6 +9,7 @@ from ..core.smtlib import * from ..core.executor import TerminateState from ..utils.helpers import issymbolic from ..binary import CGCElf +from ..platforms.platform import Platform from contextlib import closing import StringIO import logging @@ -73,7 +73,7 @@ class Socket(object): return len(buf) -class Decree(object): +class Decree(Platform): ''' A simple Decree Operating System. This class emulates the most common Decree system calls @@ -88,7 +88,7 @@ class Decree(object): CGC_SIZE_MAX=4294967295 CGC_FD_SETSIZE=32 - def __init__(self, programs): + def __init__(self, programs, **kwargs): ''' Builds a Decree OS :param cpus: CPU for this platform @@ -97,6 +97,7 @@ class Decree(object): :todo: fix deps? ''' programs = programs.split(",") + super(Decree, self).__init__(path=programs[0], **kwargs) self.clocks = 0 self.files = [] self.syscall_trace = [] @@ -144,7 +145,7 @@ class Decree(object): #Install event forwarders for proc in self.procs: - forward_signals(self, proc) + self.forward_events_from(proc) def _mk_proc(self): @@ -155,7 +156,7 @@ class Decree(object): return self.procs[self._current] def __getstate__(self): - state = {} + state = super(Decree, self).__getstate__() state['clocks'] = self.clocks state['input'] = self.input.buffer state['output'] = self.output.buffer @@ -166,7 +167,6 @@ class Decree(object): state['rwait'] = self.rwait state['twait'] = self.twait state['timers'] = self.timers - state['syscall_trace'] = self.syscall_trace return state @@ -175,6 +175,7 @@ class Decree(object): :todo: some asserts :todo: fix deps? (last line) """ + super(Decree, self).__setstate__(state) self.input = Socket() self.input.buffer = state['input'] self.output = Socket() @@ -206,7 +207,7 @@ class Decree(object): #Install event forwarders for proc in self.procs: - forward_signals(self, proc) + self.forward_events_from(proc) def _read_string(self, cpu, buf): """ diff --git a/manticore/platforms/linux.py b/manticore/platforms/linux.py index 7c1ecff..c010dc5 100644 --- a/manticore/platforms/linux.py +++ b/manticore/platforms/linux.py @@ -8,8 +8,6 @@ import ctypes #Remove in favor of binary.py from elftools.elf.elffile import ELFFile - -from ..utils.event import Signal, forward_signals from ..utils.helpers import issymbolic from ..core.cpu.abstractcpu import Interruption, Syscall, ConcretizeArgument from ..core.cpu.cpufactory import CpuFactory @@ -284,7 +282,7 @@ class Linux(Platform): This class emulates the most common Linux system calls ''' - def __init__(self, program, argv=None, envp=None): + def __init__(self, program, argv=None, envp=None, **kwargs): ''' Builds a Linux OS platform :param string program: The path to ELF binary @@ -293,7 +291,7 @@ class Linux(Platform): :ivar files: List of active file descriptors :type files: list[Socket] or list[File] ''' - super(Linux, self).__init__(program) + super(Linux, self).__init__(path=program, **kwargs) self.program = program self.clocks = 0 @@ -382,7 +380,7 @@ class Linux(Platform): #Install event forwarders for proc in self.procs: - forward_signals(self, proc) + self.forward_events_from(proc) def _mk_proc(self, arch): if arch in {'i386', 'armv7'}: @@ -396,7 +394,7 @@ class Linux(Platform): return self.procs[self._current] def __getstate__(self): - state = {} + state = super(Linux, self).__getstate__() state['clocks'] = self.clocks state['input'] = self.input.buffer state['output'] = self.output.buffer @@ -437,6 +435,8 @@ class Linux(Platform): :todo: some asserts :todo: fix deps? (last line) """ + super(Linux, self).__setstate__(state) + self.input = Socket() self.input.buffer = state['input'] self.output = Socket() @@ -481,7 +481,7 @@ class Linux(Platform): #Install event forwarders for proc in self.procs: - forward_signals(self, proc) + self.forward_events_from(proc) def _init_arm_kernel_helpers(self): ''' diff --git a/manticore/platforms/platform.py b/manticore/platforms/platform.py index 0778b20..71d0723 100644 --- a/manticore/platforms/platform.py +++ b/manticore/platforms/platform.py @@ -1,5 +1,5 @@ -from manticore.utils.event import Signal +from manticore.utils.event import Eventful from itertools import islice, imap import inspect @@ -21,12 +21,22 @@ class ConcretizeSyscallArgument(OSException): super(ConcretizeSyscallArgument, self).__init__(message) -class Platform(object): +class Platform(Eventful): ''' Base class for all operating system platforms. ''' - def __init__(self, path): - self._path = path + def __init__(self, path, **kwargs): + super(Platform, self).__init__(**kwargs) + self._path = path #Not clear why all platforms must have a "path" def invoke_model(self, model, prefix_args=None): self._function_abi.invoke(model, prefix_args) + + def __setstate__(self, state): + super(Platform, self).__setstate__(state) + self._path = state['path'] + + def __getstate__(self): + state = super(Platform, self).__getstate__() + state['path'] = self._path + return state diff --git a/manticore/platforms/windows.py b/manticore/platforms/windows.py index 1e02eda..4ce5c65 100644 --- a/manticore/platforms/windows.py +++ b/manticore/platforms/windows.py @@ -9,7 +9,6 @@ from ..core.cpu.abstractcpu import Interruption, Syscall from ..core.state import ForkState, TerminateState from ..utils.helpers import issymbolic from ..platforms.platform import * -from ..utils.event import Signal, forward_signals from ..binary.pe import minidump @@ -63,10 +62,11 @@ class Windows(Platform): def _mk_memory(self): return Memory32() - def __init__(self, path, additional_context = None, snapshot_folder=None): + def __init__(self, path, additional_context = None, snapshot_folder=None, **kwargs): ''' Builds a Windows OS platform ''' + super(Windows, self).__init__(path,**kwargs) self.clocks = 0 self.files = [] self.syscall_trace = [] @@ -178,7 +178,7 @@ class Windows(Platform): #Install event forwarders for proc in self.procs: - forward_signals(self, proc) + self.forward_events_from(proc) @property @@ -190,7 +190,7 @@ class Windows(Platform): return self.procs[self._current] def __getstate__(self): - state = {} + state = super(Windows, self).__getstate__() state['clocks'] = self.clocks state['procs'] = self.procs state['current'] = self._current @@ -206,6 +206,7 @@ class Windows(Platform): :todo: some asserts :todo: fix deps? (last line) """ + super(Windows, self).__setstate__(state) self.procs = state['procs'] self._current = state['current'] self.running = state['running'] @@ -216,7 +217,7 @@ class Windows(Platform): #Install event forwarders for proc in self.procs: - forward_signals(self, proc) + self.forward_events_from(proc) def _read_string(self, cpu, buf): """ diff --git a/manticore/utils/event.py b/manticore/utils/event.py index 151b882..72939ed 100644 --- a/manticore/utils/event.py +++ b/manticore/utils/event.py @@ -1,168 +1,80 @@ import inspect -from weakref import WeakSet, WeakKeyDictionary +from weakref import ref, WeakSet, WeakKeyDictionary, WeakValueDictionary from types import MethodType -# Inspired by: -# http://code.activestate.com/recipes/577980-improved-signalsslots-implementation-in-python/ - -class SignalDisconnectedError(RuntimeError): - pass - - -def forward_signals(dest, source, arg=False): +class Eventful(object): ''' - Replicate and forward all the signals from source to dest + Abstract class for objects emitting and receiving events + An eventful object can: + - publish an event with arbitrary arguments to its subscribers + - let foreign objects subscribe their methods to events emitted here + - forward events to/from other eventful objects ''' - #Import all signals from state - for signal_name in source.__dict__: - signal = getattr(source, signal_name, None) - if isinstance(signal, Signal): - proxy = getattr(dest, signal_name, Signal()) - proxy.when(source, signal, arg) - setattr(dest, signal_name, proxy) - -def _manage_signals(obj, enabled): - ''' - Enable or disable all signals at obj - ''' - #Import all signals from state - for signal_name in dir(obj): - signal = getattr(obj, signal_name) - if isinstance(signal, Signal): - if enabled: - signal.enable() - else: - signal.disable() - - -def enable_signals(obj): - _manage_signals(obj, True) - -def disable_signals(obj): - _manage_signals(obj, False) - - -class Signal(object): - ''' - The Signal class is an approximation of Qt's signals+slot system. Each event - that an object would like to produce requires a Signal() object. All - interested parties on the event must register themselves as receivers via - connect() or the '+=' operator. - - The event source calls emit() to produce an event (or treat it as a - callable). All registered receivers will receive it, synchronously. - ''' - - def __init__(self, description=None): - ''' - Create a Signal() object. Pass 'True' to constructor if locking around - emit() is required. - ''' - self.description = description - self._functions = WeakSet() - self._methods = WeakKeyDictionary() + def __init__(self, *args, **kwargs): + # A dictionary from "event name" -> callback methods + # Note that several methods can be associated with the same object + self._signals = dict() + # a set of sink eventful objects (see forward_events_from()) self._forwards = WeakKeyDictionary() - self.disabled = False + super(Eventful, self).__init__(*args, **kwargs) - def disable(self): - self.disabled = True - def enable(self): - self.disabled = False - + def __setstate__(self, state): + ''' It wont get serialized by design, user is responsible to reconnect''' + self._signals = dict() + self._forwards = WeakKeyDictionary() + return True - def __len__(self): - return len(self._functions) + len(self._methods) + def __getstate__(self): + return {} - def __call__(self, *args, **kwargs): - return self.emit(*args, **kwargs) + def _unref(self, robj): + # this is called when an object that has subscribed to events emitted + # here has recently been garbage collected + # This simply removes all callback methods associated with that object + # Also if no more callbacks at all for an event name it deletes the event entry + remove = set() + for name, bucket in self._signals.iteritems(): + if robj in bucket: + del bucket[robj] + if len(bucket) == 0: + remove.add(name) + for name in remove: + del self._signals[name] - def emit(self, *args, **kwargs): - 'Invoke the signal with |args| and |kwargs|' - results = [] + def _get_signal_bucket(self, name): + #Each event name has a bucket of callback methods + #A bucket is a dictionary obj -> set(method1, method2...) + return self._signals.setdefault(name, dict()) - if self.disabled: - return results + def publish(self, name, *args, **kwargs): + bucket = self._get_signal_bucket(name) + for robj, methods in bucket.items(): + for callback in methods: + callback(robj(), *args, **kwargs) - for f in self._functions: - if '__predicate__' in f.__dict__: - if not f.__dict__['__predicate__'](): - continue - - results.append(f(*args, **kwargs)) - - for obj, funcs in self._methods.items(): - for f in funcs: - if '__predicate__' in f.__dict__: - if not f.__dict__['__predicate__'](): - continue - - results.append(f(obj, *args, **kwargs)) - - return results - - def connect(self, dest, predicate=None): - ''' - Connect |dest| to the signal. If |predicate| is set, it is treated as a - nullary callable whose return value determines if the signal is fired. - - NOTE: Passing identical values to multiple invocations of connect() with - different values of predicate will overwrite previous predicates and - persist the last-used value. - - To achieve a similar effect, wrap |dest| in a function. - - ''' - assert callable(dest) - if inspect.ismethod(dest): - obj, impl = dest.__self__, dest.__func__ - - if predicate is not None: - impl.__dict__['__predicate__'] = predicate - - self._methods.setdefault(obj, set()).add(impl) - else: - if predicate is not None: - dest.__dict__['__predicate__'] = predicate - - self._functions.add(dest) - - for signal, methods in self._forwards.items(): - for method in methods: - signal.connect(method) - self._forwards.clear() - - def when(self, obj, signal, arg=False): - ''' This forwards signal from obj ''' - #will reemit forwarded signal prepending obj to arguments - if arg: - method = MethodType(lambda *args, **kwargs: self.emit(*args, **kwargs), obj) - else: - method = self.emit - if len(self): - signal.connect(method) - else: - self._forwards.setdefault(signal, set()).add(method) - - def __iadd__(self, dest): - self.connect(dest) - return self - - def disconnect(self, dest): - try: - if inspect.ismethod(dest): - obj, impl = dest.__self__, dest.__func__ - self._methods[obj].remove(impl) + #The include_source flag indicates to prepend the source of the event in + # the callback signature. This is set on forward_events_from/to + for sink, include_source in self._forwards.items(): + if include_source: + sink.publish(name, self, *args, **kwargs) else: - self._functions.remove(dest) - except KeyError: - raise SignalDisconnectedError() + sink.publish(name, *args, **kwargs) - def __isub__(self, dest): - self.disconnect(dest) - return self + def subscribe(self, name, method): + if not inspect.ismethod(method): + raise TypeError + obj, callback = method.__self__, method.__func__ + bucket = self._get_signal_bucket(name) + robj = ref(obj, self._unref) #see unref() for explanation + bucket.setdefault(robj, set()).add(callback) - def reset(self): - self._functions.clear() - self._methods.clear() - self._forwards.clear() + def forward_events_from(self, source, include_source=False): + if not isinstance(source, Eventful): + raise TypeError + source.forward_events_to(self, include_source=include_source) + def forward_events_to(self, sink, include_source=False): + ''' This forwards signal to sink ''' + if not isinstance(sink, Eventful): + raise TypeError + self._forwards[sink] = include_source diff --git a/tests/test_binaries.py b/tests/test_binaries.py index ef693c7..0f8b0ea 100644 --- a/tests/test_binaries.py +++ b/tests/test_binaries.py @@ -61,6 +61,7 @@ class IntegrationTest(unittest.TestCase): '--procs', '4', filename, '+++++++++'], stdout=output) + self.assertTrue(time.time()-t < 20) def test_cli_verbosity(self): diff --git a/tests/test_events.py b/tests/test_events.py new file mode 100644 index 0000000..26dca3c --- /dev/null +++ b/tests/test_events.py @@ -0,0 +1,77 @@ + +import unittest + +from manticore.utils.event import Eventful + +class A(Eventful): + def do_stuff(self): + self.publish("eventA",1, 'a') + +class B(Eventful): + def __init__(self, child, **kwargs): + super(B, self).__init__(**kwargs) + self.child = child + self.forward_events_from(child) + + def do_stuff(self): + self.publish("eventB", 2, 'b') + + +class C(): + def __init__(self): + self.received = [] + def callback(self, *args): + self.received.append(args) + +class ManticoreDriver(unittest.TestCase): + _multiprocess_can_split_ = True + def setUp(self): + self.state = {} + + def tearDown(self): + pass + + def test_weak_references(self): + a = A() + self.assertSequenceEqual( map(len, (a._signals, a._forwards)), (0, 0) ) + + b = B(a) + self.assertSequenceEqual( map(len, (a._signals, a._forwards)), (0, 1) ) + self.assertSequenceEqual( map(len, (b._signals, b._forwards)), (0, 0) ) + + c = C() + b.subscribe('eventA', c.callback) + + self.assertSequenceEqual( map(len, (a._signals, a._forwards)), (0, 1) ) + self.assertSequenceEqual( map(len, (b._signals, b._forwards)), (1, 0) ) + + b.subscribe('eventB', c.callback) + + self.assertSequenceEqual( map(len, (a._signals, a._forwards)), (0, 1) ) + self.assertSequenceEqual( map(len, (b._signals, b._forwards)), (2, 0) ) + + + del c + + self.assertSequenceEqual( map(len, (a._signals, a._forwards)), (0, 1) ) + self.assertSequenceEqual( map(len, (b._signals, b._forwards)), (0, 0) ) + + del b + + self.assertSequenceEqual( map(len, (a._signals, a._forwards)), (0, 0) ) + + def test_basic(self): + a = A() + b = B(a) + c = C() + + b.subscribe('eventA', c.callback) + b.subscribe('eventB', c.callback) + + a.do_stuff() + self.assertSequenceEqual(c.received, [(1, 'a')]) + + b.do_stuff() + self.assertSequenceEqual(c.received, [(1, 'a'), (2, 'b')]) + + diff --git a/tests/test_signal.py b/tests/test_signal.py deleted file mode 100644 index 5a2c435..0000000 --- a/tests/test_signal.py +++ /dev/null @@ -1,70 +0,0 @@ - -import unittest - -from manticore.utils.event import Signal - -class Sender(object): - def __init__(self): - self.sig = Signal() - self.sig2 = Signal() - -class ManticoreDriver(unittest.TestCase): - _multiprocess_can_split_ = True - def setUp(self): - self.state = {} - - def tearDown(self): - pass - - def setReceived(self, key, val): - self.state[key] = val - - def setReceived2(self, key, val): - self.state[key] = val - - - def test_basic(self): - s = Sender() - - def recv(): - self.state['received'] = True - - self.state['received'] = False - - s.sig += recv - - s.sig() - - self.assertEqual(self.state['received'], True) - - def test_method(self): - s = Sender() - s.sig += self.setReceived - - s.sig('received', True) - - self.assertEqual(self.state['received'], True) - - def test_disconnect(self): - s = Sender() - - s.sig.connect(self.setReceived) - - s.sig -= self.setReceived - - s.sig('received', True) - - self.assertNotIn('received', self.state) - - def test_predicate(self): - s = Sender() - - s.sig.connect(self.setReceived) - s.sig2.connect(self.setReceived2, lambda: False) - - s.sig('true', True) - s.sig2('false', True) - - self.assertEqual(self.state['true'], True) - self.assertNotIn('false', self.state) - diff --git a/tests/test_state.py b/tests/test_state.py index f460a97..7dadae8 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -1,7 +1,6 @@ import unittest - +from manticore.utils.event import Eventful from manticore.platforms import linux -from manticore.utils.event import Signal from manticore.core.state import State from manticore.core.smtlib import BitVecVariable, ConstraintSet @@ -19,25 +18,15 @@ class FakeMemory(object): class FakeCpu(object): def __init__(self): - self.will_decode_instruction = Signal() - self.will_execute_instruction = Signal() - self.did_execute_instruction = Signal() - self.will_emulate_instruction = Signal() - self.did_emulate_instruction = Signal() - - self.will_read_register = Signal() - self.will_write_register = Signal() - self.will_read_memory = Signal() - self.will_write_memory = Signal() - self._memory = FakeMemory() @property def memory(self): return self._memory -class FakePlatform(object): +class FakePlatform(Eventful): def __init__(self): + super(FakePlatform, self).__init__() self._constraints = None self.procs = [FakeCpu()] diff --git a/tests/test_workspace.py b/tests/test_workspace.py index 4e8c4a8..c26bb6f 100644 --- a/tests/test_workspace.py +++ b/tests/test_workspace.py @@ -4,10 +4,10 @@ import unittest from multiprocessing.managers import SyncManager from manticore.platforms import linux -from manticore.utils.event import Signal from manticore.core.state import State from manticore.core.smtlib import BitVecVariable, ConstraintSet from manticore.core.workspace import * +from manticore.utils.event import Eventful manager = SyncManager() manager.start(lambda: signal.signal(signal.SIGINT, signal.SIG_IGN)) @@ -24,18 +24,9 @@ class FakeMemory(object): def constraints(self, constraints): self._constraints = constraints -class FakeCpu(object): +class FakeCpu(Eventful): def __init__(self): - self.will_decode_instruction = Signal() - self.will_execute_instruction = Signal() - self.did_execute_instruction = Signal() - self.will_emulate_instruction = Signal() - self.did_emulate_instruction = Signal() - - self.will_read_register = Signal() - self.will_write_register = Signal() - self.will_read_memory = Signal() - self.will_write_memory = Signal() + super(FakeCpu, self).__init__() self._memory = FakeMemory()