mipsevm: LoadUnicorn, update hook mem-write and hook code

mipsevm/patch.go

@@ -8,23 +8,8 @@ import (
 	"io"
 )
 
-type state struct {
-	// TODO VM state
-
-	pc uint32
-}
-
-func (s *state) SetMemRange(addr uint32, r io.Reader) error {
-	// TODO write memory
-	return nil
-}
-
-func (s *state) SetRegister(i uint8, v uint32) {
-	// TODO
-}
-
-func LoadELF(f *elf.File) (*state, error) {
-	s := &state{}
+func LoadELF(f *elf.File) (*State, error) {
+	s := &State{}
 
 	for i, prog := range f.Progs {
 		if prog.Type == 0x70000003 { // MIPS_ABIFLAGS
@@ -47,7 +32,7 @@ func LoadELF(f *elf.File) (*state, error) {
 		if prog.Vaddr+prog.Memsz >= uint64(1<<32) {
 			return nil, fmt.Errorf("program %d out of 32-bit mem range: %x - %x (size: %x)", i, prog.Vaddr, prog.Vaddr+prog.Memsz, prog.Memsz)
 		}
-		if err := s.SetMemRange(uint32(prog.Vaddr), r); err != nil {
+		if err := s.SetMemoryRange(uint32(prog.Vaddr), r); err != nil {
 			return nil, fmt.Errorf("failed to read program segment %d: %w", i, err)
 		}
 	}
@@ -55,7 +40,7 @@ func LoadELF(f *elf.File) (*state, error) {
 	return s, nil
 }
 
-func patchVM(f *elf.File, st *state) error {
+func patchVM(f *elf.File, st *State) error {
 	symbols, err := f.Symbols()
 	if err != nil {
 		return fmt.Errorf("failed to read symbols data, cannot patch program: %w", err)
@@ -74,14 +59,14 @@ func patchVM(f *elf.File, st *state) error {
 			// MIPS32 patch: ret (pseudo instruction)
 			// 03e00008 = jr $ra = ret (pseudo instruction)
 			// 00000000 = invalid, make sure it never enters the actual function
-			if err := st.SetMemRange(uint32(s.Value), bytes.NewReader([]byte{
+			if err := st.SetMemoryRange(uint32(s.Value), bytes.NewReader([]byte{
 				0x03, 0xe0, 0x00, 0x08,
 				0, 0, 0, 0,
 			})); err != nil {
 				return fmt.Errorf("failed to patch Go runtime.gcenable: %w", err)
 			}
 		case "runtime.MemProfileRate":
-			if err := st.SetMemRange(uint32(s.Value), bytes.NewReader(make([]byte, 4))); err != nil { // disable mem profiling, to avoid a lot of unnecessary floating point ops
+			if err := st.SetMemoryRange(uint32(s.Value), bytes.NewReader(make([]byte, 4))); err != nil { // disable mem profiling, to avoid a lot of unnecessary floating point ops
 				return err
 			}
 		}
@@ -89,12 +74,12 @@ func patchVM(f *elf.File, st *state) error {
 
 	// setup stack pointer
 	sp := uint32(0xc0_00_00_00)
-	st.SetRegister(29, sp)
+	st.Registers[29] = sp
 
 	storeMem := func(addr uint32, v uint32) {
 		var dat [4]byte
 		binary.BigEndian.PutUint32(dat[:], v)
-		_ = st.SetMemRange(addr, bytes.NewReader(dat[:]))
+		_ = st.SetMemoryRange(addr, bytes.NewReader(dat[:]))
 	}
 
 	// init argc, argv, aux on stack
@@ -107,7 +92,7 @@ func patchVM(f *elf.File, st *state) error {
 	storeMem(sp+4*7, sp+4*9) // auxv[3] = address of 16 bytes containing random value
 	storeMem(sp+4*8, 0)      // auxv[term] = 0
 
-	_ = st.SetMemRange(sp+4*9, bytes.NewReader([]byte("4;byfairdiceroll"))) // 16 bytes of "randomness"
+	_ = st.SetMemoryRange(sp+4*9, bytes.NewReader([]byte("4;byfairdiceroll"))) // 16 bytes of "randomness"
 
 	return nil
 }

mipsevm/state.go

@@ -8,8 +8,10 @@ import (
 
 const (
 	pageAddrSize = 10
+	pageKeySize  = 32 - pageAddrSize
 	pageSize     = 1 << pageAddrSize
 	pageAddrMask = pageSize - 1
+	maxPageCount = 1 << pageKeySize
 )
 
 type Page [pageSize]byte
@@ -51,14 +53,6 @@ type State struct {
 // Misc exit/step data = TBD
 // + proof(s) for memory leaf nodes
 
-func (s *State) ApplyRegisterDiff(regs [32]uint32, hi, lo uint32) {
-	for i := 0; i < 32; i++ {
-		s.Registers[i] = regs[i]
-	}
-	s.Hi = hi
-	s.Lo = lo
-}
-
 func (s *State) SetMemory(addr uint32, v uint32, size uint32) {
 	for i := size; i > 0; i-- {
 		pageIndex := addr >> pageAddrSize
@@ -74,6 +68,23 @@ func (s *State) SetMemory(addr uint32, v uint32, size uint32) {
 	}
 }
 
+func (s *State) SetMemoryRange(addr uint32, r io.Reader) error {
+	for {
+		pageIndex := addr >> pageAddrSize
+		pageAddr := addr & pageAddrMask
+		p, ok := s.Memory[pageIndex]
+		if !ok {
+			p = &Page{}
+			s.Memory[pageIndex] = p
+		}
+		n, err := r.Read(p[pageAddr:])
+		if err != nil {
+			return err
+		}
+		addr += uint32(n)
+	}
+}
+
 type memReader struct {
 	state *State
 	addr  uint32
@@ -107,7 +118,7 @@ func (r *memReader) Read(dest []byte) (n int, err error) {
 	return n, nil
 }
 
-func (s *State) ReadMemory(addr uint32, count uint32) io.Reader {
+func (s *State) ReadMemoryRange(addr uint32, count uint32) io.Reader {
 	return &memReader{state: s, addr: addr, count: count}
 }
 

mipsevm/unicorn.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"math"
 
 	uc "github.com/unicorn-engine/unicorn/bindings/go/unicorn"
 )
@@ -13,8 +14,28 @@ func NewUnicorn() (uc.Unicorn, error) {
 }
 
 func LoadUnicorn(st *State, mu uc.Unicorn) error {
-	// TODO mmap each page of state into unicorn
-	// TODO set all registers
+	// mmap and write each page of memory state into unicorn
+	for pageIndex, page := range st.Memory {
+		addr := uint64(pageIndex) << pageAddrSize
+		if err := mu.MemMap(addr, pageSize); err != nil {
+			return fmt.Errorf("failed to mmap page at addr 0x%x: %w", addr, err)
+		}
+		if err := mu.MemWrite(addr, page[:]); err != nil {
+			return fmt.Errorf("failed to write page at addr 0x%x: %w", addr, err)
+		}
+	}
+	// write all registers into unicorn, including PC, LO, HI
+	regValues := make([]uint64, 32+3)
+	// TODO: do we have to sign-extend registers before writing them to unicorn, or are the trailing bits unused?
+	for i, v := range st.Registers {
+		regValues[i] = uint64(v)
+	}
+	regValues[32] = uint64(st.PC)
+	regValues[33] = uint64(st.Lo)
+	regValues[34] = uint64(st.Hi)
+	if err := mu.RegWriteBatch(regBatchKeys(), regValues); err != nil {
+		return fmt.Errorf("failed to write registers: %w", err)
+	}
 	return nil
 }
 
@@ -33,9 +54,9 @@ func HookUnicorn(st *State, mu uc.Unicorn, stdOut, stdErr io.Writer) error {
 			count, _ := mu.RegRead(uc.MIPS_REG_A2)
 			switch fd {
 			case 1:
-				_, _ = io.Copy(stdOut, st.ReadMemory(uint32(addr), uint32(count)))
+				_, _ = io.Copy(stdOut, st.ReadMemoryRange(uint32(addr), uint32(count)))
 			case 2:
-				_, _ = io.Copy(stdErr, st.ReadMemory(uint32(addr), uint32(count)))
+				_, _ = io.Copy(stdErr, st.ReadMemoryRange(uint32(addr), uint32(count)))
 			default:
 				// ignore other output data
 			}
@@ -82,19 +103,31 @@ func HookUnicorn(st *State, mu uc.Unicorn, stdOut, stdErr io.Writer) error {
 	}
 
 	_, err = mu.HookAdd(uc.HOOK_MEM_WRITE, func(mu uc.Unicorn, access int, addr64 uint64, size int, value int64) {
-		//rt := value
-		//rs := addr64 & 3
-		//addr := uint32(addr64 & 0xFFFFFFFC)
-		// TODO write to state memory
+		if addr64 > math.MaxUint32 {
+			panic("invalid addr")
+		}
+		if size < 0 || size > 4 {
+			panic("invalid mem size")
+		}
+		st.SetMemory(uint32(addr64), uint32(size), uint32(value))
 	}, 0, 0x80000000)
 	if err != nil {
 		return fmt.Errorf("failed to set up mem-write hook: %w", err)
 	}
 
+	regBatch := regBatchKeys()
 	_, err = mu.HookAdd(uc.HOOK_CODE, func(mu uc.Unicorn, addr uint64, size uint32) {
-		steps += 1
-
-		// TODO: diff all registers
+		st.Step += 1
+		batch, err := mu.RegReadBatch(regBatch)
+		if err != nil {
+			panic(fmt.Errorf("failed to read register batch: %w", err))
+		}
+		for i := 0; i < 32; i++ {
+			st.Registers[i] = uint32(batch[i])
+		}
+		st.PC = uint32(batch[32])
+		st.Lo = uint32(batch[33])
+		st.Hi = uint32(batch[34])
 	}, 0, 0x80000000)
 	if err != nil {
 		return fmt.Errorf("failed to set up instruction hook: %w", err)
@@ -102,3 +135,12 @@ func HookUnicorn(st *State, mu uc.Unicorn, stdOut, stdErr io.Writer) error {
 
 	return nil
 }
+
+func regBatchKeys() []int {
+	var batch []int
+	for i := 0; i < 32; i++ {
+		batch = append(batch, uc.MIPS_REG_ZERO+i)
+	}
+	batch = append(batch, uc.MIPS_REG_PC, uc.MIPS_REG_LO, uc.MIPS_REG_HI)
+	return batch
+}