github.com/dylandreimerink/gobpfld@v0.6.1-0.20220205171531-e79c330ad608/cmd/examples/tailcall/main.go (about)

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"embed"
     6  	"fmt"
     7  	"os"
     8  	"os/signal"
     9  	"runtime"
    10  	"time"
    11  
    12  	"github.com/dylandreimerink/gobpfld"
    13  	"github.com/dylandreimerink/gobpfld/bpftypes"
    14  	"golang.org/x/sys/unix"
    15  )
    16  
    17  //go:embed src/xdp
    18  var f embed.FS
    19  
    20  func main() {
    21  	elfFileBytes, err := f.ReadFile("src/xdp")
    22  	if err != nil {
    23  		fmt.Fprintf(os.Stderr, "error opening ELF file: %s\n", err.Error())
    24  		os.Exit(1)
    25  	}
    26  
    27  	elf, err := gobpfld.LoadProgramFromELF(bytes.NewReader(elfFileBytes), gobpfld.ELFParseSettings{})
    28  	if err != nil {
    29  		fmt.Fprintf(os.Stderr, "error while reading ELF file: %s\n", err.Error())
    30  		os.Exit(1)
    31  	}
    32  
    33  	tailProgs := []string{"ipv4_prog", "ipv6_prog", "tcp_prog", "udp_prog"}
    34  
    35  	entryProgram := elf.Programs["entry_prog"].(*gobpfld.ProgramXDP)
    36  
    37  	log, err := entryProgram.Load(gobpfld.ProgXDPLoadOpts{
    38  		VerifierLogLevel: bpftypes.BPFLogLevelVerbose,
    39  	})
    40  
    41  	fmt.Printf("BPF Verifier log:\n%s\n", log)
    42  
    43  	if err != nil {
    44  		fmt.Fprintf(os.Stderr, "error while loading program 'entry': %s\n", err.Error())
    45  		os.Exit(1)
    46  	}
    47  
    48  	// Before we attach the entry program to the network device we load all tail programs and set their fd's
    49  	// in the 'tails' map.
    50  	tailMap := elf.Maps["tails"].(*gobpfld.ProgArrayMap)
    51  	for i, progName := range tailProgs {
    52  		tailProg := elf.Programs[progName].(*gobpfld.ProgramXDP)
    53  
    54  		log, err := tailProg.Load(gobpfld.ProgXDPLoadOpts{
    55  			VerifierLogLevel: bpftypes.BPFLogLevelVerbose,
    56  		})
    57  
    58  		fmt.Printf("BPF Verifier log:\n%s\n", log)
    59  
    60  		if err != nil {
    61  			fmt.Fprintf(os.Stderr, "error while loading program '%s': %s\n", progName, err.Error())
    62  			os.Exit(1)
    63  		}
    64  
    65  		err = tailMap.Set(int32(i), tailProg)
    66  		if err != nil {
    67  			fmt.Fprintf(os.Stderr, "error while setting prog array fd '%s': %s\n", progName, err.Error())
    68  			os.Exit(1)
    69  		}
    70  	}
    71  
    72  	sigChan := make(chan os.Signal, 1)
    73  	signal.Notify(sigChan, os.Interrupt, unix.SIGTERM, unix.SIGINT)
    74  
    75  	err = entryProgram.Attach(gobpfld.ProgXDPAttachOpts{
    76  		InterfaceName: "lo",
    77  		Replace:       true,
    78  	})
    79  
    80  	if err != nil {
    81  		fmt.Fprintf(os.Stderr, "error while attaching program to loopback device: %s\n", err.Error())
    82  		os.Exit(1)
    83  	}
    84  
    85  	// We need to know how many CPU's this machine has to size our value buffers correctly
    86  	// runtime.NumCPU returns the usable number of CPU's for this process.
    87  	// This number can be different from the number available to the kernel if this process
    88  	// has custom CPU affinity / scheduling. To avoid this the /proc/cpuinfo "file" should be
    89  	// parsed which seems the most reliable method for CPU count detection.
    90  	// But this is not (yet) included in gobpfld.
    91  	numCPUs := runtime.NumCPU()
    92  
    93  	ipStats := elf.Maps["ip_proto_stats"].(*gobpfld.HashMap)
    94  	tcpStats := elf.Maps["tcp_stats"].(*gobpfld.HashMap)
    95  	udpStats := elf.Maps["udp_stats"].(*gobpfld.HashMap)
    96  
    97  	type trafficStats struct {
    98  		packets uint64
    99  		bytes   uint64
   100  	}
   101  
   102  	ticker := time.Tick(1 * time.Second)
   103  	for {
   104  		select {
   105  		case <-ticker:
   106  			printStats := func(k, v interface{}) error {
   107  				var sum trafficStats
   108  				for _, row := range *v.(*[]trafficStats) {
   109  					sum.bytes += row.bytes
   110  					sum.packets += row.packets
   111  				}
   112  
   113  				var key interface{}
   114  				switch t := k.(type) {
   115  				case *uint8:
   116  					key = *t
   117  				case *uint16:
   118  					key = *t
   119  				}
   120  
   121  				fmt.Printf(" %d: pkts: %d, bytes: %d\n", key, sum.packets, sum.bytes)
   122  
   123  				return nil
   124  			}
   125  
   126  			var (
   127  				protoNum uint8
   128  				portNum  uint16
   129  			)
   130  
   131  			stats := make([]trafficStats, numCPUs)
   132  
   133  			fmt.Println("-------------------------")
   134  			fmt.Println("IP Proto stats:")
   135  			iter := ipStats.Iterator()
   136  
   137  			// gobpfld.MapIterForEach(ipStats.Iterator(), &protoNum, &stats, printStats)
   138  			err := iter.Init(&protoNum, &stats)
   139  			if err != nil {
   140  				fmt.Fprintf(os.Stderr, "init: %s\n", err.Error())
   141  				continue
   142  			}
   143  
   144  			var updated bool
   145  			for updated, err = iter.Next(); updated && err == nil; updated, err = iter.Next() {
   146  				err = printStats(&protoNum, &stats)
   147  				if err != nil {
   148  					fmt.Fprintf(os.Stderr, "printStats: %s\n", err.Error())
   149  					continue
   150  				}
   151  			}
   152  
   153  			fmt.Println("TCP Proto stats:")
   154  			gobpfld.MapIterForEach(tcpStats.Iterator(), &portNum, &stats, printStats)
   155  
   156  			fmt.Println("UDP Proto stats:")
   157  			gobpfld.MapIterForEach(udpStats.Iterator(), &portNum, &stats, printStats)
   158  
   159  		case <-sigChan:
   160  			fmt.Println("Detaching XPD program and stopping")
   161  
   162  			err = entryProgram.XDPLinkDetach(gobpfld.BPFProgramXDPLinkDetachSettings{
   163  				All: true,
   164  			})
   165  			if err != nil {
   166  				fmt.Fprintf(os.Stderr, "error while detaching program: %s\n", err.Error())
   167  				os.Exit(1)
   168  			}
   169  
   170  			os.Exit(0)
   171  		}
   172  	}
   173  }