main.go

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"log"
	"os"
	"os/signal"
	"syscall"
	"fmt"
	"flag"
	"encoding/json"

	"github.com/cilium/ebpf/link"
	"github.com/cilium/ebpf/ringbuf"
	"github.com/cilium/ebpf/rlimit"

	"github.com/davecgh/go-spew/spew"
)

const ChunkSize = 0x400
const MaxChunks = 16

type Filter struct {
	Interface string `json:"interface"`
	Method    string `json:"method"`
}

type Config struct {
	PackageName  string          `json:"package_name"`
	Args         bool            `json:"args"`
	FileName     string          `json:"file_name"`
	CustomMethods MethodMappings `json:"customMethods"`
	Filters      []Filter        `json:"filters"`
}

func parseFlags() (Config, string) {
	var config Config
	var configFile string

	flag.StringVar(&configFile, "c", "", "specify the configuration file")
	flag.StringVar(&config.PackageName, "p", "", "specify the package name to trace")
	flag.BoolVar(&config.Args, "a", false, "whether to trace arguments")
	flag.StringVar(&config.FileName, "f", "", "specify the log output file")

	flag.Parse()

	return config, configFile
}

func parseConfigFile(filename string) (Config, error) {
	var config Config
	file, err := os.Open(filename)
	if err != nil {
		return config, err
	}
	defer file.Close()

	decoder := json.NewDecoder(file)
	if err := decoder.Decode(&config); err != nil {
		return config, err
	}

	return config, nil
}

func mergeConfigs(flagConfig, fileConfig Config) Config {
	if flagConfig.PackageName != "" {
		fileConfig.PackageName = flagConfig.PackageName
	}
	if flagConfig.Args {
		fileConfig.Args = flagConfig.Args
	}
	if flagConfig.FileName != "" {
		fileConfig.FileName = flagConfig.FileName
	}
	return fileConfig
}

func mergeChunks(chunks [][]byte) []byte {
	mergedData := make([]byte, 0, len(chunks)*ChunkSize)

    for _, chunk := range chunks {
        mergedData = append(mergedData, chunk...)
    }

    return mergedData
}

func shouldFilter(interfaceToken, methodName string, filters []Filter) bool {
	for _, filter := range filters {
		if filter.Interface == interfaceToken && (filter.Method == "" || filter.Method == methodName) {
			return true
		}
	}
	return false
}

func main() {

	// 解析命令行参数
	flagConfig, configFile := parseFlags()

	// 解析配置文件
	var fileConfig Config
	if configFile != "" {
		var err error
		fileConfig, err = parseConfigFile(configFile)
		if err != nil {
			log.Fatalf("读取配置文件失败: %v", err)
		}
	}
	
	// 合并命令行参数和配置文件参数
	conf := mergeConfigs(flagConfig, fileConfig)

	var logger *log.Logger
	if conf.FileName != "" {
		file, err := os.OpenFile(conf.FileName, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
		if err != nil {
			fmt.Println("无法打开日志文件:", err)
			return
		}
		defer file.Close()
		logger = log.New(file, "", log.Ldate|log.Ltime|log.Lshortfile)
	} else {
		// 如果未指定日志文件，则将日志输出到控制台
		logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile)
	}

	// Name of the kernel function to trace.
	fn := "binder_transaction"

	// Subscribe to signals for terminating the program.
	stopper := make(chan os.Signal, 1)
	signal.Notify(stopper, os.Interrupt, syscall.SIGTERM)

	// Load default method mappings from embedded JSON
	if err := LoadMethodMappings(methodsJSON); err != nil {
		logger.Fatal(err)
	}

	// Override default method mappings with custom methods from config file
	OverrideMethodMappings(conf.CustomMethods)

	if err := LoadPackageMappings(); err != nil {
		logger.Fatal(err)
	}

	// Allow the current process to lock memory for eBPF resources.
	if err := rlimit.RemoveMemlock(); err != nil {
		logger.Fatal(err)
	}

	// Load pre-compiled programs and maps into the kernel.
	objs := bpfObjects{}
	if err := loadBpfObjects(&objs, nil); err != nil {
		logger.Fatalf("loading objects: %v", err)
	}
	defer objs.Close()

	var config_key uint32 = 0;
	config_value := bpfTraceConfig{};
	uid, err := GetUidByPackageName(conf.PackageName)
	if err != nil {
		uid = 0;
	}
	config_value.Uid = uint32(uid);
	if err := objs.TraceConfigMap.Put(config_key, config_value); err != nil {
		logger.Fatalf("写入 BPF 配置映射失败: %v", err)
	}

	// Open a Kprobe at the entry point of the kernel function and attach the
	// pre-compiled program. Each time the kernel function enters, the program
	// will emit an event containing pid and command of the execved task.
	kp, err := link.Kprobe(fn, objs.KprobeBinderTransaction, nil)
	if err != nil {
		logger.Fatalf("opening kprobe: %s", err)
	}
	defer kp.Close()

	// Open a ringbuf reader from userspace RINGBUF map described in the
	// eBPF C program.
	rd, err := ringbuf.NewReader(objs.TraceEventMap)
	if err != nil {
		logger.Fatalf("opening ringbuf reader: %s", err)
	}
	defer rd.Close()

	// Close the reader when the process receives a signal, which will exit
	// the read loop.
	go func() {
		<-stopper

		if err := rd.Close(); err != nil {
			logger.Fatalf("closing ringbuf reader: %s", err)
		}
	}()

	logger.Println("Waiting for events..")

	// bpfEvent is generated by bpf2go.
	var event bpfTraceEvent
	var transactionBuffers = make(map[uint64][][]byte)
	for {
		record, err := rd.Read()
		if err != nil {
			if errors.Is(err, ringbuf.ErrClosed) {
				logger.Println("Received signal, exiting..")
				return
			}
			logger.Printf("reading from reader: %s", err)
			continue
		}

		// Parse the ringbuf event entry into a bpfEvent structure.
		if err := binary.Read(bytes.NewBuffer(record.RawSample), binary.LittleEndian, &event); err != nil {
			logger.Printf("parsing ringbuf event: %s", err)
			continue
		}

		if _, exists := transactionBuffers[event.TransactionId]; !exists {
			var totalChunks = (event.DataSize + ChunkSize - 1) / ChunkSize
			transactionBuffers[event.TransactionId] = make([][]byte, totalChunks)
		}
		transactionBuffers[event.TransactionId][event.ChunkIndex] = append([]byte(nil), event.ChunkData[:]...)
		
		var complete = true
		for _, chunk := range transactionBuffers[event.TransactionId] {
			if chunk == nil {
				complete = false
				break
			}
		}
		
		if complete {
			completeData := mergeChunks(transactionBuffers[event.TransactionId])
			delete(transactionBuffers, event.TransactionId)

			if len(completeData) <= 16 {
				continue
			}
			if len(completeData) > ChunkSize*MaxChunks {
				continue
			}

			parcelData := completeData[:event.DataSize-1]

			interfaceToken, err := ExtractInterfaceName(parcelData)

			if err != nil {
				logger.Println("Error parsing parcel:", err)
				continue
			}

			methodName, err := GetMethodName(interfaceToken, int(event.Code))
			if err != nil {
				methodName = fmt.Sprintf("%d", event.Code)
			}

			if shouldFilter(interfaceToken, methodName, conf.Filters) {
				continue
			}

			packageName, err := GetPackageNameByUid(int(event.Uid))
			if err != nil {
				packageName = fmt.Sprintf("%d", event.Uid)
			}

			if conf.Args {
				logger.Printf("(pid:%d, uid:%d, package:%s) -> (interface:%s, method:%s)\n%s", event.Pid, event.Uid, packageName, interfaceToken, methodName, spew.Sdump(parcelData))
			} else {
				logger.Printf("(pid:%d, uid:%d, package:%s) -> (interface:%s, method:%s)\n", event.Pid, event.Uid, packageName, interfaceToken, methodName)
			}
		}
	}
}