# built-in imports
from enum import Enum, auto
import enum_tools.documentation
# third-party imports
from triton import TritonContext
@enum_tools.documentation.document_enum
class ThreadState(Enum):
RUNNING = auto() # doc: Normal state
DEAD = auto() # doc: State after pthread_exit
JOINING = auto() # doc: State after pthread_join
LOCKED = auto() # doc: State after a pthread_lock & co
[docs]
class ThreadContext(object):
"""
Thread data structure holding all information related to it.
Purposely used to save registers and to restore them in a
TritonContext.
"""
def __init__(self, tid: int):
"""
:param tid: thread id
"""
self.cregs = dict() # context of concrete registers
self.sregs = dict() # context of symbolic registers
self._join_th_id = None # joined thread id
self.tid = tid # the thread id
self.count = 0 # Number of instructions executed until scheduling
self.state = ThreadState.RUNNING
# FIXME: Keep the thread_scheduling and automated the reset on restore
[docs]
def save(self, tt_ctx: TritonContext) -> None:
"""
Save the current thread state from the current execution.
That implies keeping a reference on symbolic and concrete
registers.
:param tt_ctx: current TritonContext to save
:type tt_ctx: `TritonContext <https://triton.quarkslab.com/documentation/doxygen/py_TritonContext_page.html>`_
"""
# Save symbolic registers
self.sregs = tt_ctx.getSymbolicRegisters()
# Save concrete registers
for r in tt_ctx.getParentRegisters():
self.cregs.update({r.getId(): tt_ctx.getConcreteRegisterValue(r)})
[docs]
def restore(self, tt_ctx: TritonContext) -> None:
"""
Restore a thread state in the given TritonContext
:param tt_ctx: context in which to restor the current thread state
:type tt_ctx: `TritonContext <https://triton.quarkslab.com/documentation/doxygen/py_TritonContext_page.html>`_
"""
# Restore concrete registers
for rid, v in self.cregs.items():
tt_ctx.setConcreteRegisterValue(tt_ctx.getRegister(rid), v)
# Restore symbolic registers
for rid, e in self.sregs.items():
tt_ctx.assignSymbolicExpressionToRegister(e, tt_ctx.getRegister(rid))
[docs]
def kill(self) -> None:
"""
Kill the current thread. Called when exiting the thread.
:return:
"""
self.state = ThreadState.DEAD
[docs]
def is_dead(self) -> bool:
"""
Returns whether the thread is killed or not
:return: boolean indicating if the thread is dead or not
"""
return self.state == ThreadState.DEAD
[docs]
def join_thread(self, th_id: int) -> None:
"""
Put the thread in a join state where waits for
another thread.
:param th_id: id of the thread to join
:return: None
"""
self._join_th_id = th_id
self.state = ThreadState.JOINING
[docs]
def is_waiting_to_join(self) -> bool:
"""
Checks whether the thread is waiting to join
another one.
:return: boolean on whether it waits for another thread
"""
return self.state == ThreadState.JOINING
[docs]
def cancel_join(self) -> None:
"""
Cancel a join operation.
:return: None
"""
self._join_th_id = None
self.state = ThreadState.RUNNING
[docs]
def is_main_thread(self) -> bool:
"""
Returns whether it is the main thread
(namely its id is 0)
:return: bool
"""
return self.tid == 0
[docs]
def is_running(self) -> bool:
"""
Return if the thread is properly running or not.
:return: True if the thread is running
"""
return self.state == ThreadState.RUNNING