From 778c0ca5c319a4dd2b34eae62062846e572a2411 Mon Sep 17 00:00:00 2001
From: Malfurious <m@lfurio.us>
Date: Tue, 20 Feb 2024 12:44:35 -0500
Subject: payload: Improve recursion performance

There is a small network of mutually-recursive helper functions which
produce the main outputs for Payload objects (the length, bytes, etc.).
The runtime performance of this code can suffer as a Payload grows to
contain more and more items.

These issues are heavily mitigated by implementing memoization within
one of these functions (which propagates the benefit to the rest of the
call tree).  Memo dictionary is only used for a single operation
(lifetime) to avoid the possibility of bad cached results.

Signed-off-by: Malfurious <m@lfurio.us>
---
 sploit/payload/payload.py | 36 ++++++++++++++++++++++--------------
 1 file changed, 22 insertions(+), 14 deletions(-)

(limited to 'sploit/payload/payload.py')

diff --git a/sploit/payload/payload.py b/sploit/payload/payload.py
index a59eed2..b88e323 100644
--- a/sploit/payload/payload.py
+++ b/sploit/payload/payload.py
@@ -84,6 +84,7 @@ class Payload(IndexTbl):
         """Return human-readable Payload."""
         FMT = "\n{:<20} {:<20} {:<20}"
         s = f"{len(self.__entries__)} items, {len(self)} bytes @ {hex(self)}"
+        memo = {}
 
         if len(self.__entries__) > 0:
             s += FMT.format("ADDRESS", "SYMBOL", "DATA")
@@ -93,8 +94,8 @@ class Payload(IndexTbl):
             key = "(unkeyed)" if key is None else str(key)
             key = f"[{key}]" if isinstance(value, IndexEntry) else key
 
-            addr = self.__addrof(i)
-            data = str(self.__bytesof(i))
+            addr = self.__addrof(i, memo)
+            data = str(self.__bytesof(i, memo))
             if len(data) > _REPR_DATA_LEN:
                 data = data[:_REPR_DATA_LEN] + " ..."
 
@@ -104,7 +105,8 @@ class Payload(IndexTbl):
 
     def __bytes__(self):
         """Return calculated payload bytes."""
-        x = [ self.__bytesof(i) for i in range(len(self.__entries__)) ]
+        memo = {}
+        x = [ self.__bytesof(i, memo) for i in range(len(self.__entries__)) ]
         return b"".join(x)
 
     def __call__(self, *args):
@@ -146,13 +148,13 @@ class Payload(IndexTbl):
 
     def __getindex__(self, index):
         """Return payload index value or address."""
-        value, _ = self.__valueof(index)
+        value, _ = self.__valueof(index, {})
         return value
 
     def __setindex__(self, index, value):
         """Set payload index value."""
         try:
-            addr = self.__addrof(index)
+            addr = self.__addrof(index, {})
         except KeyError:
             addr = self.end()
         value = self.__prep_insertion(value, addr)
@@ -164,27 +166,33 @@ class Payload(IndexTbl):
 
     # Payload helpers
 
-    def __valueof(self, index):
+    def __valueof(self, index, memo):
         """Return a tuple: (addr of value, literal value) for index."""
         value = self.__entries__[index]
-        addr = self.__addrof(index)
+        addr = self.__addrof(index, memo)
         if isinstance(value, IndexEntry):
             value @= addr
             return value, value
         return addr, value
 
-    def __addrof(self, index):
+    def __addrof(self, index, memo):
         """Return address (base + offset) for index."""
         index = self.__entries__.key2idx(index)
-        sizes = [ len(self.__bytesof(i)) for i in range(index) ]
+        sizes = [ len(self.__bytesof(i, memo)) for i in range(index) ]
         return self.base + sum(sizes)
 
-    def __bytesof(self, index):
+    def __bytesof(self, index, memo):
         """Return byte output for index."""
-        _, value = self.__valueof(index)
-        if isinstance(value, PayloadEntry):
-            return value.payload_bytes(self)
-        return bytes(value)
+        try:
+            return memo[index]
+        except KeyError:
+            _, value = self.__valueof(index, memo)
+            if isinstance(value, PayloadEntry):
+                data = value.payload_bytes(self)
+            else:
+                data = bytes(value)
+            memo[index] = data
+            return data
 
     def __prep_insertion(self, value, addr):
         """Initialize or type coerce input value for payload insert."""
-- 
cgit v1.2.3