diff options
Diffstat (limited to 'library/lua/memscan.lua')
| -rw-r--r-- | library/lua/memscan.lua | 455 |
1 files changed, 455 insertions, 0 deletions
diff --git a/library/lua/memscan.lua b/library/lua/memscan.lua new file mode 100644 index 00000000..95b9197b --- /dev/null +++ b/library/lua/memscan.lua @@ -0,0 +1,455 @@ +-- Utilities for offset scan scripts. + +local _ENV = mkmodule('memscan') + +local utils = require('utils') + +-- A length-checked view on a memory buffer + +CheckedArray = CheckedArray or {} + +function CheckedArray.new(type,saddr,eaddr) + local data = df.reinterpret_cast(type,saddr) + local esize = data:sizeof() + local count = math.floor((eaddr-saddr)/esize) + local obj = { + type = type, start = saddr, size = count*esize, + esize = esize, data = data, count = count + } + setmetatable(obj, CheckedArray) + return obj +end + +function CheckedArray:__len() + return self.count +end +function CheckedArray:__index(idx) + if type(idx) == number then + if idx >= self.count then + error('Index out of bounds: '..tostring(idx)) + end + return self.data[idx] + else + return CheckedArray[idx] + end +end +function CheckedArray:__newindex(idx, val) + if idx >= self.count then + error('Index out of bounds: '..tostring(idx)) + end + self.data[idx] = val +end +function CheckedArray:addr2idx(addr, round) + local off = addr - self.start + if off >= 0 and off < self.size and (round or (off % self.esize) == 0) then + return math.floor(off / self.esize), off + end +end +function CheckedArray:idx2addr(idx) + if idx >= 0 and idx < self.count then + return self.start + idx*self.esize + end +end + +-- Search methods + +function CheckedArray:find(data,sidx,eidx,reverse) + local dcnt = #data + sidx = math.max(0, sidx or 0) + eidx = math.min(self.count, eidx or self.count) + if (eidx - sidx) >= dcnt and dcnt > 0 then + return dfhack.with_temp_object( + df.new(self.type, dcnt), + function(buffer) + for i = 1,dcnt do + buffer[i-1] = data[i] + end + local cnt = eidx - sidx - dcnt + local step = self.esize + local sptr = self.start + sidx*step + local ksize = dcnt*step + if reverse then + local idx, addr = dfhack.internal.memscan(sptr + cnt*step, cnt, -step, buffer, ksize) + if idx then + return sidx + cnt - idx, addr + end + else + local idx, addr = dfhack.internal.memscan(sptr, cnt, step, buffer, ksize) + if idx then + return sidx + idx, addr + end + end + end + ) + end +end +function CheckedArray:find_one(data,sidx,eidx,reverse) + local idx, addr = self:find(data,sidx,eidx,reverse) + if idx then + -- Verify this is the only match + if reverse then + eidx = idx+#data-1 + else + sidx = idx+1 + end + if self:find(data,sidx,eidx,reverse) then + return nil + end + end + return idx, addr +end +function CheckedArray:list_changes(old_arr,old_val,new_val,delta) + if old_arr.type ~= self.type or old_arr.count ~= self.count then + error('Incompatible arrays') + end + local eidx = self.count + local optr = old_arr.start + local nptr = self.start + local esize = self.esize + local rv + local sidx = 0 + while true do + local idx = dfhack.internal.diffscan(optr, nptr, sidx, eidx, esize, old_val, new_val, delta) + if not idx then + break + end + rv = rv or {} + rv[#rv+1] = idx + sidx = idx+1 + end + return rv +end +function CheckedArray:filter_changes(prev_list,old_arr,old_val,new_val,delta) + if old_arr.type ~= self.type or old_arr.count ~= self.count then + error('Incompatible arrays') + end + local eidx = self.count + local optr = old_arr.start + local nptr = self.start + local esize = self.esize + local rv + for i=1,#prev_list do + local idx = prev_list[i] + if idx < 0 or idx >= eidx then + error('Index out of bounds: '..idx) + end + if dfhack.internal.diffscan(optr, nptr, idx, idx+1, esize, old_val, new_val, delta) then + rv = rv or {} + rv[#rv+1] = idx + end + end + return rv +end + +-- A raw memory area class + +MemoryArea = MemoryArea or {} +MemoryArea.__index = MemoryArea + +function MemoryArea.new(astart, aend) + local obj = { + start_addr = astart, end_addr = aend, size = aend - astart, + int8_t = CheckedArray.new('int8_t',astart,aend), + uint8_t = CheckedArray.new('uint8_t',astart,aend), + int16_t = CheckedArray.new('int16_t',astart,aend), + uint16_t = CheckedArray.new('uint16_t',astart,aend), + int32_t = CheckedArray.new('int32_t',astart,aend), + uint32_t = CheckedArray.new('uint32_t',astart,aend) + } + setmetatable(obj, MemoryArea) + return obj +end + +function MemoryArea:__gc() + df.delete(self.buffer) +end + +function MemoryArea:__tostring() + return string.format('<MemoryArea: %x..%x>', self.start_addr, self.end_addr) +end +function MemoryArea:contains_range(start,size) + return start >= self.start_addr and (start+size) <= self.end_addr +end +function MemoryArea:contains_obj(obj,count) + local size, base = df.sizeof(obj) + return size and base and self:contains_range(base, size*(count or 1)) +end + +function MemoryArea:clone() + local buffer = df.new('int8_t', self.size) + local _, base = buffer:sizeof() + local rv = MemoryArea.new(base, base+self.size) + rv.buffer = buffer + return rv +end +function MemoryArea:copy_from(area2) + if area2.size ~= self.size then + error('Size mismatch') + end + dfhack.internal.memmove(self.start_addr, area2.start_addr, self.size) +end +function MemoryArea:delete() + setmetatable(self, nil) + df.delete(self.buffer) + for k,v in pairs(self) do self[k] = nil end +end + +-- Static data segment search + +local function find_data_segment() + local data_start, data_end + + for i,mem in ipairs(dfhack.internal.getMemRanges()) do + if data_end then + if mem.start_addr == data_end and mem.read and mem.write then + data_end = mem.end_addr + else + break + end + elseif mem.read and mem.write + and (string.match(mem.name,'/dwarfort%.exe$') + or string.match(mem.name,'/Dwarf_Fortress$')) + then + data_start = mem.start_addr + data_end = mem.end_addr + end + end + + return data_start, data_end +end + +function get_data_segment() + local s, e = find_data_segment() + if s and e then + return MemoryArea.new(s, e) + end +end + +-- Register a found offset, or report an error. + +function found_offset(name,val) + local cval = dfhack.internal.getAddress(name) + + if not val then + print('Could not find offset '..name) + if not cval and not utils.prompt_yes_no('Continue with the script?') then + error('User quit') + end + return + end + + if df.isvalid(val) then + _,val = val:sizeof() + end + + print(string.format('Found offset %s: %x', name, val)) + + if cval then + if cval ~= val then + error(string.format('Mismatch with the current value: %x',val)) + end + else + dfhack.internal.setAddress(name, val) + end +end + +-- Offset of a field within struct + +function field_ref(handle,...) + local items = table.pack(...) + for i=1,items.n-1 do + handle = handle[items[i]] + end + return handle:_field(items[items.n]) +end + +function field_offset(type,...) + local handle = df.reinterpret_cast(type,1) + local _,addr = df.sizeof(field_ref(handle,...)) + return addr-1 +end + +function MemoryArea:object_by_field(addr,type,...) + if not addr then + return nil + end + local base = addr - field_offset(type,...) + local obj = df.reinterpret_cast(type, base) + if not self:contains_obj(obj) then + obj = nil + end + return obj, base +end + +-- Validation + +function is_valid_vector(ref,size) + local ints = df.reinterpret_cast('uint32_t', ref) + return ints[0] <= ints[1] and ints[1] <= ints[2] + and (size == nil or (ints[1] - ints[0]) % size == 0) +end + +-- Difference search helper + +DiffSearcher = DiffSearcher or {} +DiffSearcher.__index = DiffSearcher + +function DiffSearcher.new(area) + local obj = { area = area } + setmetatable(obj, DiffSearcher) + return obj +end + +function DiffSearcher:begin_search(type) + self.type = type + self.old_value = nil + self.search_sets = nil + if not self.save_area then + self.save_area = self.area:clone() + end +end +function DiffSearcher:advance_search(new_value,delta) + if self.search_sets then + local nsets = #self.search_sets + local ovec = self.save_area[self.type] + local nvec = self.area[self.type] + local new_set + if nsets > 0 then + local last_set = self.search_sets[nsets] + new_set = nvec:filter_changes(last_set,ovec,self.old_value,new_value,delta) + else + new_set = nvec:list_changes(ovec,self.old_value,new_value,delta) + end + if new_set then + self.search_sets[nsets+1] = new_set + self.old_value = new_value + self.save_area:copy_from(self.area) + return #new_set, new_set + end + else + self.old_value = new_value + self.search_sets = {} + self.save_area:copy_from(self.area) + return #self.save_area[self.type], nil + end +end +function DiffSearcher:reset() + self.search_sets = nil + if self.save_area then + self.save_area:delete() + self.save_area = nil + end +end +function DiffSearcher:idx2addr(idx) + return self.area[self.type]:idx2addr(idx) +end + +-- Interactive search utility + +function DiffSearcher:find_interactive(prompt,data_type,condition_cb) + enum = enum or {} + + -- Loop for restarting search from scratch + while true do + print('\n'..prompt) + + self:begin_search(data_type) + + local found = false + local ccursor = 0 + + -- Loop through choices + while true do + print('') + + local ok, value, delta = condition_cb(ccursor) + + ccursor = ccursor + 1 + + if not ok then + break + end + + -- Search for it in the memory + local cnt, set = self:advance_search(value, delta) + if not cnt then + dfhack.printerr(' Converged to zero candidates; probably a mistake somewhere.') + break + elseif set and cnt == 1 then + -- To confirm, wait for two 1-candidate results in a row + if found then + local addr = self:idx2addr(set[1]) + print(string.format(' Confirmed address: %x\n', addr)) + return addr, set[1] + else + found = true + end + end + + print(' '..cnt..' candidates remaining.') + end + + if not utils.prompt_yes_no('\nRetry search from the start?') then + return nil + end + end +end + +function DiffSearcher:find_menu_cursor(prompt,data_type,choices,enum) + enum = enum or {} + + return self:find_interactive( + prompt, data_type, + function(ccursor) + local choice + + -- Select the next value to search for + if type(choices) == 'function' then + choice = choices(ccursor) + + if not choice then + return false + end + else + choice = choices[(ccursor % #choices) + 1] + end + + -- Ask the user to select it + if enum ~= 'noprompt' then + local cname = enum[choice] or choice + if type(choice) == 'string' and type(cname) == 'number' then + choice, cname = cname, choice + end + if cname ~= choice then + cname = cname..' ('..choice..')' + end + + print(' Please select: '..cname) + if not utils.prompt_yes_no(' Continue?', true) then + return false + end + end + + return true, choice + end + ) +end + +function DiffSearcher:find_counter(prompt,data_type,delta,action_prompt) + delta = delta or 1 + + return self:find_interactive( + prompt, data_type, + function(ccursor) + if ccursor > 0 then + print(" "..(action_prompt or 'Please do the action.')) + end + if not utils.prompt_yes_no(' Continue?', true) then + return false + end + return true, nil, delta + end + ) +end + +return _ENV |
