switch metadata tables to lazy parsing
This commit is contained in:
parent
6e552048bd
commit
ce9a1aefa2
|
@ -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)
|
||||
|
|
|
@ -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()
|
||||
|
|
|
@ -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 = {}
|
||||
|
|
Loading…
Reference in New Issue