optimize shellcode and recover main address at runtime

macho-go/pkg/ios/macho/edit.go

@@ -408,14 +408,12 @@ func (mc *MachoContext) ReworkForObjc() {
 
 			// __bss section is dynamically allocated at the end to or something, hmmge
 			// assume that it order correctly, which it should if compiled and not modified
+			// each section has their addr field which we can use that with segment virtual address
+			// to calculate the offset of the last section from segment starts
+			// then use the size of section to calculate the end of segment in file
 			sections := segment.Sections()
 			last := sections[len(sections)-1]
-			data_end = int(last.Offset()) + int(last.Size())
-
-			if last.Offset() == 0 {
-				before_last := sections[len(sections)-2]
-				data_end += int(before_last.Offset()) + int(before_last.Size())
-			}
+			data_end = int(last.Addr() - segment.Vmaddr() + segment.Fileoff() + last.Size())
 		}
 		ptr += int64(cmd.Cmdsize())
 	}
@@ -431,62 +429,54 @@ func (mc *MachoContext) ReworkForObjc() {
 	// we need to store the return address, and parameters passed to main
 	// we also store our header address to not calculate many times
 
-	/*
-	  adr x8, 0
-	  sub sp, sp, #0x30
-	  str x30, [sp]
-	  movz x9, #0x3d68 ; offset at this point
-	  sub x8, x8, x9
-	  str x8, [sp, #0x8]
-	  str x0, [sp, #0x10]
-	  str x1, [sp, #0x18]
-	  str x2, [sp, #0x20]
-	  str x3, [sp, #0x28]
+	// must expand stack to store arguments passed
+	// must use newly allocated stack region
+	// must save return address
+	// must recover stack to same value before calling main
+	// must recover link register before calling main
 
-	  movz x9, #0x81d8 ; offset to end of __DATA
-	  add x9, x8, x9
-	  ldr x9, [x9]
-	  blr x9
-	  ldr x8, [sp, #0x8]
-	  ldr x0, [sp, #0x10]
-	  ldr x1, [sp, #0x18]
-	  ldr x2, [sp, #0x20]
-	  ldr x3, [sp, #0x28]
-	  movz x9, #0x3e3c  ; offset to original main
-	  add x9, x8, x9
-	  blr x9
-	  ldr x30, [sp]
-	  add sp, sp, #0x10
-	  ret
+	// we use shorthand store/load multiple
+	// arm also has different indexing instruction, so be careful
+	// https://developer.arm.com/documentation/102374/0101/Loads-and-stores---addressing
+	/*
+  adr x8, 0
+  # x9 = (offset end of __DATA) - (offset shellcode)
+  movz x9, #0x9999
+  add x8, x8, x9
+
+  stp x30, x8, [sp], #-0x10
+  stp x3, x2, [sp], #-0x10
+  stp x1, x0, [sp], #-0x10
+
+  # custom intializer
+  ldr x9, [x8]
+  blr x9
+
+  ldp x1, x0, [sp, #0x10]!
+  ldp x3, x2, [sp, #0x10]!
+  ldp x30, x8, [sp, #0x10]!
+
+  # original main
+  # link register is set so jump only
+  ldr x9, [x8, #8]
+  br x9
 	*/
 
 	// TODO: fix to work with offset larger than 0xffff
 	shellcode := []uint32{
 		0x10000008,
-		0xD100C3FF,
-		0xF90003FE,
-		0, // movz_shellcode_offset,
-		0xCB090108,
-		0xF90007E8,
-		0xF9000BE0,
-		0xF9000FE1,
-		0xF90013E2,
-		0xF90017E3,
-		0, // movz_data_end_offset,
-		0x8B090109,
-		0xF9400129,
+		0, // x9 = (offset end of __DATA) - (offset shellcode)
+		0x8B090108,
+		0xA8BF23FE,
+		0xA8BF0BE3,
+		0xA8BF03E1,
+		0xF9400109,
 		0xD63F0120,
-		0xF94007E8,
-		0xF9400BE0,
-		0xF9400FE1,
-		0xF94013E2,
-		0xF94017E3,
-		0, // movz_main_offset,
-		0x8B090109,
-		0xD63F0120,
-		0xF94003FE,
-		0x910043FF,
-		0xD65F03C0,
+		0xA9C103E1,
+		0xA9C10BE3,
+		0xA9C123FE,
+		0xF9400509,
+		0xD61F0120,
 	}
 
 	ins_size_byte := 4
@@ -497,20 +487,22 @@ func (mc *MachoContext) ReworkForObjc() {
 		return uint32(uint32(v)<<5 | uint32(0x694)<<21 | uint32(0x09))
 	}
 
-	movz_shellcode_offset := encode_movz(shellcode_offset)
-	movz_main_offset := encode_movz(main_offset)
-	movz_data_end_offset := encode_movz(data_end)
+	// movz_shellcode_offset := encode_movz(shellcode_offset)
+	// movz_main_offset := encode_movz(main_offset)
+	// movz_data_end_offset := encode_movz(data_end)
+	movz_calculate_offset := encode_movz(data_end - shellcode_offset)
 
-	shellcode[3] = movz_shellcode_offset
-	shellcode[10] = movz_data_end_offset
-	shellcode[19] = movz_main_offset
+	shellcode[1] = movz_calculate_offset
+	// shellcode[10] = movz_data_end_offset
+	// shellcode[19] = movz_main_offset
 
 	fmt.Printf("// shellcode_offset=%x\n", shellcode_offset)
 	fmt.Printf("// main_offset=%x\n", main_offset)
 	fmt.Printf("// data_end=%x\n", data_end)
-	fmt.Printf("// movz_shellcode_offset=%x\n", movz_shellcode_offset)
-	fmt.Printf("// movz_main_offset=%x\n", movz_main_offset)
-	fmt.Printf("// movz_data_end_offset=%x\n", movz_data_end_offset)
+	fmt.Printf("// movz_calculate_offset=%x\n", movz_calculate_offset)
+	// fmt.Printf("// movz_shellcode_offset=%x\n", movz_shellcode_offset)
+	// fmt.Printf("// movz_main_offset=%x\n", movz_main_offset)
+	// fmt.Printf("// movz_data_end_offset=%x\n", movz_data_end_offset)
 	fmt.Printf("// lc_main_offset=%x\n", lc_main_offset)
 
 	offset := int64(shellcode_offset)

macho-go/pkg/ios/macho/macho.go

@@ -64,6 +64,10 @@ func (mc *MachoContext) Segments() []Segment {
 	return mc.segments
 }
 
+func (mc *MachoContext) Main() uint64 {
+	return mc.entryoff
+}
+
 func (mc *MachoContext) WriteEnabled() bool {
 	return mc.file != nil
 }