From e8fd456af0611287666d95e7dbc8a9091fb931dc Mon Sep 17 00:00:00 2001
From: neok-m4700 <neok-m4700@users.noreply.github.com>
Date: Fri, 12 Feb 2021 20:46:21 +0100
Subject: [PATCH] avoid fetching all registers when using gereg(...)

---
 ptrace/binding/__init__.py |  2 +-
 ptrace/binding/func.py     |  1 +
 ptrace/debugger/process.py | 18 +++++++-----------
 3 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/ptrace/binding/__init__.py b/ptrace/binding/__init__.py
index eb0374c..027bd62 100644
--- a/ptrace/binding/__init__.py
+++ b/ptrace/binding/__init__.py
@@ -1,7 +1,7 @@
 from ptrace.binding.func import (   # noqa
     HAS_PTRACE_SINGLESTEP, HAS_PTRACE_EVENTS,
     HAS_PTRACE_IO, HAS_PTRACE_SIGINFO, HAS_PTRACE_GETREGS,
-    HAS_PTRACE_GETREGSET, REGISTER_NAMES,
+    HAS_PTRACE_GETREGSET, REGISTER_NAMES, REGISTER_OFFSETS,
     ptrace_attach, ptrace_traceme,
     ptrace_detach, ptrace_kill,
     ptrace_cont, ptrace_syscall,
diff --git a/ptrace/binding/func.py b/ptrace/binding/func.py
index bc4bf4d..5215fa9 100644
--- a/ptrace/binding/func.py
+++ b/ptrace/binding/func.py
@@ -23,6 +23,7 @@
 else:
     raise NotImplementedError("Unknown OS!")
 REGISTER_NAMES = tuple(name for name, type in ptrace_registers_t._fields_)
+REGISTER_OFFSETS = {n: getattr(ptrace_registers_t, n).offset for n, t in ptrace_registers_t._fields_}
 
 HAS_PTRACE_SINGLESTEP = True
 HAS_PTRACE_EVENTS = False
diff --git a/ptrace/debugger/process.py b/ptrace/debugger/process.py
index db4573d..2e63447 100644
--- a/ptrace/debugger/process.py
+++ b/ptrace/debugger/process.py
@@ -4,9 +4,9 @@
     HAS_PTRACE_GETREGSET,
     ptrace_attach, ptrace_detach,
     ptrace_cont, ptrace_syscall,
-    ptrace_setregs,
+    ptrace_setregs, ptrace_peekuser,
     ptrace_peektext, ptrace_poketext,
-    REGISTER_NAMES)
+    REGISTER_NAMES, REGISTER_OFFSETS)
 from ptrace.os_tools import HAS_PROC, RUNNING_BSD, RUNNING_PYTHON3
 from ptrace.tools import dumpRegs
 from ptrace.cpu_info import CPU_WORD_SIZE
@@ -49,7 +49,7 @@
 if HAS_PTRACE_GETREGS or HAS_PTRACE_GETREGSET:
     from ptrace.binding import ptrace_getregs
 else:
-    from ptrace.binding import ptrace_peekuser, ptrace_registers_t
+    from ptrace.binding import ptrace_registers_t
 if HAS_DISASSEMBLER:
     from ptrace.disasm import disassemble, disassembleOne, MAX_INSTR_SIZE
 if HAS_PROC:
@@ -416,13 +416,10 @@ def getregs(self):
         if HAS_PTRACE_GETREGS or HAS_PTRACE_GETREGSET:
             return ptrace_getregs(self.pid)
         else:
-            # FIXME: Optimize getreg() when used with this function
             words = []
-            nb_words = sizeof(ptrace_registers_t) // CPU_WORD_SIZE
-            for offset in range(nb_words):
-                word = ptrace_peekuser(self.pid, offset * CPU_WORD_SIZE)
-                bytes = word2bytes(word)
-                words.append(bytes)
+            for name in REGISTER_OFFSETS:
+                word = ptrace_peekuser(self.pid, REGISTER_OFFSETS[name])
+                words.append(word2bytes(word))
             bytes = ''.join(words)
             return bytes2type(bytes, ptrace_registers_t)
 
@@ -434,8 +431,7 @@ def getreg(self, name):
             mask = None
         if name not in REGISTER_NAMES:
             raise ProcessError(self, "Unknown register: %r" % name)
-        regs = self.getregs()
-        value = getattr(regs, name)
+        value = ptrace_peekuser(self.pid, REGISTER_OFFSETS[name])
         value >>= shift
         if mask:
             value &= mask