ipsw/pkg/dyld/disass.go

package dyld

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"path/filepath"
	"slices"
	"strings"

	"github.com/blacktop/arm64-cgo/disassemble"
	"github.com/blacktop/go-macho/pkg/swift"
	"github.com/blacktop/ipsw/internal/demangle"
	"github.com/blacktop/ipsw/pkg/disass"
	"github.com/blacktop/ipsw/pkg/symbols"
)

type DyldDisass struct {
	f       *File
	cfg     *disass.Config
	tr      *disass.Triage
	dylibs  []*CacheImage
	decoder disassemble.Decoder
}

func NewDyldDisass(f *File, cfg *disass.Config) *DyldDisass {
	return &DyldDisass{f: f, cfg: cfg}
}

func (d DyldDisass) Demangle() bool {
	return d.cfg.Demangle
}
func (d DyldDisass) Quite() bool {
	return d.cfg.Quiet
}
func (d DyldDisass) Color() bool {
	return d.cfg.Color
}
func (d DyldDisass) AsJSON() bool {
	return d.cfg.AsJSON
}
func (d DyldDisass) Data() []byte {
	return d.cfg.Data
}
func (d DyldDisass) StartAddr() uint64 {
	return d.cfg.StartAddress
}
func (d DyldDisass) Middle() uint64 {
	return d.cfg.Middle
}
func (d DyldDisass) ReadAddr(addr uint64) (uint64, error) {
	ptr, err := d.f.ReadPointerAtAddress(addr)
	if err != nil {
		return 0, err
	}
	return d.f.SlideInfo.SlidePointer(ptr), nil
}

func (d DyldDisass) Dylibs() []*CacheImage {
	return d.dylibs
}

func (d DyldDisass) hasDep(img *CacheImage) bool {
	return slices.Contains(d.dylibs, img)
}

// Contains returns true if Triage immediates contains a given address and will return the instruction address
func (d DyldDisass) Contains(address uint64) (bool, uint64) {
	for loc, addr := range d.tr.Addresses {
		if addr == address {
			return true, loc
		}
	}
	return false, 0
}

func (d DyldDisass) HasLoc(location uint64) (bool, uint64) {
	for loc, addr := range d.tr.Addresses {
		if loc == location {
			return true, addr
		}
	}
	return false, 0
}

// IsLocation returns if given address is a local branch location within the disassembled function
func (d DyldDisass) IsLocation(imm uint64) bool {
	if _, ok := d.tr.Locations[imm]; ok {
		return true
	}
	return false
}

// IsBranchLocation returns if given address is branch to a location instruction
func (d DyldDisass) IsBranchLocation(addr uint64) (bool, uint64) {
	for loc, addrs := range d.tr.Locations {
		if slices.Contains(addrs, addr) {
			return true, loc
		}
	}
	return false, 0
}

// IsData returns if given address is a data variable address referenced in the disassembled function
func (d DyldDisass) IsData(addr uint64) (bool, *disass.AddrDetails) {
	if detail, ok := d.tr.Details[addr]; ok {
		if detail.Flags.IsCstringLiterals() {
			return true, &detail
		}
		segment := strings.ToLower(detail.Segment)
		if strings.Contains(segment, "data") && !strings.EqualFold(detail.Section, "__got") {
			return true, &detail
		}
	}
	return false, nil
}

// IsPointer returns if given address is a pointer to another address
func (d DyldDisass) IsPointer(imm uint64) (bool, *disass.AddrDetails) {
	if deet, ok := d.tr.Details[imm]; ok {
		if deet.Pointer > 0 {
			return true, &deet
		}
	}
	return false, nil
}

func operandRegister(instr *disassemble.Inst, idx int) (disassemble.Register, bool) {
	if instr == nil || int(instr.NumOps) <= idx || instr.Operands[idx].NumRegisters == 0 {
		return disassemble.REG_NONE, false
	}
	return instr.Operands[idx].Registers[0], true
}

func operandImmediate(instr *disassemble.Inst, idx int) (uint64, bool) {
	if instr == nil || int(instr.NumOps) <= idx {
		return 0, false
	}
	return instr.Operands[idx].Immediate, true
}

// Triage walks a function and analyzes all immediates
func (d *DyldDisass) Triage() error {
	var instrValue uint32
	var prevInstr disassemble.Inst
	var hasPrev bool

	d.tr = &disass.Triage{
		Addresses: make(map[uint64]uint64),
		Locations: make(map[uint64][]uint64),
	}
	startAddr := d.StartAddr()
	r := bytes.NewReader(d.Data())

	// instructions, err := disassemble.GetInstructions(d.startAddr(), d.data())
	// if err != nil {
	// 	return err
	// }

	// for _, block := range instructions.Blocks() {
	// 	for _, i := range block {
	// 		fmt.Printf("%#08x:  %s\t%s\n", uint64(i.Address), disassemble.GetOpCodeByteString(i.Raw), i)
	// 	}
	// 	fmt.Println()
	// }
	// fmt.Println("DONE")

	// extract all immediates
	for {
		err := binary.Read(r, binary.LittleEndian, &instrValue)

		if err == io.EOF {
			break
		}

		var instruction disassemble.Inst
		if err := d.decoder.DecomposeInto(startAddr, instrValue, &instruction); err != nil {
			startAddr += uint64(binary.Size(uint32(0)))
			continue
		}

		if disass.IsBranchOp(instruction.Operation) {
			for idx := 0; idx < int(instruction.NumOps); idx++ {
				op := instruction.Operands[idx]
				if op.Class == disassemble.LABEL {
					d.tr.Addresses[instruction.Address] = uint64(op.Immediate)
				}
			}
		} else if disass.IsLoadLiteral(&instruction) {
			if imm, ok := operandImmediate(&instruction, 1); ok {
				d.tr.Addresses[instruction.Address] = imm
			}
		} else if hasPrev && prevInstr.Operation == disassemble.ARM64_ADRP &&
			(instruction.Operation == disassemble.ARM64_ADD ||
				instruction.Operation == disassemble.ARM64_LDR ||
				instruction.Operation == disassemble.ARM64_LDRB ||
				instruction.Operation == disassemble.ARM64_LDRSW ||
				instruction.Operation == disassemble.ARM64_STRB) {
			adrpRegister, ok := operandRegister(&prevInstr, 0)
			if ok {
				adrpImm, ok := operandImmediate(&prevInstr, 1)
				if ok {
					srcRegister, ok := operandRegister(&instruction, 1)
					if ok {
						validPattern := true
						if adrpRegister == srcRegister {
							switch instruction.Operation {
							case disassemble.ARM64_LDR, disassemble.ARM64_LDRB, disassemble.ARM64_LDRSW, disassemble.ARM64_STRB:
								if imm, ok := operandImmediate(&instruction, 1); ok {
									adrpImm += imm
								} else {
									validPattern = false
								}
							case disassemble.ARM64_ADD:
								if imm, ok := operandImmediate(&instruction, 2); ok {
									adrpImm += imm
								} else {
									validPattern = false
								}
							}
						}
						if validPattern {
							d.tr.Addresses[instruction.Address] = adrpImm
						}
					}
				}
			}
		}

		// lookup adrp/ldr or add address as a cstring or symbol name
		// if (operation == "ldr" || operation == "add") && prevInstruction.Operation().String() == "adrp" {
		// 	if operands := i.Instruction.Operands(); operands != nil && prevInstruction.Operands() != nil {
		// 		adrpRegister := prevInstruction.Operands()[0].Reg[0]
		// 		adrpImm := prevInstruction.Operands()[1].Immediate
		// 		if operation == "ldr" && adrpRegister == operands[1].Reg[0] {
		// 			adrpImm += operands[1].Immediate
		// 		} else if operation == "add" && adrpRegister == operands[1].Reg[0] {
		// 			adrpImm += operands[2].Immediate
		// 		}
		// 		d.tr.addresses[i.Instruction.Address()] = adrpImm
		// 	}

		// } else if i.Instruction.Group() == arm64.GROUP_BRANCH_EXCEPTION_SYSTEM { // check if branch location is a function
		// 	if operands := i.Instruction.Operands(); operands != nil {
		// 		for _, operand := range operands {
		// 			if operand.OpClass == arm64.LABEL {
		// 				d.tr.addresses[i.Instruction.Address()] = operand.Immediate
		// 			}
		// 		}
		// 	}
		// } else if i.Instruction.Group() == arm64.GROUP_DATA_PROCESSING_IMM || i.Instruction.Group() == arm64.GROUP_LOAD_STORE {
		// 	operation := i.Instruction.Operation()
		// 	if operation == arm64.ARM64_LDR || operation == arm64.ARM64_ADR {
		// 		if operands := i.Instruction.Operands(); operands != nil {
		// 			for _, operand := range operands {
		// 				if operand.OpClass == arm64.LABEL {
		// 					d.tr.addresses[i.Instruction.Address()] = operand.Immediate
		// 				}
		// 			}
		// 		}
		// 	}
		// }

		prevInstr = instruction
		hasPrev = true
		startAddr += uint64(binary.Size(uint32(0)))
	}

	if !d.Quite() {
		d.tr.Details = make(map[uint64]disass.AddrDetails)

		for addr, imm := range d.tr.Addresses {
			ptr := uint64(0)
			image, err := d.f.GetImageContainingVMAddr(imm)
			if err != nil {
				ptr, _ = d.f.ReadPointerAtAddress(imm)
				ptr = d.f.SlideInfo.SlidePointer(ptr)
				image, err = d.f.GetImageContainingVMAddr(ptr)
				if err != nil {
					continue
				}
			}

			if !d.hasDep(image) {
				d.dylibs = append(d.dylibs, image)
			}

			m, err := image.GetMacho()
			if err != nil {
				return err
			}
			defer m.Close()

			if fn, err := m.GetFunctionForVMAddr(addr); err == nil {
				d.tr.Function = &fn
				if d.tr.Function.StartAddr <= imm && imm < d.tr.Function.EndAddr {
					d.tr.Locations[imm] = append(d.tr.Locations[imm], addr)
					continue
				}
			}

			if ptr > 0 {
				if c := m.FindSectionForVMAddr(ptr); c != nil {
					d.tr.Details[imm] = disass.AddrDetails{
						Image:   filepath.Base(image.Name),
						Segment: c.Seg,
						Section: c.Name,
						Pointer: ptr,
						Flags:   c.Flags,
					}
				}
			} else {
				if c := m.FindSectionForVMAddr(imm); c != nil {
					d.tr.Details[imm] = disass.AddrDetails{
						Image:   filepath.Base(image.Name),
						Segment: c.Seg,
						Section: c.Name,
						Pointer: ptr,
						Flags:   c.Flags,
					}
				}
			}
		}
	}

	return nil
}

// ImageDependencies recursively returns all the image's loaded dylibs and those dylibs' loaded dylibs etc
// func (f *File) ImageDependencies(imageName string) error {

// 	image, err := f.Image(imageName)
// 	if err != nil {
// 		return err
// 	}

// 	m, err := image.GetPartialMacho()
// 	if err != nil {
// 		return err
// 	}
// 	defer m.Close()

// 	imports := m.ImportedLibraries()
// 	if len(imports) == 0 {
// 		return nil
// 	}

// 	for _, imp := range imports {
// 		if !utils.StrSliceHas(image.Analysis.Dependencies, imp) {
// 			image.Analysis.Dependencies = append(image.Analysis.Dependencies, imp)
// 			if err := f.ImageDependencies(imp); err != nil {
// 				return err
// 			}
// 			image.Analysis.Dependencies = utils.Unique(image.Analysis.Dependencies)
// 		}
// 	}

// 	return nil
// }

// IsFunctionStart checks if address is at a function start and returns symbol name
func (d DyldDisass) IsFunctionStart(addr uint64) (bool, string) {
	image, err := d.f.Image(d.cfg.Image)
	if err != nil {
		return false, err.Error()
	}
	m, err := image.GetMacho()
	if err != nil {
		return false, err.Error()
	}
	for _, fn := range m.GetFunctions() {
		if addr == fn.StartAddr {
			if symName, ok := d.f.AddressToSymbol.Get(addr); ok {
				if d.Demangle() {
					if strings.HasPrefix(symName, "_$s") { // TODO: better detect swift symbols
						symName, _ = swift.Demangle(symName)
						return ok, symName
					}
					return ok, demangle.Do(symName, false, false)
				}
				return ok, symName
			}
			return true, fmt.Sprintf("sub_%x", addr)
		}
	}
	return false, ""
}

// FindSymbol returns symbol from the addr2symbol map for a given virtual address
func (d DyldDisass) FindSymbol(addr uint64) (string, bool) {
	if symName, ok := d.f.AddressToSymbol.Get(addr); ok {
		return symbols.FormatSymbol(symName, d.cfg.Demangle), true
	}
	if symName, ok := d.resolveStubSymbol(addr); ok {
		return symbols.FormatSymbol(symName, d.cfg.Demangle), true
	}
	return "", false
}

func (d DyldDisass) resolveStubSymbol(addr uint64) (string, bool) {
	if d.f == nil {
		return "", false
	}
	if image, err := d.f.GetImageContainingTextAddr(addr); err == nil {
		if _, targetName, err := image.ResolveStubAtAddr(addr); err == nil && targetName != "" {
			symName := jumpSymbolName(targetName)
			d.f.AddressToSymbol.Set(addr, symName)
			return symName, true
		}
	}
	target, ok := d.stubIslandTarget(addr)
	if !ok {
		return "", false
	}
	if targetName, ok := d.resolveTargetSymbolName(target); ok {
		symName := jumpSymbolName(targetName)
		d.f.AddressToSymbol.Set(addr, symName)
		return symName, true
	}
	symName := fmt.Sprintf("%s%x", symbols.PrefixStubFallback, target)
	d.f.AddressToSymbol.Set(addr, symName)
	return symName, true
}

func (d DyldDisass) stubIslandTarget(addr uint64) (uint64, bool) {
	if len(d.f.islandStubs) == 0 {
		if err := d.f.ParseStubIslands(); err != nil {
			return 0, false
		}
	}
	target, ok := d.f.islandStubs[addr]
	if !ok {
		return 0, false
	}
	if _, _, err := d.f.GetOffset(target); err == nil || d.f.SlideInfo == nil {
		return target, true
	}
	return d.f.SlideInfo.SlidePointer(target), true
}

func (d DyldDisass) resolveTargetSymbolName(target uint64) (string, bool) {
	if symName, ok := d.f.AddressToSymbol.Get(target); ok && symName != "" {
		return symName, true
	}
	image, err := d.f.GetImageContainingTextAddr(target)
	if err != nil {
		return "", false
	}
	m, err := image.GetMacho()
	if err != nil {
		return "", false
	}
	defer m.Close()
	if syms, err := m.FindAddressSymbols(target); err == nil {
		for _, sym := range syms {
			if sym.Name != "" && sym.Name != "<redacted>" {
				return sym.Name, true
			}
		}
	}
	if name, err := image.FindLocalSymbolAtAddr(target); err == nil && name != "" {
		return name, true
	}
	return "", false
}

func jumpSymbolName(name string) string {
	name = strings.TrimPrefix(name, symbols.PrefixStubHelper)
	if strings.HasPrefix(name, symbols.PrefixJump) {
		return name
	}
	return symbols.PrefixJump + name
}

func (d DyldDisass) FindSwiftString(addr uint64) (string, bool) {
	if str, ok := d.f.AddressToSymbol.Get(addr); ok {
		return str, true
	}
	return "", false
}

func (d DyldDisass) GetCString(addr uint64) (string, error) {
	return d.f.GetCString(addr)
}