ebpf/ebpftracer/tracer/socket.go

package tracer

import (
	"fmt"
	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/btf"
	"k8s.io/klog/v2"
	"net"
	"os"
	"runtime"
	"syscall"
	"time"
)

func init() {
	enable_ebpf_protocol(PROTO_HTTP1)
	enable_ebpf_protocol(PROTO_HTTP2)
	enable_ebpf_protocol(PROTO_TLS_HTTP1)
	enable_ebpf_protocol(PROTO_TLS_HTTP2)
	enable_ebpf_protocol(PROTO_DUBBO)
	enable_ebpf_protocol(PROTO_SOFARPC)
	enable_ebpf_protocol(PROTO_MYSQL)
	enable_ebpf_protocol(PROTO_POSTGRESQL)
	enable_ebpf_protocol(PROTO_REDIS)
	enable_ebpf_protocol(PROTO_KAFKA)
	enable_ebpf_protocol(PROTO_MQTT)
	enable_ebpf_protocol(PROTO_DNS)
}

func MapInit(collectionSpec *ebpf.CollectionSpec, opts *ebpf.CollectionOptions) {
	set_offset_map(collectionSpec, opts)
	set_conf_map_default(collectionSpec, opts)

	//offsetData := make([]any, runtime.NumCPU())
	//for i := range offsetData {
	//	offsetData[i] = testStruct{
	//		test_id: 99999,
	//	}
	//}
	//if bpf_table_set_value(collectionSpec, opts, "test_heap", offsetData) != ETR_OK {}
	//insert_output_prog_to_map(collectionSpec, opts)
}

func MapInsert(collection *ebpf.Collection) {
	insert_output_prog_to_map(collection)
	insert_adapt_kern_uid_to_map(collection)

	// Update protocol filter array
	update_protocol_filter_array(collection)

}

func insert_output_prog_to_map(collection *ebpf.Collection) {
	// jmp for tracepoints
	__insert_output_prog_to_map(collection, MAP_PROGS_JMP_TP_NAME, PROG_DATA_SUBMIT_NAME_FOR_TP, PROG_DATA_SUBMIT_TP_IDX)
	__insert_output_prog_to_map(collection, MAP_PROGS_JMP_TP_NAME, PROG_OUTPUT_DATA_NAME_FOR_TP, PROG_OUTPUT_DATA_TP_IDX)
	__insert_output_prog_to_map(collection, MAP_PROGS_JMP_TP_NAME, PROG_IO_EVENT_NAME_FOR_TP, PROG_IO_EVENT_TP_IDX)

	// jmp for kprobe/uprobe
	__insert_output_prog_to_map(collection, MAP_PROGS_JMP_KP_NAME, PROG_DATA_SUBMIT_NAME_FOR_KP, PROG_DATA_SUBMIT_KP_IDX)
	__insert_output_prog_to_map(collection, MAP_PROGS_JMP_KP_NAME, PROG_OUTPUT_DATA_NAME_FOR_KP, PROG_OUTPUT_DATA_KP_IDX)
}

func __insert_output_prog_to_map(collection *ebpf.Collection, mapName string, progName string, key uint32) {
	// find in programs
	prog, ok := collection.Programs[progName]
	fmt.Println(prog, ok)
	if ok {
		progFd := prog.FD()
		//fmt.Println("progFd", progFd)
		if bpf_table_set_value(collection, mapName, key, uint32(progFd)) != ETR_OK {
			//fmt.Println("no")
		} else {
			//fmt.Println("ok")
		}
	}
}

func update_protocol_filter_array(collection *ebpf.Collection) {
	for i := 0; i < PROTO_NUM; i++ {
		if bpf_table_set_value(collection, MAP_PROTO_FILTER_NAME, uint32(i), EbpfConfigProtocolFilter[i]) != ETR_OK {
			//fmt.Println("no")
		} else {
			//fmt.Println("ok")
		}
	}
}

//func update_allow_port_bitmap(collection *ebpf.Collection) {
//	for i := 0; i < PROTO_NUM; i++ {
//		if bpf_table_set_value(collection, MAP_ALLOW_PORT_BITMAP_NAME, 0, &allow_port_bitmap) != ETR_OK {
//			fmt.Println("no")
//		} else {
//			fmt.Println("ok")
//		}
//	}
//}

func insert_adapt_kern_uid_to_map(collection *ebpf.Collection) {
	pid := os.Getpid()
	tid := syscall.Gettid()
	adaptKernUID := uint64(pid)<<32 | uint64(tid)
	if bpf_table_set_value(collection, MAP_ADAPT_KERN_UID_NAME, 0, uint32(adaptKernUID)) != ETR_OK {
		fmt.Println("no")
	} else {
		fmt.Println("ok")
	}
}

func enable_ebpf_protocol(protocol int) {
	if protocol < PROTO_NUM {
		EbpfConfigProtocolFilter[protocol] = 1
	}
}

func Offset() {
	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", offsetInferServerAddr, offsetInferServerPort))
	if err != nil {
		fmt.Errorf("failed to create server listener: %v", err)
		return
	}
	if err := kernelOffsetInferServer(listener); err != nil {
		fmt.Printf("Error in kernel_offset_infer_server: %v\n", err)
	}
	if err := kernelOffsetInferClient(); err != nil {
		fmt.Printf("Error in kernel_offset_infer_client: %v\n", err)
	}
	defer listener.Close()
}

func set_offset_map(collectionSpec *ebpf.CollectionSpec, opts *ebpf.CollectionOptions) {
	// 解析BTF数据
	if update_offset_map_from_btf_vmlinux(collectionSpec, opts) != ETR_OK {
		klog.Infof("[eBPF Kernel Adapt] Set offsets map from btf_vmlinux, not support.\n")
		if update_offset_map_default(collectionSpec, opts) != ETR_OK {
			klog.Infof("Fatal error, failed to update default offset\n")
		}
	} else {
		klog.Infof("[eBPF Kernel Adapt] Set offsets map from btf_vmlinux, success.\n")
	}
}

func set_conf_map_default(collectionSpec *ebpf.CollectionSpec, opts *ebpf.CollectionOptions) {
	uidBase := uint64(time.Now().UnixNano()/int64(time.Millisecond)) & 0xffffffffffffff
	numCPU := runtime.NumCPU()
	tConf := make([]any, numCPU)
	for i := range tConf {
		socketID := uint64(i)<<56 | uint64(uidBase)
		tracerConf := traceConf{
			SocketID:               socketID,
			CoroutineTraceID:       socketID,
			ThreadTraceID:          socketID,
			DataLimitMax:           4096,
			GoTracingTimeout:       120,
			IOEventCollectMode:     1,
			IOEventMinimalDuartion: 1000000,
		}
		tConf[i] = tracerConf
	}
	if bpf_table_pre_set_value(collectionSpec, opts, MAP_TRACE_CONF_NAME, tConf) != ETR_OK {
		klog.Infof("[eBPF Kernel Adapt] Set config map from btf_vmlinux, not support.\n")
	} else {
		klog.Infof("[eBPF Kernel Adapt] Set config map from btf_vmlinux, success.\n")
	}
}

// 解析BTF数据
func update_offset_map_from_btf_vmlinux(collectionSpec *ebpf.CollectionSpec, opts *ebpf.CollectionOptions) int {
	btfSpec, err := btf.LoadKernelSpec()
	if err != nil || btfSpec == nil {
		return ETR_NOTSUPP
	}

	copied_seq_offs := kernel_struct_field_offset(btfSpec, "tcp_sock", "copied_seq")
	write_seq_offs := kernel_struct_field_offset(btfSpec, "tcp_sock", "write_seq")
	files_offs := kernel_struct_field_offset(btfSpec, "task_struct", "files")
	sk_flags_offs := kernel_struct_field_offset(btfSpec, "sock", "__sk_flags_offset")

	if sk_flags_offs == ETR_NOTEXIST {
		sk_flags_offs = kernel_struct_field_offset(btfSpec, "sock", "sk_pacing_shift")
		if sk_flags_offs > 0 {
			sk_flags_offs -= 1
		}
	}

	struct_files_struct_fdt_offset := kernel_struct_field_offset(btfSpec, "files_struct", "fdt")
	struct_files_private_data_offset := kernel_struct_field_offset(btfSpec, "file", "private_data")
	struct_file_f_inode_offset := kernel_struct_field_offset(btfSpec, "file", "f_inode")
	struct_inode_i_mode_offset := kernel_struct_field_offset(btfSpec, "inode", "i_mode")
	struct_file_dentry_offset_1 := kernel_struct_field_offset(btfSpec, "file", "f_path")
	struct_file_dentry_offset_2 := kernel_struct_field_offset(btfSpec, "path", "dentry")
	if struct_file_dentry_offset_1 < 0 ||
		struct_file_dentry_offset_2 < 0 {
		return ETR_NOTSUPP
	}

	struct_file_dentry_offset := struct_file_dentry_offset_1 + struct_file_dentry_offset_2
	struct_dentry_name_offset_1 := kernel_struct_field_offset(btfSpec, "dentry", "d_name")
	struct_dentry_name_offset_2 := kernel_struct_field_offset(btfSpec, "qstr", "name")
	if struct_dentry_name_offset_1 < 0 ||
		struct_dentry_name_offset_2 < 0 {
		return ETR_NOTSUPP
	}

	struct_dentry_name_offset := struct_dentry_name_offset_1 + struct_dentry_name_offset_2
	struct_sock_family_offset := kernel_struct_field_offset(btfSpec, "sock_common", "skc_family")
	struct_sock_saddr_offset := kernel_struct_field_offset(btfSpec, "sock_common", "skc_rcv_saddr")
	struct_sock_daddr_offset := kernel_struct_field_offset(btfSpec, "sock_common", "skc_daddr")
	struct_sock_ip6saddr_offset := kernel_struct_field_offset(btfSpec, "sock_common", "skc_v6_rcv_saddr")
	struct_sock_ip6daddr_offset := kernel_struct_field_offset(btfSpec, "sock_common", "skc_v6_daddr")
	struct_sock_dport_offset := kernel_struct_field_offset(btfSpec, "sock_common", "skc_dport")
	struct_sock_sport_offset := kernel_struct_field_offset(btfSpec, "sock_common", "skc_num")
	struct_sock_skc_state_offset := kernel_struct_field_offset(btfSpec, "sock_common", "skc_state")
	struct_sock_common_ipv6only_offset := kernel_struct_field_offset(btfSpec, "sock_common", "skc_flags")

	klog.Infof("Offsets from BTF vmlinux:\n")
	klog.Infof("    copied_seq_offs: 0x%x\n", copied_seq_offs)
	klog.Infof("    write_seq_offs: 0x%x\n", write_seq_offs)
	klog.Infof("    files_offs: 0x%x\n", files_offs)
	klog.Infof("    sk_flags_offs: 0x%x\n", sk_flags_offs)
	klog.Infof("    struct_files_struct_fdt_offset: 0x%x\n", struct_files_struct_fdt_offset)
	klog.Infof("    struct_files_private_data_offset: 0x%x\n", struct_files_private_data_offset)
	klog.Infof("    struct_file_f_inode_offset: 0x%x\n", struct_file_f_inode_offset)
	klog.Infof("    struct_inode_i_mode_offset: 0x%x\n", struct_inode_i_mode_offset)
	klog.Infof("    struct_file_dentry_offset: 0x%x\n", struct_file_dentry_offset)
	klog.Infof("    struct_dentry_name_offset: 0x%x\n", struct_dentry_name_offset)
	klog.Infof("    struct_sock_family_offset: 0x%x\n", struct_sock_family_offset)
	klog.Infof("    struct_sock_saddr_offset: 0x%x\n", struct_sock_saddr_offset)
	klog.Infof("    struct_sock_daddr_offset: 0x%x\n", struct_sock_daddr_offset)
	klog.Infof("    struct_sock_ip6saddr_offset: 0x%x\n", struct_sock_ip6saddr_offset)
	klog.Infof("    struct_sock_ip6daddr_offset: 0x%x\n", struct_sock_ip6daddr_offset)
	klog.Infof("    struct_sock_dport_offset: 0x%x\n", struct_sock_dport_offset)
	klog.Infof("    struct_sock_sport_offset: 0x%x\n", struct_sock_sport_offset)
	klog.Infof("    struct_sock_skc_state_offset: 0x%x\n", struct_sock_skc_state_offset)
	klog.Infof("    struct_sock_common_ipv6only_offset: 0x%x\n", struct_sock_common_ipv6only_offset)

	if copied_seq_offs < 0 || write_seq_offs < 0 || files_offs < 0 ||
		sk_flags_offs < 0 || struct_files_struct_fdt_offset < 0 ||
		struct_files_private_data_offset < 0 ||
		struct_file_f_inode_offset < 0 || struct_inode_i_mode_offset < 0 ||
		struct_inode_i_mode_offset < 0 || struct_file_dentry_offset < 0 ||
		struct_dentry_name_offset < 0 || struct_sock_family_offset < 0 ||
		struct_sock_saddr_offset < 0 || struct_sock_daddr_offset < 0 ||
		struct_sock_ip6saddr_offset < 0 ||
		struct_sock_ip6daddr_offset < 0 || struct_sock_dport_offset < 0 ||
		struct_sock_sport_offset < 0 || struct_sock_skc_state_offset < 0 ||
		struct_sock_common_ipv6only_offset < 0 {
		return ETR_NOTSUPP
	}

	offset := bpfOffsetParam{}
	offset.Ready = 1
	offset.TaskFilesOffset = uint32(files_offs)
	offset.SockFlagsOffset = uint32(sk_flags_offs)
	offset.TcpSockCopiedSeqOffset = uint32(copied_seq_offs)
	offset.TcpSockWriteSeqOffset = uint32(write_seq_offs)
	offset.StructFilesStructFdtOffset = uint32(struct_files_struct_fdt_offset)
	offset.StructFilesPrivateDataOffset = uint32(struct_files_private_data_offset)
	offset.StructFileFInodeOffset = uint32(struct_file_f_inode_offset)
	offset.StructInodeIModeOffset = uint32(struct_inode_i_mode_offset)
	offset.StructFileDentryOffset = uint32(struct_file_dentry_offset)
	offset.StructDentryNameOffset = uint32(struct_dentry_name_offset)
	offset.StructSockFamilyOffset = uint32(struct_sock_family_offset)
	offset.StructSockSaddrOffset = uint32(struct_sock_saddr_offset)
	offset.StructSockDaddrOffset = uint32(struct_sock_daddr_offset)
	offset.StructSockIp6saddrOffset = uint32(struct_sock_ip6saddr_offset)
	offset.StructSockIp6daddrOffset = uint32(struct_sock_ip6daddr_offset)
	offset.StructSockDportOffset = uint32(struct_sock_dport_offset)
	offset.StructSockSportOffset = uint32(struct_sock_sport_offset)
	offset.StructSockSkcStateOffset = uint32(struct_sock_skc_state_offset)
	offset.StructSockCommonIpv6onlyOffset = uint32(struct_sock_common_ipv6only_offset)

	if update_offsets_table(collectionSpec, opts, offset) != ETR_OK {
		return ETR_UPDATE_MAP_FAILD
	}
	return ETR_OK
}

func update_offset_map_default(collectionSpec *ebpf.CollectionSpec, opts *ebpf.CollectionOptions) int {
	offset := bpfOffsetParam{}

	offset.StructFilesStructFdtOffset = 0x20
	offset.StructFilesPrivateDataOffset = 0xc8
	offset.StructFileFInodeOffset = 0x20
	offset.StructInodeIModeOffset = 0x00
	offset.StructFileDentryOffset = 0x18
	offset.StructDentryNameOffset = 0x28
	offset.StructSockFamilyOffset = 0x10
	offset.StructSockSaddrOffset = 0x4
	offset.StructSockDaddrOffset = 0x0
	offset.StructSockIp6saddrOffset = 0x48
	offset.StructSockIp6daddrOffset = 0x38
	offset.StructSockDportOffset = 0xc
	offset.StructSockSportOffset = 0xe
	offset.StructSockSkcStateOffset = 0x12
	offset.StructSockCommonIpv6onlyOffset = 0x13

	if update_offsets_table(collectionSpec, opts, offset) != ETR_OK {
		return ETR_UPDATE_MAP_FAILD
	}
	return ETR_OK
}

func update_offsets_table(collectionSpec *ebpf.CollectionSpec, opts *ebpf.CollectionOptions, offset bpfOffsetParam) int {
	numCPU := runtime.NumCPU()
	offsetData := make([]any, numCPU)
	for i := range offsetData {
		offsetData[i] = offset
	}
	if bpf_table_pre_set_value(collectionSpec, opts, MAP_MEMBERS_OFFSET_NAME, offsetData) != ETR_OK {
		return ETR_UPDATE_MAP_FAILD
	}
	return ETR_OK
}