ebpf/containers/stack.go

package containers

import (
	"debug/dwarf"
	"debug/elf"
	debugelf "debug/elf"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sort"
	"strconv"

	"github.com/coroot/coroot-node-agent/ebpftracer"
	"github.com/coroot/coroot-node-agent/ebpftracer/tracer"
	tracerelf "github.com/coroot/coroot-node-agent/ebpftracer/tracer"
	"github.com/coroot/coroot-node-agent/proc"
	"golang.org/x/arch/arm64/arm64asm"
	"golang.org/x/arch/x86/x86asm"
)

type uprobesDef struct {
	Name       string
	Offset     uint64
	EntAddress uint64
	RetAddress uint64
}

func (c *Container) stackTrace(tracer *ebpftracer.Tracer, pid uint32) error {
	ENV_PID := os.Getenv("FILTER_PID")
	if ENV_PID != "" {
		filterPid, _ := strconv.ParseInt(ENV_PID, 10, 64)
		if filterPid != int64(pid) {
			return nil
		}
	}
	binType := "dotnet"
	MatchString := ".*HandleFunc|.*main.*|testfun.*|.*serverHandler.*|.*ServeHTTP.*"
	dbgpath := ""

	WHITE_LIST := os.Getenv("WHITE_LIST")
	BIN_TYPE := os.Getenv("BIN_TYPE")
	DBG_PATH := os.Getenv("DBG_PATH")
	if WHITE_LIST != "" {
		MatchString = WHITE_LIST
	}

	if DBG_PATH != "" {
		dbgpath = DBG_PATH
	}

	if BIN_TYPE != "" {
		binType = BIN_TYPE
	}

	fmt.Println("UprobesMatchString:::init", MatchString)
	path := proc.Path(uint32(pid), "exe")

	if dbgpath != "" {
		c.Uprobes, _ = c.getJavaAOTUprobes(binType, path, dbgpath, MatchString)
	} else {
		c.Uprobes, _ = c.getUprobes(path, MatchString)
	}

	p := c.processes[uint32(pid)]
	if p == nil {
		return nil
	}

	if p.stackUprobesChecked {
		return nil
	}

	c.UprobesMap = map[string]tracerelf.Uprobe{}
	fmt.Println("UprobesMap:::init")
	for _, up := range c.Uprobes {
		fmt.Println("UprobesMap:::", up.Funcname, up.Address, up.AbsOffset)
		c.UprobesMap[fmt.Sprintf("%s-%s", up.Funcname, up.Address+up.AbsOffset)] = up
	}
	codeType := c.GetCodeTypeFromCache(pid)
	tracer.InitKProcInfo(pid, c.instanceID, uint16(codeType))
	p.uprobes = append(p.uprobes, tracer.AttachStackUprobes(path, c.Uprobes)...)
	p.stackUprobesChecked = true

	return nil
}

func (c *Container) getJavaAOTUprobes(binType, path string, dbgpath string, MatchString string) ([]tracer.Uprobe, error) {
	uprobes := []tracer.Uprobe{}

	elfFile, err := elf.Open(path)

	funSection := ".text"

	if binType == "dotnet" {
		funSection = "__managedcode"
	}

	textSection := elfFile.Section(funSection)
	if textSection == nil {
		fmt.Println("no text section", nil)
		return nil, nil
	}
	textSectionData, err := textSection.Data()
	if err != nil {
		fmt.Println("failed to read text section", err)
		return nil, nil
	}
	textSectionLen := uint64(len(textSectionData) - 1)

	dwarfFile, err := elf.Open(dbgpath)

	if err != nil {
		log.Fatal(err)
	}

	dwarfData, err := dwarfFile.DWARF()
	if err != nil {
		log.Fatal(err)
	}

	entryReader := dwarfData.Reader()

	// var targetAddress uint64

	listEntry := make(map[dwarf.Offset]uprobesDef)
	SpecListEntry := []dwarf.Entry{}

	for {
		entry, err := entryReader.Next()
		if err == io.EOF {
			// We've reached the end of DWARF entries
			break
		}
		if err != nil {
			log.Fatalf("Error reading entry: %v", err)
		}
		if entry == nil {
			log.Println("Warning: a nil entry was returned with no error")
			break
		}
		if entry.Tag == dwarf.TagSubprogram {
			// fmt.Printf("entry address: %x, %d\n", entry.Offset, entry.Children)
			funName, _ := entry.Val(dwarf.AttrName).(string)
			found, _ := regexp.MatchString(MatchString, funName)
			if found {
				entAddress, _ := entry.Val(dwarf.AttrLowpc).(uint64)
				retAddress, _ := entry.Val(dwarf.AttrHighpc).(uint64)
				// fmt.Printf("Function %s address: %x, %x\n", funName, address, entry.Offset)
				uprobes := uprobesDef{}
				uprobes.EntAddress = entAddress
				uprobes.RetAddress = retAddress
				uprobes.Offset = uint64(entry.Offset)
				uprobes.Name = funName

				listEntry[entry.Offset] = uprobes
			}

			specAddr, _ := entry.Val(dwarf.AttrSpecification).(dwarf.Offset)
			lowpc := entry.Val(dwarf.AttrLowpc)
			if lowpc != nil && specAddr > 0 && lowpc.(uint64) > 0 {
				// fmt.Printf("AttrSpecification address: %x, %x\n", specAddr, entry.Offset)
				SpecListEntry = append(SpecListEntry, *entry)
			}
		}
	}

	for _, v := range SpecListEntry {
		specAddr, _ := v.Val(dwarf.AttrSpecification).(dwarf.Offset)
		// fmt.Printf("SpecListEntrySpecListEntrySpecListEntry Attach Function: %x\n", specAddr)
		_, ok := listEntry[specAddr]
		if ok {
			vv := listEntry[specAddr]
			entAddr := v.Val(dwarf.AttrLowpc)
			if entAddr != nil {
				vv.EntAddress = entAddr.(uint64)
			}
			retAddr := v.Val(dwarf.AttrHighpc)
			if retAddr != nil {
				switch retAddr.(type) {
				case uint64:
					vv.RetAddress = uint64(retAddr.(uint64))
				case int64:
					vv.RetAddress = uint64(retAddr.(int64))
				default:
					fmt.Println("Unknown type")
				}
			}
			listEntry[specAddr] = vv
		}
	}

	for _, v := range listEntry {
		fmt.Printf("Need Attach Function %s address: %x, %x\n", v.Name, v.EntAddress, v.RetAddress)

		sStart := v.EntAddress - textSection.Addr
		sSize := v.RetAddress
		if v.RetAddress > v.EntAddress {
			sSize = v.RetAddress - v.EntAddress
		}
		sEnd := sStart + sSize
		if sEnd > textSectionLen {
			continue
		}
		sBytes := textSectionData[sStart:sEnd]
		rbpOffsets := getRbpEnterOffsets(elfFile.Machine, sBytes)
		returnOffsets := getReturnOffsets(elfFile.Machine, sBytes)

		if rbpOffsets != 0 {
			uprobes = append(uprobes, tracer.Uprobe{
				Funcname:  v.Name,               // 函数名
				Location:  tracer.AtDotNetEntry, // 入口
				Address:   v.EntAddress,         // 函数地址
				AbsOffset: uint64(rbpOffsets),   // 函数相对 ELF 偏移
				RelOffset: 0,                    // 函数真实偏移
			})
		} else {
			// 函数入口加入待 attach 列表
			uprobes = append(uprobes, tracer.Uprobe{
				Funcname:  v.Name,         // 函数名
				Location:  tracer.AtEntry, // 入口
				Address:   v.EntAddress,   // 函数地址
				AbsOffset: 0,              // 函数相对 ELF 偏移
				RelOffset: 0,              // 函数真实偏移
			})
		}

		for _, offset := range returnOffsets {
			uprobes = append(uprobes, tracer.Uprobe{
				Funcname:  v.Name,
				Location:  tracer.AtRet,
				Address:   v.EntAddress,
				AbsOffset: uint64(offset),
				RelOffset: 0,
			})
		}
	}
	return uprobes, nil
}

func (c *Container) getUprobes(path string, MatchString string) ([]tracer.Uprobe, error) {
	uprobes := []tracer.Uprobe{}

	binFile, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	// cache := map[string]interface{}{}
	// 解析 elf 文件
	elfFile, _ := debugelf.NewFile(binFile)
	// 获取所有符号表
	symbols, _ := elfFile.Symbols()

	sort.Slice(symbols, func(i, j int) bool { return symbols[i].Value < symbols[j].Value })

	c.Symbols = symbols
	// 符号表组装成键值 map，方便使用
	symnames := map[string]debugelf.Symbol{}
	for _, symbol := range symbols {
		fmt.Println(symbol.Name, symbol)
		symnames[symbol.Name] = symbol
	}

	textSection := elfFile.Section(".text")
	if textSection == nil {
		fmt.Println("no text section", nil)
		return nil, nil
	}
	textSectionData, err := textSection.Data()
	if err != nil {
		fmt.Println("failed to read text section", err)
		return nil, nil
	}
	textSectionLen := uint64(len(textSectionData) - 1)

	//  遍历符号表
	for _, symbol := range symbols {
		if debugelf.ST_TYPE(symbol.Info) != debugelf.STT_FUNC {
			continue
		}
		// fmt.Println("Hello FunName: ", symbol.Name)
		// 使用正则表达式匹配函数白名单列表
		found, err := regexp.MatchString(MatchString, symbol.Name)
		// found, err := regexp.MatchString("main.*", symbol.Name)

		if err != nil {
			log.Fatal(err)
		}

		if found {
			// 匹配到了加入 attachFuncs 列表
			fmt.Printf("Fuck This: %s, %x\n", symbol.Name, symbol.Value)
			// attachFuncs = append(attachFuncs, symbol.Name)
			// 根据函数名拿到当前函数的符号结构体
			sym := symnames[symbol.Name]
			if err != nil {
				fmt.Printf("symnames[symbol.Name]", symbol.Name, err)
				return nil, err
			}

			address := sym.Value
			for _, p := range elfFile.Progs {
				if p.Type != elf.PT_LOAD || (p.Flags&elf.PF_X) == 0 {
					continue
				}

				if p.Vaddr <= sym.Value && sym.Value < (p.Vaddr+p.Memsz) {
					address = sym.Value - p.Vaddr + p.Off
					break
				}
			}

			// 函数入口加入待 attach 列表
			uprobes = append(uprobes, tracer.Uprobe{
				Funcname:  symbol.Name,    // 函数名
				Location:  tracer.AtEntry, // 入口
				Address:   address,        // 函数地址
				AbsOffset: 0,              // 函数相对 ELF 偏移
				RelOffset: 0,              // 函数真实偏移
				Wanted:    true,
			})

			sStart := sym.Value - textSection.Addr
			sEnd := sStart + sym.Size
			if sEnd > textSectionLen {
				continue
			}
			sBytes := textSectionData[sStart:sEnd]
			returnOffsets := getReturnOffsets(elfFile.Machine, sBytes)

			for _, offset := range returnOffsets {
				uprobes = append(uprobes, tracer.Uprobe{
					Funcname:  symbol.Name,
					Location:  tracer.AtRet,
					Address:   address,
					AbsOffset: uint64(offset),
					RelOffset: 0,
				})
			}
		}
	}
	return uprobes, nil
}

func getRbpEnterOffsets(machine elf.Machine, instructions []byte) int {
	switch machine {
	case elf.EM_X86_64:
		for i := 0; i < len(instructions); {
			ins, err := x86asm.Decode(instructions[i:], 64)
			if err == nil && ins.Op == x86asm.LEA && ins.Args[0].String() == "RBP" {
				fmt.Printf("getRbpEnterOffsets: %v, %s, %s\n", ins, ins.Args[0].String(), ins.Args[1].String())
				return i
			}
			i += ins.Len
		}
	case elf.EM_AARCH64:
		for i := 0; i < len(instructions); {
			ins, err := arm64asm.Decode(instructions[i:])
			if err == nil && ins.Op == arm64asm.RET {
				return i
			}
			i += 4
		}
	}
	return 0
}

func getReturnOffsets(machine elf.Machine, instructions []byte) []int {
	var res []int
	switch machine {
	case elf.EM_X86_64:
		for i := 0; i < len(instructions); {
			ins, err := x86asm.Decode(instructions[i:], 64)
			if err == nil && ins.Op == x86asm.RET {
				res = append(res, i)
			}
			i += ins.Len
		}
	case elf.EM_AARCH64:
		for i := 0; i < len(instructions); {
			ins, err := arm64asm.Decode(instructions[i:])
			if err == nil && ins.Op == arm64asm.RET {
				res = append(res, i)
			}
			i += 4
		}
	}
	return res
}