v3 candidate
This commit is contained in:
Andy Dixon
@@ -0,0 +1,541 @@
|
||||
package disasm
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// cfg.go builds control-flow graphs and cross-references on top of the
|
||||
// per-architecture decoders. It performs a recursive-descent sweep from an
|
||||
// entry point: it follows direct branches and the fall-through after calls
|
||||
// (calls leave the function, so their targets are recorded as xrefs rather
|
||||
// than followed), splits the decoded instructions into basic blocks, and
|
||||
// records every direct code/data reference it can resolve statically.
|
||||
|
||||
// Block is one basic block: a maximal straight-line instruction run.
|
||||
type Block struct {
|
||||
Start int64 `json:"start"`
|
||||
End int64 `json:"end"` // address after the last instruction
|
||||
Instrs []map[string]any `json:"instrs"`
|
||||
Succs []int64 `json:"succs"` // successor block start addresses
|
||||
Kind string `json:"kind"` // ret | jmp | cond | call | fallthrough
|
||||
}
|
||||
|
||||
// Xref is a single static reference from one address to another.
|
||||
type Xref struct {
|
||||
From int64 `json:"from"`
|
||||
To int64 `json:"to"`
|
||||
Kind string `json:"kind"` // call | jump | data
|
||||
}
|
||||
|
||||
// CFG is the analysis result for one entry point.
|
||||
type CFG struct {
|
||||
Entry int64 `json:"entry"`
|
||||
Arch string `json:"arch"`
|
||||
Blocks []*Block `json:"blocks"`
|
||||
Calls []int64 `json:"calls"` // distinct call targets (in or out of range)
|
||||
Xrefs []Xref `json:"xrefs"` // all resolved references
|
||||
Unknown int `json:"unknown"` // indirect/unresolved branch+call count
|
||||
}
|
||||
|
||||
// Function is a discovered function and its structural metrics.
|
||||
type Function struct {
|
||||
Name string `json:"name"`
|
||||
Entry int64 `json:"entry"`
|
||||
End int64 `json:"end"`
|
||||
Blocks []*Block `json:"blocks"`
|
||||
Instructions int `json:"instructions"`
|
||||
Complexity int `json:"complexity"`
|
||||
Calls []int64 `json:"calls"`
|
||||
CalledBy []int64 `json:"calledBy"`
|
||||
Unknown int `json:"unknown"`
|
||||
}
|
||||
|
||||
// CallEdge is one direct call from a discovered function.
|
||||
type CallEdge struct {
|
||||
Caller int64 `json:"caller"`
|
||||
CallSite int64 `json:"callSite"`
|
||||
Callee int64 `json:"callee"`
|
||||
External bool `json:"external"`
|
||||
}
|
||||
|
||||
// Seed is a known or suspected function entry supplied by executable metadata.
|
||||
type Seed struct {
|
||||
Name string `json:"name"`
|
||||
Entry int64 `json:"entry"`
|
||||
}
|
||||
|
||||
// Program is a bounded recursive function-discovery result for one byte range.
|
||||
type Program struct {
|
||||
Entry int64 `json:"entry"`
|
||||
Base int64 `json:"base"`
|
||||
End int64 `json:"end"`
|
||||
Arch string `json:"arch"`
|
||||
Functions []*Function `json:"functions"`
|
||||
Edges []CallEdge `json:"edges"`
|
||||
ExternalCalls []int64 `json:"externalCalls"`
|
||||
Instructions int `json:"instructions"`
|
||||
Unknown int `json:"unknown"`
|
||||
Truncated bool `json:"truncated"`
|
||||
}
|
||||
|
||||
type rawInstr struct {
|
||||
addr int64
|
||||
length int
|
||||
mnem string
|
||||
class string
|
||||
target int64
|
||||
hasTgt bool
|
||||
data int64
|
||||
hasDat bool
|
||||
m map[string]any
|
||||
}
|
||||
|
||||
// BuildCFG decodes from entry (within buf, which starts at address base) and
|
||||
// returns the control-flow graph plus discovered references.
|
||||
func BuildCFG(buf []byte, base int64, arch string, entry int64, maxInstr int) *CFG {
|
||||
if maxInstr <= 0 {
|
||||
maxInstr = 8192
|
||||
}
|
||||
cfg := &CFG{Entry: entry, Arch: arch}
|
||||
decoded := map[int64]*rawInstr{}
|
||||
inRange := func(a int64) bool { return a >= base && a < base+int64(len(buf)) }
|
||||
|
||||
var stack []int64
|
||||
push := func(a int64) {
|
||||
if inRange(a) {
|
||||
stack = append(stack, a)
|
||||
}
|
||||
}
|
||||
push(entry)
|
||||
|
||||
for len(stack) > 0 && len(decoded) < maxInstr {
|
||||
addr := stack[len(stack)-1]
|
||||
stack = stack[:len(stack)-1]
|
||||
for inRange(addr) && len(decoded) < maxInstr {
|
||||
if _, seen := decoded[addr]; seen {
|
||||
break
|
||||
}
|
||||
ri := decodeAt(buf, addr, base, arch)
|
||||
if ri == nil {
|
||||
break
|
||||
}
|
||||
decoded[addr] = ri
|
||||
next := addr + int64(ri.length)
|
||||
switch ri.class {
|
||||
case "ret":
|
||||
addr = -1
|
||||
case "jmp":
|
||||
if ri.hasTgt {
|
||||
addr = ri.target
|
||||
} else {
|
||||
cfg.Unknown++
|
||||
addr = -1
|
||||
}
|
||||
case "cond":
|
||||
if ri.hasTgt {
|
||||
push(ri.target)
|
||||
} else {
|
||||
cfg.Unknown++
|
||||
}
|
||||
addr = next
|
||||
case "call":
|
||||
if !ri.hasTgt {
|
||||
cfg.Unknown++
|
||||
}
|
||||
addr = next
|
||||
default:
|
||||
addr = next
|
||||
}
|
||||
if addr < 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Collect references and call targets.
|
||||
callSet := map[int64]bool{}
|
||||
for _, ri := range decoded {
|
||||
if ri.hasTgt {
|
||||
switch ri.class {
|
||||
case "call":
|
||||
cfg.Xrefs = append(cfg.Xrefs, Xref{From: ri.addr, To: ri.target, Kind: "call"})
|
||||
if !callSet[ri.target] {
|
||||
callSet[ri.target] = true
|
||||
cfg.Calls = append(cfg.Calls, ri.target)
|
||||
}
|
||||
case "jmp", "cond":
|
||||
cfg.Xrefs = append(cfg.Xrefs, Xref{From: ri.addr, To: ri.target, Kind: "jump"})
|
||||
}
|
||||
}
|
||||
if ri.hasDat {
|
||||
cfg.Xrefs = append(cfg.Xrefs, Xref{From: ri.addr, To: ri.data, Kind: "data"})
|
||||
}
|
||||
}
|
||||
sort.Slice(cfg.Calls, func(i, j int) bool { return cfg.Calls[i] < cfg.Calls[j] })
|
||||
sort.Slice(cfg.Xrefs, func(i, j int) bool {
|
||||
if cfg.Xrefs[i].From != cfg.Xrefs[j].From {
|
||||
return cfg.Xrefs[i].From < cfg.Xrefs[j].From
|
||||
}
|
||||
return cfg.Xrefs[i].To < cfg.Xrefs[j].To
|
||||
})
|
||||
|
||||
cfg.Blocks = buildBlocks(decoded, inRange)
|
||||
return cfg
|
||||
}
|
||||
|
||||
// DiscoverFunctions recursively follows direct calls from entry, building a
|
||||
// bounded call graph over buf. Indirect calls are counted but cannot be
|
||||
// discovered statically. maxInstr is a total budget across all functions.
|
||||
func DiscoverFunctions(buf []byte, base int64, arch string, entry int64, maxFunctions, maxInstr int) *Program {
|
||||
return DiscoverFunctionsFromSeeds(buf, base, arch, []Seed{{Entry: entry}}, maxFunctions, maxInstr)
|
||||
}
|
||||
|
||||
// DiscoverFunctionsFromSeeds analyzes named and unnamed roots plus every
|
||||
// direct in-range call target reached from them.
|
||||
func DiscoverFunctionsFromSeeds(buf []byte, base int64, arch string, seeds []Seed, maxFunctions, maxInstr int) *Program {
|
||||
if maxFunctions <= 0 || maxFunctions > 4096 {
|
||||
maxFunctions = 256
|
||||
}
|
||||
if maxInstr <= 0 || maxInstr > 1_000_000 {
|
||||
maxInstr = 65_536
|
||||
}
|
||||
end := base + int64(len(buf))
|
||||
entry := base
|
||||
if len(seeds) > 0 {
|
||||
entry = seeds[0].Entry
|
||||
}
|
||||
program := &Program{Entry: entry, Base: base, End: end, Arch: arch}
|
||||
inRange := func(addr int64) bool { return addr >= base && addr < end }
|
||||
queue := make([]int64, 0, len(seeds))
|
||||
queued := map[int64]bool{}
|
||||
names := map[int64]string{}
|
||||
for _, seed := range seeds {
|
||||
if !inRange(seed.Entry) {
|
||||
continue
|
||||
}
|
||||
if !queued[seed.Entry] {
|
||||
queued[seed.Entry] = true
|
||||
queue = append(queue, seed.Entry)
|
||||
}
|
||||
if seed.Name != "" {
|
||||
names[seed.Entry] = seed.Name
|
||||
}
|
||||
}
|
||||
if len(queue) == 0 {
|
||||
program.Truncated = true
|
||||
return program
|
||||
}
|
||||
|
||||
discovered := map[int64]*Function{}
|
||||
callerSets := map[int64]map[int64]bool{}
|
||||
externalSet := map[int64]bool{}
|
||||
|
||||
for len(queue) > 0 && len(discovered) < maxFunctions && program.Instructions < maxInstr {
|
||||
fnEntry := queue[0]
|
||||
queue = queue[1:]
|
||||
if discovered[fnEntry] != nil {
|
||||
continue
|
||||
}
|
||||
remaining := maxInstr - program.Instructions
|
||||
perFunction := remaining
|
||||
if perFunction > 8192 {
|
||||
perFunction = 8192
|
||||
}
|
||||
cfg := BuildCFG(buf, base, arch, fnEntry, perFunction)
|
||||
fn := summarizeFunction(cfg, names[fnEntry])
|
||||
discovered[fnEntry] = fn
|
||||
program.Functions = append(program.Functions, fn)
|
||||
program.Instructions += fn.Instructions
|
||||
program.Unknown += fn.Unknown
|
||||
|
||||
for _, xref := range cfg.Xrefs {
|
||||
if xref.Kind != "call" {
|
||||
continue
|
||||
}
|
||||
external := !inRange(xref.To)
|
||||
program.Edges = append(program.Edges, CallEdge{
|
||||
Caller: fnEntry, CallSite: xref.From, Callee: xref.To, External: external,
|
||||
})
|
||||
if external {
|
||||
externalSet[xref.To] = true
|
||||
continue
|
||||
}
|
||||
if callerSets[xref.To] == nil {
|
||||
callerSets[xref.To] = map[int64]bool{}
|
||||
}
|
||||
callerSets[xref.To][fnEntry] = true
|
||||
if !queued[xref.To] {
|
||||
queued[xref.To] = true
|
||||
queue = append(queue, xref.To)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(queue) > 0 || program.Instructions >= maxInstr {
|
||||
program.Truncated = true
|
||||
}
|
||||
for _, fn := range program.Functions {
|
||||
for caller := range callerSets[fn.Entry] {
|
||||
fn.CalledBy = append(fn.CalledBy, caller)
|
||||
}
|
||||
sort.Slice(fn.CalledBy, func(i, j int) bool { return fn.CalledBy[i] < fn.CalledBy[j] })
|
||||
}
|
||||
for target := range externalSet {
|
||||
program.ExternalCalls = append(program.ExternalCalls, target)
|
||||
}
|
||||
sort.Slice(program.Functions, func(i, j int) bool { return program.Functions[i].Entry < program.Functions[j].Entry })
|
||||
sort.Slice(program.Edges, func(i, j int) bool {
|
||||
if program.Edges[i].Caller != program.Edges[j].Caller {
|
||||
return program.Edges[i].Caller < program.Edges[j].Caller
|
||||
}
|
||||
if program.Edges[i].CallSite != program.Edges[j].CallSite {
|
||||
return program.Edges[i].CallSite < program.Edges[j].CallSite
|
||||
}
|
||||
return program.Edges[i].Callee < program.Edges[j].Callee
|
||||
})
|
||||
sort.Slice(program.ExternalCalls, func(i, j int) bool { return program.ExternalCalls[i] < program.ExternalCalls[j] })
|
||||
return program
|
||||
}
|
||||
|
||||
func summarizeFunction(cfg *CFG, name string) *Function {
|
||||
if name == "" {
|
||||
name = unknownFunctionName(cfg.Entry)
|
||||
}
|
||||
fn := &Function{
|
||||
Name: name, Entry: cfg.Entry, Blocks: cfg.Blocks,
|
||||
Calls: append([]int64(nil), cfg.Calls...), Unknown: cfg.Unknown,
|
||||
}
|
||||
edges := 0
|
||||
for _, block := range cfg.Blocks {
|
||||
if block.End > fn.End {
|
||||
fn.End = block.End
|
||||
}
|
||||
fn.Instructions += len(block.Instrs)
|
||||
edges += len(block.Succs)
|
||||
}
|
||||
fn.Complexity = edges - len(cfg.Blocks) + 2
|
||||
if fn.Complexity < 1 {
|
||||
fn.Complexity = 1
|
||||
}
|
||||
return fn
|
||||
}
|
||||
|
||||
func unknownFunctionName(entry int64) string {
|
||||
return "sub_" + strconv.FormatInt(entry, 16)
|
||||
}
|
||||
|
||||
func buildBlocks(decoded map[int64]*rawInstr, inRange func(int64) bool) []*Block {
|
||||
if len(decoded) == 0 {
|
||||
return nil
|
||||
}
|
||||
addrs := make([]int64, 0, len(decoded))
|
||||
for a := range decoded {
|
||||
addrs = append(addrs, a)
|
||||
}
|
||||
sort.Slice(addrs, func(i, j int) bool { return addrs[i] < addrs[j] })
|
||||
|
||||
// Leaders: first instr, every branch target, and the instruction after any
|
||||
// branch/call/ret.
|
||||
leaders := map[int64]bool{addrs[0]: true}
|
||||
for _, a := range addrs {
|
||||
ri := decoded[a]
|
||||
next := a + int64(ri.length)
|
||||
switch ri.class {
|
||||
case "ret":
|
||||
if _, ok := decoded[next]; ok {
|
||||
leaders[next] = true
|
||||
}
|
||||
case "jmp":
|
||||
if ri.hasTgt {
|
||||
if _, ok := decoded[ri.target]; ok {
|
||||
leaders[ri.target] = true
|
||||
}
|
||||
}
|
||||
if _, ok := decoded[next]; ok {
|
||||
leaders[next] = true
|
||||
}
|
||||
case "cond", "call":
|
||||
if ri.hasTgt {
|
||||
if _, ok := decoded[ri.target]; ok {
|
||||
leaders[ri.target] = true
|
||||
}
|
||||
}
|
||||
if _, ok := decoded[next]; ok {
|
||||
leaders[next] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var blocks []*Block
|
||||
var cur *Block
|
||||
flush := func() {
|
||||
if cur != nil {
|
||||
blocks = append(blocks, cur)
|
||||
cur = nil
|
||||
}
|
||||
}
|
||||
for idx, a := range addrs {
|
||||
ri := decoded[a]
|
||||
if leaders[a] {
|
||||
flush()
|
||||
cur = &Block{Start: a}
|
||||
}
|
||||
if cur == nil {
|
||||
cur = &Block{Start: a}
|
||||
}
|
||||
cur.Instrs = append(cur.Instrs, ri.m)
|
||||
next := a + int64(ri.length)
|
||||
cur.End = next
|
||||
|
||||
// Decide whether this instruction terminates the block.
|
||||
isLast := idx == len(addrs)-1
|
||||
nextIsLeader := false
|
||||
if !isLast {
|
||||
if leaders[addrs[idx+1]] {
|
||||
nextIsLeader = true
|
||||
}
|
||||
}
|
||||
terminator := ri.class == "ret" || ri.class == "jmp" || ri.class == "cond"
|
||||
if terminator || nextIsLeader || isLast {
|
||||
cur.Kind = blockKind(ri.class)
|
||||
switch ri.class {
|
||||
case "ret":
|
||||
// no successors
|
||||
case "jmp":
|
||||
if ri.hasTgt && inRange(ri.target) {
|
||||
cur.Succs = []int64{ri.target}
|
||||
}
|
||||
case "cond":
|
||||
if ri.hasTgt && inRange(ri.target) {
|
||||
cur.Succs = append(cur.Succs, ri.target)
|
||||
}
|
||||
if _, ok := decoded[next]; ok {
|
||||
cur.Succs = append(cur.Succs, next)
|
||||
}
|
||||
default:
|
||||
if _, ok := decoded[next]; ok {
|
||||
cur.Succs = []int64{next}
|
||||
}
|
||||
}
|
||||
flush()
|
||||
}
|
||||
}
|
||||
flush()
|
||||
return blocks
|
||||
}
|
||||
|
||||
func blockKind(class string) string {
|
||||
switch class {
|
||||
case "ret", "jmp", "cond", "call":
|
||||
return class
|
||||
default:
|
||||
return "fallthrough"
|
||||
}
|
||||
}
|
||||
|
||||
// decodeAt decodes a single instruction at address addr (buf starts at base).
|
||||
func decodeAt(buf []byte, addr, base int64, arch string) *rawInstr {
|
||||
off := int(addr - base)
|
||||
if off < 0 || off >= len(buf) {
|
||||
return nil
|
||||
}
|
||||
end := off + 16 // max x86 instruction length; fixed-width arches need <= 4
|
||||
if end > len(buf) {
|
||||
end = len(buf)
|
||||
}
|
||||
res := Disassemble(buf[off:end], addr, arch, 1)
|
||||
if len(res) == 0 {
|
||||
return nil
|
||||
}
|
||||
m := res[0]
|
||||
ri := &rawInstr{addr: addr, m: m}
|
||||
ri.length = byteCount(m["bytes"].(string))
|
||||
if ri.length <= 0 {
|
||||
ri.length = 1
|
||||
}
|
||||
ri.mnem, _ = m["mnemonic"].(string)
|
||||
ops, _ := m["operands"].(string)
|
||||
ri.class = classify(ri.mnem)
|
||||
if t, ok := directTarget(ops); ok {
|
||||
ri.target, ri.hasTgt = t, true
|
||||
}
|
||||
if dataReferencing(ri.class) {
|
||||
if d, ok := dataTarget(ops); ok {
|
||||
ri.data, ri.hasDat = d, true
|
||||
}
|
||||
}
|
||||
return ri
|
||||
}
|
||||
|
||||
func byteCount(s string) int {
|
||||
if s == "" {
|
||||
return 0
|
||||
}
|
||||
return len(strings.Fields(s))
|
||||
}
|
||||
|
||||
func classify(mnem string) string {
|
||||
switch mnem {
|
||||
case "ret", "retf", "iret":
|
||||
return "ret"
|
||||
case "jmp", "jmpf", "b", "br":
|
||||
return "jmp"
|
||||
case "call", "callf", "bl", "blr", "jal", "jalr":
|
||||
return "call"
|
||||
case "cbz", "cbnz", "tbz", "tbnz", "loop", "loope", "loopne", "jecxz", "jrcxz":
|
||||
return "cond"
|
||||
}
|
||||
if strings.HasPrefix(mnem, "b.") { // AArch64 conditional branch
|
||||
return "cond"
|
||||
}
|
||||
if strings.HasPrefix(mnem, "j") { // x86 Jcc (jmp already handled above)
|
||||
return "cond"
|
||||
}
|
||||
if strings.HasPrefix(mnem, "beq") || strings.HasPrefix(mnem, "bne") { // MIPS
|
||||
return "cond"
|
||||
}
|
||||
return "normal"
|
||||
}
|
||||
|
||||
func dataReferencing(class string) bool {
|
||||
// RIP-relative resolved targets show up as "; 0x.." comments on any
|
||||
// instruction; only mine them on non-control-flow ops as data xrefs.
|
||||
return class == "normal" || class == "call"
|
||||
}
|
||||
|
||||
// directTarget extracts a branch/call target when the final operand is a bare
|
||||
// absolute address (a direct, statically known target). Register/memory
|
||||
// indirect operands return false.
|
||||
func directTarget(operands string) (int64, bool) {
|
||||
if i := strings.Index(operands, ";"); i >= 0 {
|
||||
operands = operands[:i]
|
||||
}
|
||||
if c := strings.LastIndex(operands, ","); c >= 0 {
|
||||
operands = operands[c+1:]
|
||||
}
|
||||
operands = strings.TrimSpace(operands)
|
||||
return parseHexAddr(operands)
|
||||
}
|
||||
|
||||
// dataTarget extracts a resolved address from a "; 0x.." trailing comment
|
||||
// (e.g. an x86 RIP-relative memory operand).
|
||||
func dataTarget(operands string) (int64, bool) {
|
||||
i := strings.Index(operands, ";")
|
||||
if i < 0 {
|
||||
return 0, false
|
||||
}
|
||||
return parseHexAddr(strings.TrimSpace(operands[i+1:]))
|
||||
}
|
||||
|
||||
func parseHexAddr(s string) (int64, bool) {
|
||||
if !strings.HasPrefix(s, "0x") {
|
||||
return 0, false
|
||||
}
|
||||
v, err := strconv.ParseUint(s[2:], 16, 64)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return int64(v), true
|
||||
}
|
||||
Reference in New Issue
Block a user