switch metadata tables to lazy parsing

This commit is contained in:
Shiz 2020-06-02 02:26:58 +02:00
parent 6e552048bd
commit ce9a1aefa2
3 changed files with 59 additions and 29 deletions

View File

@ -119,22 +119,31 @@ class CLICodedToken(Type):
def __init__(self, types):
self.types = types
def parse(self, input, context):
metadata = context.user.metadata
def _get_size(self, metadata):
row_counts = []
for t in self.types:
row_counts.append(metadata.row_counts.get(t, 0))
max_rows = max(row_counts)
table_bits = math.ceil(math.log(len(self.types), 2))
if max_rows < (1 << 16 - table_bits):
size = 2
return table_bits, 2
else:
size = 4
return table_bits, 4
def parse(self, input, context):
metadata = context.user.metadata
table_bits, size = self._get_size(metadata)
val = int.from_bytes(input.read(size), 'little')
row = val >> table_bits
tag = val & ((1 << table_bits) - 1)
return CLIToken(row=row, table=self.types[tag])
def sizeof(self, value, context):
metadata = context.user.metadata
_, size = self._get_size(metadata)
return size
class CLISectionReference(Struct, nocopy=True):
rva = UInt(32)
size = UInt(32)
@ -143,46 +152,66 @@ class CLITableIndex(Type):
def __init__(self, type):
self.type = type
def parse(self, input, context):
metadata = context.user.metadata
def _get_size(self, metadata):
row_count = metadata.row_counts.get(self.type, 0)
if row_count < (1 << 16):
size = 2
return 2
else:
size = 4
return 4
def parse(self, input, context):
metadata = context.user.metadata
size = self._get_size(metadata)
return int.from_bytes(input.read(size), 'little')
def sizeof(self, input, context):
metadata = context.user.metadata
return self._get_size(metadata)
class CLITableRange(Type):
def __init__(self, type):
self.type = type
def parse(self, input, context):
metadata = context.user.metadata
def _get_size(self, metadata):
row_count = metadata.row_counts.get(self.type, 0)
if row_count < (1 << 16):
size = 2
return 2
else:
size = 4
return 4
def parse(self, input, context):
metadata = context.user.metadata
size = self._get_size(metadata)
return int.from_bytes(input.read(size), 'little')
def sizeof(self, input, context):
metadata = context.user.metadata
return self._get_size(metadata)
class CLIStreamIndex(Type):
def __init__(self, type, child=None):
self.type = type
self.child = child
def parse(self, input, context):
metadata = context.user.metadata
def _get_size(self, metadata):
big_mapping = {
CLIStreamType.String: CLIHeapFlags.BigStringStream,
CLIStreamType.Blob: CLIHeapFlags.BigBlobStream,
CLIStreamType.GUID: CLIHeapFlags.BigGUIDStream,
}
flag = big_mapping.get(self.type, None)
if flag:
if metadata.heap_flags & flag:
size = 4
else:
size = 2
if not flag:
return 2
elif metadata.heap_flags & flag:
return 4
else:
size = 2
return 2
def parse(self, input, context):
metadata = context.user.metadata
size = self._get_size(metadata)
return int.from_bytes(input.read(size), 'little')
def sizeof(self, input, context):
metadata = context.user.metadata
return self._get_size(metadata)

View File

@ -89,7 +89,7 @@ class CLIFile:
self.streams[type] = stream
def get_table(self, t):
return (wrap(self, t, i + 1, x) for i, x in enumerate(self.streams[CLIStreamType.Metadata].tables[t]))
return (wrap(self, t, i + 1, x()) for i, x in enumerate(self.streams[CLIStreamType.Metadata].tables[t]))
def get_table_size(self, t):
return len(self.streams[CLIStreamType.Metadata].tables[t])
@ -97,7 +97,7 @@ class CLIFile:
def get_table_entry(self, t, i):
if not i:
return None
return wrap(self, t, i, self.streams[CLIStreamType.Metadata].tables[t][i - 1])
return wrap(self, t, i, self.streams[CLIStreamType.Metadata].tables[t][i - 1]())
def get_string_at(self, i):
buf = bytearray()

View File

@ -1,6 +1,6 @@
import enum
import uuid
from destruct import Type, Struct, Arr
from destruct import sizeof, Type, Struct, Arr, Lazy
from .common import CLIStreamType, CLITableType, CLIHeapFlags
from .tables import TABLE_PARSERS
@ -56,14 +56,15 @@ class CLIMetadataStream(Struct, nocopy=True):
def on_present(self, spec, context):
spec.row_counts.count = len(self.present)
spec.tables.count = len(self.present)
context.user.metadata = self
def on_row_counts(self, spec, context):
counts = {}
for p, count in zip(self.present, self.row_counts):
spec.tables.child.append(Arr(TABLE_PARSERS[p], count=count))
counts[p] = count
self.row_counts = counts
context.user.metadata = self
self.row_counts = {p: count for p, count in zip(self.present, self.row_counts)}
for p in self.present:
elem = TABLE_PARSERS[p]
elem_size = sizeof(elem, None, context)
spec.tables.child.append(Arr(Lazy(elem, elem_size), count=self.row_counts[p]))
def on_tables(self, spec, context):
tables = {}