aacomp.lua

function bt2aob(bt)
  local r = {}
  for i=1,#bt do r[i]=string.format('%02X',bt[i])end
  return table.concat(r,' ' )
end


local fmt, cat = string.format, table.concat

local loadcode = assert(load == loadstring and loadstring or load or loadstring,'no loadcode')
local function pcall_ret(ok, ...)if ok then return ... else return nil,... end end
local function ret_pcall(...)return pcall_ret(pcall(...))end
local function eval(chunks,ctx)
  local ok, r = pcall(loadcode, chunks, '_', nil, ctx or _G)
  if type(r)=='function'then return ret_pcall(r)else return nil, r end
end
local function eval_exp(exp,ctx)return ret_pcall(eval,'return '..exp,ctx)end

local function of_int(int,...)if int then return math.tointeger(int),of_int(...)end end

local function calc_rip(Adr, ofs_rip, ofs_end, rip_size, Read_Adr)
  ofs_rip, ofs_end, rip_size = of_int(ofs_rip or 4, ofs_end or 5, rip_size or 4) -- defaults
  if not (ofs_rip and ofs_end and rip_size)then
    return nil,'format: base ofs_rip ofs_end ?rip_size ?read_adr'
  end
  local adr, read_adr = getAddressSafe(Adr),getAddressSafe(Read_Adr)
  if not adr then return nil, Adr..' not known' end
  local read =
    rip_size == 4 and function(addr)return readInteger(addr,true)end or
    rip_size == 1 and function(addr)
      addr = readBytes(addr)
      if addr and addr>0x7f then addr = 0x100 - addr end
      return addr
    end
  if not read then return nil,'only 1 or 4 bytes rip'end
  local RIP = read_adr and read(read_adr)
  RIP = not RIP and not read_adr and read(adr + ofs_rip)
  if not RIP then return nil,'no RIP from '..Adr end
  return RIP + adr + ofs_end
end
_G.calc_rip = calc_rip

local cmd = 'COMP'
unregisterAutoAssemblerCommand(cmd)
registerAutoAssemblerCommand(cmd,function(s,sc)
  local sym, exp = s:match'^%s*([_#@%.%w]+)%s*,%s*(.-)%s*$'
  if not sym then exp = s end
  if not exp or exp:len()==0  then return nil,'empty expression'end
  local r, errmsg = eval_exp(exp)
  errmsg = not r and errmsg and tostring(errmsg) or 'unknown error'
  if sym and sym:len()>0 then
    if of_int(r)then
      return r<0 and fmt('define(%s,-%X)',sym,-r)or fmt('define(%s,0%X)',sym,r)
    elseif r then
      return fmt('define(%s,%s)',sym, tostring(r))
    elseif sc then
      return fmt('define(%s,666)',sym)
    end
  elseif not sc then
    return nil,'failed '..exp..' -> '..errmsg
  end
end)

function forget_fcomp(key)
  if type(key)=='string' and type(_G.__aa_comp_result)=='table'then
    _G.__aa_comp_result['&'..key] = nil
  end
end

local function assembled_bytes(start_addr, tmpl, ctx)
  local rid = '_'..math.random(10000,99999)..'_'
  local ret, prefix, post_code, subfix = {}, {}, {}
  local code = string.format([[
  %X:
  %s
  <END_BYTES>:]], start_addr, tmpl)
    :gsub('(&[_#@%.%w]+)', function(name)
      local r = ctx[name]
      if of_int(r)then return fmt('%X',r)end
     end)
    :gsub('<%s*([_%a][_%w]*)%s*>',function(sym)
      local rsym = rid..sym
      if not ret[sym] then
        local idx = 1+#ret
        ret[sym], ret[idx] = idx, sym
        prefix[idx] = fmt('registerSymbol(%s)',rsym)
        post_code[idx] = 'un'..prefix[idx]
      end
      return rsym
    end)
    :gsub('%(%s*%%([xXdDsS])%s+([_%a][_%w]*)%s*%)',function(fms, name)
      local v = ctx[name]
      if v then return fmt('%'..fms,v)end
    end)
  code = cat({cat(prefix,'\n'), code}, '\n')
  local ok = autoAssemble(code)
  if ok then
    for i=1,#ret do
      local sym, rsym = ret[i], rid .. ret[i]
      ret[sym] = getAddressSafe(rsym)
    end
    autoAssemble(cat(post_code,'\n'))-- unreg symbols
    if ret.END_BYTES then
      local len = ret.END_BYTES - start_addr
      local bt = readBytes(start_addr, len, true)
      return bt, len, ret
    end
  end
end



if next_rip_reg_asm_id then
  next_rip_reg_asm_id = nil, unregisterAssembler(next_rip_reg_asm_id)
end
if aa_comp_reg_asm_id then
  aa_comp_reg_asm_id = nil, unregisterAssembler(aa_comp_reg_asm_id)
end
local regs = 'rax;rbx;rcx;rdx;rsi;rdi;rbp;r11;r12;r13;r14;r15;'
aa_comp_reg_asm_id = registerAssembler(function(adr, inst)
  if adr < 0x1000 then return {0x90}end
  local op, ops, ctx = inst:match'^%s*([_%a][_%w]*)%s*(.-)%s*$'
  if not op or ops:len()==0 then return end -- don't know what to do, let ce handle
  while ops and #ops>0 and ops:match'^%b()$' do ops = ops:sub(2,-2)end
  op = op:lower()
  if op=='fcomp' then
    local reg, exp = ops:match'^&([_#@%.%w]+)%s*,%s*(.-)%s*$'
    if not reg or exp:len()==0 then return end
    local r, errmsg, addr = eval_exp(exp)
    if not of_int(r)then return end
    if #reg==3 and regs:find(reg) or reg=='_' then
      _G.one_off_reg_addr_setup = { reg = reg~='_' and reg,addr = r}
    elseif not _G.__aa_comp_result or
           not _G.__aa_comp_result[reg] -- don't overwrite existed, use forget_comp to refresh
    then
      local t = _G.__aa_comp_result
      t = type(t)=='table' and t or {}
      t['&'..reg] = r
      _G.__aa_comp_result = t
    end
    return {0x90}
  elseif _G.one_off_reg_addr_setup then
    ctx, _G.one_off_reg_addr_setup = _G.one_off_reg_addr_setup
    ctx.op, ctx.lhs, ctx.rhs = op, ops:match'^(.-%[)%s*.+%s*(%].-)$'
    if _G.__aa_comp_result and ops:find'&[_#@%.%w]'then
      for k,v in pairs(_G.__aa_comp_result)do if not ctx[k]then ctx[k] = v end end
    end
    if op=='jmp' or op=='call'then
      if ops:find(ctx.reg)then
        return assembled_bytes(adr, [[
            mov  (%s reg),(%X addr)
            (%s op)  (%s reg)
          ]], ctx)
      end
    elseif op:match'^j' then --- format 1, should be jcc
      return assembled_bytes(adr, [[
        (%s op) short <__yes_path>
        <no_path>:
        jmp short <__done>
        <__yes_path>:
        jmp  (%X addr)
        <__done>:
      ]], ctx)
    elseif ctx.lhs then -- format 2
      return assembled_bytes(adr, [[
        push     (%s reg)
          mov    (%s reg),(%X addr)
          (%s op) (%s lhs)(%s reg)(%s rhs)
        pop      (%s reg)
      ]], ctx)
    end
  elseif _G.__aa_comp_result and ops:find'&[_#@%.%w]'then
    return assembled_bytes(adr, inst, _G.__aa_comp_result)
  end
end)
---[===[

--- may place the following in a memory record, replace 'jmp %X' properly

cnt = type(cnt)=='number'and cnt+0x111111 or 0x10000000

local code = string.format([[
globalalloc(__X,$1000)
__X:
mov eax,%X
mov rax,-1

//// lua53-64.lua_absindex+B:
//// - 77 13                 - ja lua53-64.lua_absindex+20

reassemble(lua53-64.lua_absindex+B) /// ce seem do it not right, ja * -> jmp *,
                                    /// but my ce version may be outdated and it already been fixed.
fcomp &rax,calc_rip('lua53-64.lua_absindex+B',1,2,1)
///  if what follow & is a int register, then it is in ONE_OFF_mode,
///  the result addr only use in next instruction and forget;
///  otherwise, it is a VARIABLE <name> stored persisted,
///  next time (not just next) &<name> will be replaced
reassemble(lua53-64.lua_absindex+B)

mov rax,-1
fcomp &_,calc_rip('lua53-64.lua_absindex+B',1,2,1)
jna 1 /// jcc in ONE_OFF_mode can have a dummy target address, result addr will replace it
mov rax,-1

//// lua53-64.lua_version+5:
//// - 48 8D 05 FC7B0600     - lea rax,[lua53-64.dll.rdata+9F28] /// also displace as lua53-64.dll+68F28

reassemble(lua53-64.lua_version+5)
mov rax,-1
fcomp &r15,calc_rip('lua53-64.lua_version+5',3,7) /// note there is 2 number following base addr,
                                                  /// the 3rd rip-size default to 4 bytes
reassemble(lua53-64.lua_version+5) //// ONE_OFF_mode value can use in reassemble
mov rax,-1
fcomp &r15,calc_rip('lua53-64.lua_version+5',3,7)
mov    rbx,[rax] /// similar to jcc, anything between [*dummy] is replaced
mov   rax,-1
fcomp &VARNAME,calc_rip('lua53-64.lua_version+5',3,7) /// using store VARIABLE, they must be integer
fcomp &B1Bytes,0x66 /// prepare to use multiple store VARIABLE
fcomp &W2Bytes,0x666
cmp   [rax+&W2Bytes],&B1Bytes /// multiple VARIABLE, note that if some value is 64bit, most won't compile in this case
call  &VARNAME
mov   rax,-1
jmp   &VARNAME+&W2Bytes
{
call  &VARNAME
mov   rax,-1
jmp   &VARNAME
}
mov rax,-1
mov rax,-1
mov rax,-1
ret
@@:
mov rax,-1
mov rax,-1
comp(XXXX,calc_rip('lua53-64.lua_version+5',3,7)) /// --> define(XXXX,<calc-result>)
jmp XXXX
mov rax,-1
mov rax,-1
//
]], cnt)

autoAssembleCheck(code)
autoAssemble(code)
--[[  edited output
__X:
07CF0000- mov eax,10333333
07CF0005- mov rax,FFFFFFFFFFFFFFFF
07CF000F- nop
07CF0010- nop
07CF0011- jmp lua53-64.lua_absindex+20
07CF001F- nop
07CF0020- ja 07CF0024
07CF0022- jmp 07CF0032
07CF0024- jmp lua53-64.lua_absindex+20
07CF0032- mov rax,FFFFFFFFFFFFFFFF
07CF003C- nop
07CF003D- jna 07CF0041
07CF003F- jmp 07CF004F
07CF0041- jmp lua53-64.lua_absindex+20
07CF004F- mov rax,FFFFFFFFFFFFFFFF
07CF0059- mov rax,lua53-64.dll.rdata+9F28
07CF0063- mov rax,FFFFFFFFFFFFFFFF
07CF006D- nop
07CF006E- mov rax,lua53-64.dll.rdata+9F28
07CF0078- mov rax,FFFFFFFFFFFFFFFF
07CF0082- nop
07CF0083- push r15
07CF0085- mov r15,lua53-64.dll.rdata+9F28
07CF008F- mov rbx,[r15]
07CF0092- pop r15
07CF0094- mov rax,FFFFFFFFFFFFFFFF
07CF009E- nop
07CF009F- nop
07CF00A0- nop
07CF00A1- cmp dword ptr [rax+00000666],66
07CF00A8- call lua53-64.dll.rdata+9F28
07CF00B8- mov rax,FFFFFFFFFFFFFFFF
07CF00C2- jmp lua53-64.dll.rdata+A58E
07CF00D0- mov rax,FFFFFFFFFFFFFFFF
07CF00DA- mov rax,FFFFFFFFFFFFFFFF
07CF00E4- mov rax,FFFFFFFFFFFFFFFF
07CF00EE- ret
07CF00EF- mov rax,FFFFFFFFFFFFFFFF
07CF00F9- mov rax,FFFFFFFFFFFFFFFF
07CF0103- jmp lua53-64.dll.rdata+9F28
07CF0111- mov rax,FFFFFFFFFFFFFFFF
07CF011B- mov rax,FFFFFFFFFFFFFFFF
--]]