main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/cheggaaa/pb"
	"github.com/cilium/ebpf"
	"github.com/cilium/ebpf/btf"
	"github.com/cilium/ebpf/link"
	"github.com/cilium/ebpf/rlimit"
	"golang.org/x/sys/unix"

	"github.com/danteslimbo/dust/internal/dust"
)

func main() {
	flags := dust.Flags{}
	flags.SetFlags()
	flags.Parse()

	if flags.ShowHelp {
		flags.PrintHelp()
		os.Exit(0)
	}

	if flags.ShowVersion {
		fmt.Printf("dust %s\n", dust.Version)
		os.Exit(0)
	}

	cfg := dust.GetConfig(&flags)

	if err := unix.Setrlimit(unix.RLIMIT_NOFILE, &unix.Rlimit{
		Cur: 8192,
		Max: 8192,
	}); err != nil {
		log.Fatalf("failed to set temporary rlimit: %s", err)
	}
	if err := rlimit.RemoveMemlock(); err != nil {
		log.Fatalf("Failed to set temporary rlimit: %s", err)
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	btfSpec, err := btf.LoadKernelSpec()
	if err != nil {
		log.Fatalf("Failed to load BTF spec: %s", err)
	}
	funcs, err := dust.GetFuncs("", btfSpec, []string{}, false)
	if err != nil {
		log.Fatalf("Failed to get skb-accepting functions: %s", err)
	}

	addr2name, _, err := dust.ParseKallsyms(funcs, true)
	if err != nil {
		log.Fatalf("Failed to parse Kallsyms: %s", err)
	} else {
		log.Printf("Parsed Kallsyms: %d syms\n", len(addr2name.Addr2NameMap))
	}

	var bpfSpec *ebpf.CollectionSpec
	bpfSpec, err = LoadKProbeRequest() // generated by `go generate`
	if err != nil {
		log.Fatalf("Failed to load bpf spec: %s", err)
	}
	if err := bpfSpec.RewriteConstants(map[string]interface{}{
		"CFG": cfg,
	}); err != nil {
		log.Fatalf("Failed to rewrite config: %v", err)
	}

	var opts ebpf.CollectionOptions
	opts.Programs.KernelTypes = btfSpec
	opts.Programs.LogLevel = ebpf.LogLevelInstruction
	opts.Programs.LogSize = ebpf.DefaultVerifierLogSize * 100

	coll, err := ebpf.NewCollectionWithOptions(bpfSpec, opts)
	if err != nil {
		var (
			ve          *ebpf.VerifierError
			verifierLog string
		)
		if errors.As(err, &ve) {
			verifierLog = fmt.Sprintf("Verifier error: %+v\n", ve)
		}

		log.Fatalf("Failed to load objects: %s\n%+v", verifierLog, err)
	}
	defer coll.Close()

	var kprobes []link.Link
	defer func() {
		dust.DetachKprobes(kprobes, 0)
	}()

	msg := "kprobe"
	log.Printf("Attaching kprobes (via %s)...\n", msg)

	ignored := 0
	bar := pb.StartNew(len(funcs))

	dustKprobes := make([]dust.Kprobe, 0, len(funcs))
	funcsByPos := dust.GetFuncsByPos(funcs)
	for pos, fns := range funcsByPos {
		fn, ok := coll.Programs[fmt.Sprintf("kprobe_dust_%d", pos)]
		if ok {
			dustKprobes = append(dustKprobes, dust.Kprobe{HookFuncs: fns, Prog: fn})
		} else {
			ignored += len(fns)
			bar.Add(len(fns))
		}
	}

	retProg, ok := coll.Programs["alloc"]
	if !ok {
		log.Fatalf("Failed to find kretprobe_alloc_request")
	}
	retLink, err := link.Kretprobe("blk_mq_alloc_request", retProg, nil)
	if err != nil {
		log.Fatalf("Failed to attach kretprobe_alloc_request: %s", err)
	}
	defer retLink.Close()

	l, i := dust.AttachKprobes(ctx, bar, dustKprobes, 10)
	kprobes = append(kprobes, l...)
	ignored += i

	bar.Finish()
	select {
	case <-ctx.Done():
		return
	default:
	}
	log.Printf("Attached (ignored %d)\n", ignored)

	log.Println("Listening for events..")

	output, err := dust.NewOutput(addr2name, false, flags.Ofile)
	if err != nil {
		log.Fatalf("Failed to create output: %s", err)
	}
	defer output.Close()

	output.PrintHeader()

	defer func() {
		select {
		case <-ctx.Done():
			log.Println("Received signal, exiting program..")
		default:
		}
	}()

	var event dust.Event
	events := coll.Maps["events"]
	c := time.Tick(time.Duration(flags.Interval) * time.Second)
	_, _ = fmt.Fprintf(os.Stderr, "Fetching requests for %d seconds...\n", flags.Interval)
	for {
		for {
			if err := events.LookupAndDelete(nil, &event); err == nil {
				break
			}

			select {
			case <-ctx.Done():
				return
			case <-time.After(time.Microsecond):
				continue
			}
		}
		output.Print(&event)
		select {
		case <-ctx.Done():
		case <-c:
			return
		default:
		}
	}
}