fix main.go

ebpf-poc/main.go

@@ -2,88 +2,105 @@
 
 package main
 
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"net"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/cilium/ebpf"
+	"github.com/cilium/ebpf/link"
+	"github.com/cilium/ebpf/ringbuf"
+	"github.com/cilium/ebpf/rlimit"
+)
+
 func main() {
-	// err := rlimit.RemoveMemlock()
+	err := rlimit.RemoveMemlock()
-	// if err != nil {
+	if err != nil {
-	// 	fmt.Printf("Failed to remove mem lock: %v\n", err)
+		fmt.Printf("Failed to remove mem lock: %v\n", err)
-	// 	return
+		return
-	// }
+	}
 
-	// var objs bpfObjects
+	var objs bpfObjects
-	// err = loadBpfObjects(&objs, nil)
+	err = loadBpfObjects(&objs, nil)
-	// if err != nil {
+	if err != nil {
-	// 	fmt.Printf("Failed to load bpf objects: %v\n", err)
+		fmt.Printf("Failed to load bpf objects: %v\n", err)
-	// 	return
+		return
-	// }
+	}
 
-	// defer objs.Close()
+	defer objs.Close()
 
-	// rootCgroup := "/sys/fs/cgroup"
+	rootCgroup := "/sys/fs/cgroup"
 
-	// l, err := link.AttachCgroup(link.CgroupOptions{
+	l, err := link.AttachCgroup(link.CgroupOptions{
-	// 	Path:    rootCgroup,
+		Path:    rootCgroup,
-	// 	Attach:  ebpf.AttachCGroupInet4Connect,
+		Attach:  ebpf.AttachCGroupInet4Connect,
-	// 	Program: objs.Connect4,
+		Program: objs.Connect4,
-	// })
+	})
-	// if err != nil {
+	if err != nil {
-	// 	fmt.Printf("Failed to attach cgroup: %v\n", err)
+		fmt.Printf("Failed to attach cgroup: %v\n", err)
-	// 	return
+		return
-	// }
+	}
-	// defer l.Close()
+	defer l.Close()
 
-	// rd, err := ringbuf.NewReader(objs.Events)
+	rd, err := ringbuf.NewReader(objs.Events)
-	// if err != nil {
+	if err != nil {
-	// 	fmt.Printf("Failed to create ringbuf reader: %v\n", err)
+		fmt.Printf("Failed to create ringbuf reader: %v\n", err)
-	// 	return
+		return
-	// }
+	}
-	// defer rd.Close()
+	defer rd.Close()
 
-	// ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
-	// defer stop() // Cleans up resources allocated by the signal package
+	defer stop() // Cleans up resources allocated by the signal package
 
-	// fmt.Println("Application started. Press Ctrl+C to exit.")
+	fmt.Println("Application started. Press Ctrl+C to exit.")
 
-	// go func() {
+	go func() {
-	// 	<-ctx.Done()
+		<-ctx.Done()
-	// 	fmt.Println("\nSignal received, detaching and cleaning up...")
+		fmt.Println("\nSignal received, detaching and cleaning up...")
-	// 	rd.Close()
+		rd.Close()
-	// }()
+	}()
 
-	// var e bpfEvent
+	var e bpfEvent
-	// for {
+	for {
-	// 	rec, err := rd.Read()
+		rec, err := rd.Read()
-	// 	if err != nil {
+		if err != nil {
-	// 		if errors.Is(err, ringbuf.ErrClosed) {
+			if errors.Is(err, ringbuf.ErrClosed) {
-	// 			return
+				return
-	// 		}
+			}
-	// 		fmt.Printf("failed to read raw event: %v\n", err)
+			fmt.Printf("failed to read raw event: %v\n", err)
-	// 		continue
+			continue
-	// 	}
+		}
 
-	// 	err = binary.Read(bytes.NewReader(rec.RawSample), binary.LittleEndian, &e)
+		err = binary.Read(bytes.NewReader(rec.RawSample), binary.LittleEndian, &e)
-	// 	if err != nil {
+		if err != nil {
-	// 		fmt.Printf("failed to read event: %v\n", err)
+			fmt.Printf("failed to read event: %v\n", err)
-	// 		continue
+			continue
-	// 	}
+		}
 
-	// 	ip := make(net.IP, 4)
+		ip := make(net.IP, 4)
-	// 	binary.LittleEndian.PutUint32(ip, e.Daddr)
+		binary.LittleEndian.PutUint32(ip, e.Daddr)
 
-	// 	name := e.Comm[:]
+		name := e.Comm[:]
-	// 	if i := bytes.IndexByte(name, 0); i != -1 {
+		if i := bytes.IndexByte(name, 0); i != -1 {
-	// 		name = name[:i] // keep bytes before the first NUL
+			name = name[:i] // keep bytes before the first NUL
-	// 	}
+		}
-	// 	comm := string(name)
+		comm := string(name)
 
-	// 	proto := ""
+		proto := ""
-	// 	switch e.Proto {
+		switch e.Proto {
-	// 	case syscall.IPPROTO_TCP:
+		case syscall.IPPROTO_TCP:
-	// 		proto = "TCP"
+			proto = "TCP"
-	// 	case syscall.IPPROTO_UDP:
+		case syscall.IPPROTO_UDP:
-	// 		proto = "UDP"
+			proto = "UDP"
-	// 	}
+		}
 
-	// 	fmt.Println()
+		fmt.Println()
-	// 	fmt.Printf("PID: %v\nUID: %v\ndAddr: %v\ndPort: %v\nProto: %v\nCommand: %v\n", e.Pid, e.Uid, ip, e.Dport, proto, comm)
+		fmt.Printf("PID: %v\nUID: %v\ndAddr: %v\ndPort: %v\nProto: %v\nCommand: %v\n", e.Pid, e.Uid, ip, e.Dport, proto, comm)
-	// }
+	}
 }