ebpf/ebpftracer/stack.go

package ebpftracer

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

	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/link"
	"github.com/coroot/coroot-node-agent/ebpftracer/tracer"
	"github.com/coroot/coroot-node-agent/proc"
	"golang.org/x/sync/semaphore"
)

func (t *Tracer) stack() error {
	// uprobes := []tracer.Uprobe{}
	if t.disableL7Tracing {
		return nil
	}

	ENV_PID := os.Getenv("FILTER_PID")
	WHITE_LIST := os.Getenv("WHITE_LIST")
	if ENV_PID == "" {
		return nil
	}

	MatchString := ".*HandleFunc|.*main.*|testfun.*|.*serverHandler.*|.*ServeHTTP.*"

	if WHITE_LIST != "" {
		MatchString = WHITE_LIST
	}

	fmt.Println("UprobesMatchString::: init", MatchString)

	pid, _ := strconv.ParseInt(ENV_PID, 10, 32)
	path := proc.Path(uint32(pid), "exe")

	t.Uprobes, _ = t.getUprobes(path, MatchString)
	t.UprobesMap = map[string]tracer.Uprobe{}
	fmt.Println("UprobesMap::: init")
	for _, up := range t.Uprobes {
		fmt.Println("UprobesMap:::", up.Funcname)
		t.UprobesMap[fmt.Sprintf("%s", up.Funcname)] = up
	}
	links := t.attachUprobes(path, t.Uprobes)
	t.links = append(t.links, links...)

	// defer t.detachUprobes(links)

	return nil
}

func (t *Tracer) 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 })

	t.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.Println("Fuck This: ", symbol.Name)
			// 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 (t *Tracer) attachUprobes(path string, uprobes []tracer.Uprobe) []link.Link {
	var links []link.Link

	ex, err := link.OpenExecutable(path)
	if err != nil {
		return nil
	}
	fmt.Println("AttachAttachAttach", path)

	for i, up := range t.uprobes {
		fmt.Println("attachingERROR---", i, up)
	}

	for i, up := range uprobes {
		fmt.Printf("attaching %d -> %d -> %s -> 0x%x -> 0x%x\n", i, len(uprobes), up.Funcname, up.AbsOffset, up.Address)
		var prog *ebpf.Program
		switch up.Location {
		case tracer.AtEntry:
			prog = t.uprobes["ent"]
		case tracer.AtRet:
			prog = t.uprobes["ret"]
		}
		fmt.Println("progprogprogprogprogprog---", prog)
		up, err := ex.Uprobe(up.Funcname, prog, &link.UprobeOptions{Address: up.Address, Offset: up.AbsOffset})
		if err != nil {
			fmt.Println("attachingERROR", err)
			// return nil
		} else {
			links = append(links, up)
		}
	}
	return links
}

func (t *Tracer) detachUprobes(links []link.Link) {
	sem := semaphore.NewWeighted(10)
	for i, closer := range links {
		fmt.Printf("detaching %d/%d\r", i+1, len(links))
		sem.Acquire(context.Background(), 1)
		go func(closer io.Closer) {
			defer sem.Release(1)
			closer.Close()
		}(closer)
	}
	fmt.Println()
}