factorio-ios/tests/check_a64_program.py

"""Compare direct guest calls, nested LR saves, recursion, tails and memory with ARM64."""
import json
import struct
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/'tools'))
from a64_program import ProgramCompiler
from a64_lift import UnsupportedInstruction

OUT = ROOT/'build/program'
OUT.mkdir(parents=True,exist_ok=True)


def bl(site,target): return 0x94000000 | (((target-site)//4)&0x3ffffff)
def body(words): return struct.pack('<'+'I'*len(words),*words)
def run(*args): subprocess.run(list(map(str,args)),cwd=ROOT,check=True)


fixtures = {
 'nested': (0x1000, {
  0x1000:[0xa9bf7bf3,0xaa0003f3,bl(0x1008,0x1100),bl(0x100c,0x1200),0x8b130000,0xa8c17bf3,0xd65f03c0],
  0x1100:[0x91000400,0xd65f03c0],
  0x1200:[0xa9bf7bf4,bl(0x1204,0x1100),0xa8c17bf4,0x91000800,0xd65f03c0]}),
 'recursive': (0x2000, {0x2000:[0xa9bf7bf3,0xb4000080,0xd1000400,
                            bl(0x200c,0x2000),0x91000400,0xa8c17bf3,0xd65f03c0]}),
 'tail': (0x3000, {0x3000:[0x14000040],0x3100:[0x91001c00,0xd65f03c0]}),
 'memory': (0x4000, {
  0x4000:[0xa9bf7bf3,0xaa0003f3,bl(0x4008,0x4100),0xaa1303e0,bl(0x4010,0x4100),0xa8c17bf3,0xd65f03c0],
  0x4100:[0xf9400008,0x91000508,0xf9000008,0xaa0803e0,0xd65f03c0]})}

assembly = ['.text','.p2align 2']
cursor = 0x1000
for name,(entry,functions) in fixtures.items():
    compiler = ProgramCompiler()
    wat = compiler.compile({address:body(words) for address,words in functions.items()},entry)
    (OUT/f'{name}.wat').write_text(wat)
    run(ROOT/'build/wabt/wat2wasm',OUT/f'{name}.wat','-o',OUT/f'{name}.wasm')
    run(ROOT/'build/wabt/wasm2c',OUT/f'{name}.wasm','-o',OUT/f'{name}.c','--module-name='+name)
    for address,words in sorted(functions.items()):
        assert address>=cursor
        assembly.append(f'.space {address-cursor}')
        if address==entry: assembly += [f'.globl _native_{name}',f'_native_{name}:']
        assembly += [f'.long 0x{word:08x}' for word in words]
        cursor=address+4*len(words)
(OUT/'oracle.S').write_text('\n'.join(assembly)+'\n')
bad_return = ProgramCompiler().compile({0x5000:body([0xd282469e,0xd65f03c0])},0x5000)
(OUT/'badreturn.wat').write_text(bad_return)
run(ROOT/'build/wabt/wat2wasm',OUT/'badreturn.wat','-o',OUT/'badreturn.wasm')
run(ROOT/'build/wabt/wasm2c',OUT/'badreturn.wasm','-o',OUT/'badreturn.c','--module-name=badreturn')
for functions,entry in [({0x1000:body([bl(0x1000,0x1100),0xd65f03c0])},0x1000),
                       ({0x1000:body([0xd65f03c0]),0x1100:body([0xd65f03c0])},0x1000),
                       ({0x1000:body([0x14000000])},0x1000),
                       ({0x1000:body([0x54000040,0xd2800000,0xd65f03c0])},0x1000)]:
    try: ProgramCompiler().compile(functions,entry)
    except UnsupportedInstruction: pass
    else: raise AssertionError('Invalid guest program accepted')

includes = '\n'.join(f'#include "{name}.h"\nextern uint64_t native_{name}(uint64_t);' for name in fixtures)
checks=[]
for name in fixtures:
    argument = 'i%51' if name=='recursive' else 'next(&seed)'
    prepare = 'memcpy(memory+2048,&expected,8); registers[0]=2048;' if name=='memory' else ''
    callarg = '(uint64_t)&expected' if name=='memory' else 'input'
    memory_check = 'assert(memcmp(memory+2048,&expected,8)==0);' if name=='memory' else ''
    checks.append(f'''
 {{
  w2c_{name} module; wasm2c_{name}_instantiate(&module);
  uint8_t *memory=w2c_{name}_memory(&module)->data;
  for(unsigned i=0;i<20000;++i) {{
   uint64_t input={argument}, expected=input, registers[32], after[32];
   for(unsigned r=0;r<32;++r) registers[r]=next(&seed);
   registers[0]=input; registers[30]=0; registers[31]=65280;
   memset(memory,0xa5,65536);
   {prepare}
   memcpy(memory,registers,sizeof(registers));
   uint64_t native=native_{name}({callarg});
   assert(w2c_{name}_run(&module,0,1000)==0);
   memcpy(after,memory,sizeof(after));
   assert(after[0]==native);
   for(unsigned r=19;r<32;++r) assert(after[r]==registers[r]);
   {memory_check}
   for(unsigned a=256;a<65536;++a) {{
    if(a>=64464 && a<65280) continue;
    if(a>=2048 && a<2056) continue;
    assert(memory[a]==0xa5);
   }}
  }}
  uint64_t start[32]={{0}}; start[0]=20; start[31]=65280;
  memcpy(memory,start,sizeof(start));
  assert(w2c_{name}_run(&module,0,0)==3);
  memcpy(memory,start,sizeof(start));
  assert(w2c_{name}_run(&module,0,1)==3);
  memcpy(memory,start,sizeof(start));
  start[30]={fixtures[name][0]}; memcpy(memory,start,sizeof(start));
  assert(w2c_{name}_run(&module,0,1000)==2);
  wasm2c_{name}_free(&module);
 }}
''')
(OUT/'verify.c').write_text(f'''
#include <assert.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
{includes}
#include "badreturn.h"
static uint64_t next(uint64_t *s) {{ *s^=*s<<13; *s^=*s>>7; *s^=*s<<17; return *s; }}
int main(void) {{
 wasm_rt_init(); uint64_t seed=0x123456789abcdef;
 {''.join(checks)}
 w2c_nested bounded; wasm2c_nested_instantiate(&bounded);
 uint64_t registers[32]={{0}}; registers[0]=3; registers[31]=65280;
 memcpy(w2c_nested_memory(&bounded)->data,registers,sizeof(registers));
 assert(w2c_nested_run(&bounded,0,15)==3);
 memcpy(w2c_nested_memory(&bounded)->data,registers,sizeof(registers));
 assert(w2c_nested_run(&bounded,0,16)==0);
 memcpy(registers,w2c_nested_memory(&bounded)->data,sizeof(registers));
 assert(registers[0]==10);
 wasm2c_nested_free(&bounded);
 w2c_badreturn bad; wasm2c_badreturn_instantiate(&bad);
 assert(w2c_badreturn_run(&bad,0,10)==4);
 wasm2c_badreturn_free(&bad);
 wasm_rt_free();
 puts("PASS: 80000 native guest-program comparisons, nested calls, recursion, tails, shared memory, guards and failure statuses");
 return 0;
}}
''')
runtime=ROOT/'third_party/wabt/wasm2c'
run('xcrun','clang','-O2','-DWASM_RT_USE_MMAP=0','-DWASM_RT_MEMCHECK_BOUNDS_CHECK=1',
    '-I'+str(OUT),'-I'+str(runtime),*(OUT/f'{name}.c' for name in [*fixtures,'badreturn']),
    OUT/'verify.c',OUT/'oracle.S',runtime/'wasm-rt-impl.c',runtime/'wasm-rt-mem-impl.c','-o',OUT/'verify')
run(OUT/'verify')
(OUT/'verification.json').write_text(json.dumps({'status':'pass','native_comparisons':80000,
    'source':'hand-authored instruction fixtures, not recovered game code','fixtures':list(fixtures),
    'scope':'bounded direct guest calls, known returns, shared registers/memory'},indent=2)+'\n')