Add additional options for build and run (#89)

Author: SirNexus

.gitignore

@@ -21,6 +21,7 @@ _output
 # Staging ground for bpf programs
 bpf/
 
-# Ignore vscode
+# Ignore IDE folders
 .vscode/
-.vagrant/
+.vagrant/
+.idea

builder/build.sh

@@ -1,7 +1,9 @@
 #! /bin/bash
-set -eu
+set -eux
 
-clang-13 -g -O2 -target bpf -D__TARGET_ARCH_x86 -Wall -c $1 -o $2
+CFLAGS=${CFLAGS:-}
+
+clang-13 -g -O2 -target bpf -D__TARGET_ARCH_x86 ${CFLAGS} -Wall -c $1 -o $2
 
 # strip debug sections (see: https://github.com/libbpf/libbpf-bootstrap/blob/94000ca67c5e7be4741c09c435c9ae1777822378/examples/c/Makefile#L65)
 llvm-strip-13 -g $2

pkg/cli/internal/commands/build/build.go

@@ -12,29 +12,51 @@ import (
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pterm/pterm"
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"oras.land/oras-go/pkg/content"
+
 	"github.com/solo-io/bumblebee/builder"
 	"github.com/solo-io/bumblebee/pkg/cli/internal/options"
 	"github.com/solo-io/bumblebee/pkg/internal/version"
 	"github.com/solo-io/bumblebee/pkg/spec"
-	"github.com/spf13/cobra"
-	"github.com/spf13/pflag"
-	"oras.land/oras-go/pkg/content"
 )
 
 type buildOptions struct {
-	BuildImage string
-	Builder    string
-	OutputFile string
+	BuildImage        string
+	Builder           string
+	OutputFile        string
 	Local      bool
+	BinaryOnly bool
+	CFlags     []string
+	BuildScript       string
+	BuildScriptOutput bool
 
 	general *options.GeneralOptions
 }
 
+func (opts *buildOptions) validate() error {
+	if !opts.Local {
+		if opts.BuildScript != "" {
+			fmt.Println("ignoring specified build script for docker build")
+		}
+		if opts.BuildScriptOutput {
+			return fmt.Errorf("cannot write build script output for docker build")
+		}
+	}
+
+	return nil
+}
+
 func addToFlags(flags *pflag.FlagSet, opts *buildOptions) {
 	flags.StringVarP(&opts.BuildImage, "build-image", "i", fmt.Sprintf("ghcr.io/solo-io/bumblebee/builder:%s", version.Version), "Build image to use when compiling BPF program")
 	flags.StringVarP(&opts.Builder, "builder", "b", "docker", "Executable to use for docker build command, default: `docker`")
 	flags.StringVarP(&opts.OutputFile, "output-file", "o", "", "Output file for BPF ELF. If left blank will default to <inputfile.o>")
-	flags.BoolVarP(&opts.Local, "local", "l", false, "Build the output binary and OCI image using local tools")
+	flags.BoolVarP(&opts.Local, "local", "l", false, "Build the output binary using local tools")
+	flags.StringVar(&opts.BuildScript, "build-script", "", "Optional path to a build script for building BPF program locally")
+	flags.BoolVar(&opts.BuildScriptOutput, "build-script-out", false, "Print local script bee will use to build the BPF program")
+	flags.BoolVar(&opts.BinaryOnly, "binary-only", false, "Only create output binary and do not package it into an OCI image")
+	flags.StringArrayVar(&opts.CFlags, "cflags", nil, "cflags to be used when compiling the BPF program, passed as environment variable 'CFLAGS'")
 }
 
 func Command(opts *options.GeneralOptions) *cobra.Command {
@@ -52,8 +74,22 @@ The bee build command has 2 main parts
 By default building is done in a docker container, however, this can be switched to local by adding the local flag:
 $ build INPUT_FILE REGISTRY_REF --local
 
+If you would prefer to build only the object file, you can include the '--binary-only' flag,
+in which case you do not need to specify a registry:
+$ build INPUT_FILE --binary-only
+
+Examples:
+You can specify multiple cflags with either a space separated string, or with multiple instances:
+$ build INPUT_FILE REGISTRY_REF --cflags="-DDEBUG -DDEBUG2"
+$ build INPUT_FILE REGISTRY_REF --cflags=-DDEBUG --cflags=-DDEBUG2
+
+You can specify your own build script for building the BPF program locally.
+the input file will be passed as argument '$1' and the output filename as '$2'.
+Use the '--build-script-out' flag to see the default build script bee uses:
+$ build INPUT_FILE REGISTRY_REF --local --build-script-out
+$ build INPUT_FILE REGISTRY_REF --local --build-script=build.sh
 `,
-		Args: cobra.ExactArgs(2),
+		Args: cobra.RangeArgs(1, 2),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			return build(cmd.Context(), args, buildOpts)
 		},
@@ -69,6 +105,9 @@ $ build INPUT_FILE REGISTRY_REF --local
 }
 
 func build(ctx context.Context, args []string, opts *buildOptions) error {
+	if err := opts.validate(); err != nil {
+		return err
+	}
 
 	inputFile := args[0]
 	outputFile := opts.OutputFile
@@ -99,8 +138,17 @@ func build(ctx context.Context, args []string, opts *buildOptions) error {
 	// Create and start a fork of the default spinner.
 	var buildSpinner *pterm.SpinnerPrinter
 	if opts.Local {
+		buildScript, err := getBuildScript(opts.BuildScript)
+		if err != nil {
+			return fmt.Errorf("could not load build script: %v", err)
+		}
+		if opts.BuildScriptOutput {
+			fmt.Printf("%s\n", buildScript)
+			return nil
+		}
+
 		buildSpinner, _ = pterm.DefaultSpinner.Start("Compiling BPF program locally")
-		if err := buildLocal(ctx, inputFile, outputFile); err != nil {
+		if err := buildLocal(ctx, opts, buildScript, inputFile, outputFile); err != nil {
 			buildSpinner.UpdateText("Failed to compile BPF program locally")
 			buildSpinner.Fail()
 			return err
@@ -116,6 +164,14 @@ func build(ctx context.Context, args []string, opts *buildOptions) error {
 	buildSpinner.UpdateText(fmt.Sprintf("Successfully compiled \"%s\" and wrote it to \"%s\"", inputFile, outputFile))
 	buildSpinner.Success() // Resolve spinner with success message.
 
+	if opts.BinaryOnly {
+		return nil
+	}
+
+	if len(args) == 1 {
+		return fmt.Errorf("must specify a registry to package the output or run with '--binary-only'")
+	}
+
 	// TODO: Figure out this hack, file.Seek() didn't seem to work
 	outputFd.Close()
 	reopened, err := os.Open(outputFile)
@@ -175,6 +231,18 @@ func getPlatformInfo(ctx context.Context) *ocispec.Platform {
 	}
 }
 
+func getBuildScript(path string) ([]byte, error) {
+	if path != "" {
+		buildScript, err := os.ReadFile(path)
+		if err != nil {
+			return nil, fmt.Errorf("could not read build script: %v", err)
+		}
+		return buildScript, nil
+	}
+
+	return builder.GetBuildScript(), nil
+}
+
 func buildDocker(
 	ctx context.Context,
 	opts *buildOptions,
@@ -190,10 +258,13 @@ func buildDocker(
 		"run",
 		"-v",
 		fmt.Sprintf("%s:/usr/src/bpf", wd),
-		opts.BuildImage,
-		inputFile,
-		outputFile,
 	}
+
+	if len(opts.CFlags) > 0 {
+		dockerArgs = append(dockerArgs, "--env", fmt.Sprintf("CFLAGS=%s", strings.Join(opts.CFlags, " ")))
+	}
+	dockerArgs = append(dockerArgs, opts.BuildImage, inputFile, outputFile)
+
 	dockerCmd := exec.CommandContext(ctx, opts.Builder, dockerArgs...)
 	byt, err := dockerCmd.CombinedOutput()
 	if err != nil {
@@ -203,12 +274,19 @@ func buildDocker(
 	return nil
 }
 
-func buildLocal(ctx context.Context, inputFile, outputFile string) error {
-	buildScript := builder.GetBuildScript()
-
+func buildLocal(
+	ctx context.Context,
+	opts *buildOptions,
+	buildScript []byte,
+	inputFile,
+	outputFile string,
+) error {
 	// Pass the script into sh via stdin, then arguments
 	// TODO: need to handle CWD gracefully
 	shCmd := exec.CommandContext(ctx, "sh", "-s", "--", inputFile, outputFile)
+	shCmd.Env = []string{
+		fmt.Sprintf("CFLAGS=%s", strings.Join(opts.CFlags, " ")),
+	}
 	stdin, err := shCmd.StdinPipe()
 	if err != nil {
 		return err
@@ -220,8 +298,8 @@ func buildLocal(ctx context.Context, inputFile, outputFile string) error {
 	}()
 
 	out, err := shCmd.CombinedOutput()
+	pterm.Info.Printf("%s\n", out)
 	if err != nil {
-		fmt.Printf("%s\n", out)
 		return err
 	}
 	return nil

pkg/cli/internal/commands/run/run.go

@@ -27,9 +27,11 @@ import (
 type runOptions struct {
 	general *options.GeneralOptions
 
-	debug  bool
-	filter []string
-	notty  bool
+	debug    bool
+	filter   []string
+	notty    bool
+	pinMaps  string
+	pinProgs string
 }
 
 const filterDescription string = "Filter to apply to output from maps. Format is \"map_name,key_name,regex\" " +
@@ -41,6 +43,8 @@ func addToFlags(flags *pflag.FlagSet, opts *runOptions) {
 	flags.BoolVarP(&opts.debug, "debug", "d", false, "Create a log file 'debug.log' that provides debug logs of loader and TUI execution")
 	flags.StringSliceVarP(&opts.filter, "filter", "f", []string{}, filterDescription)
 	flags.BoolVar(&opts.notty, "no-tty", false, "Set to true for running without a tty allocated, so no interaction will be expected or rich output will done")
+	flags.StringVar(&opts.pinMaps, "pin-maps", "", "Directory to pin maps to, left unpinned if empty")
+	flags.StringVar(&opts.pinProgs, "pin-progs", "", "Directory to pin progs to, left unpinned if empty")
 }
 
 func Command(opts *options.GeneralOptions) *cobra.Command {
@@ -119,6 +123,8 @@ func run(cmd *cobra.Command, args []string, opts *runOptions) error {
 	loaderOpts := loader.LoadOptions{
 		ParsedELF: parsedELF,
 		Watcher:   tuiApp,
+		PinMaps:   opts.pinMaps,
+		PinProgs:  opts.pinProgs,
 	}
 
 	// bail out before starting TUI if context canceled

pkg/loader/loader.go

@@ -6,17 +6,20 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"os"
+	"path/filepath"
 	"strings"
 	"time"
 
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/btf"
 	"github.com/cilium/ebpf/link"
 	"github.com/cilium/ebpf/ringbuf"
+	"golang.org/x/sync/errgroup"
+
 	"github.com/solo-io/bumblebee/pkg/decoder"
 	"github.com/solo-io/bumblebee/pkg/stats"
 	"github.com/solo-io/go-utils/contextutils"
-	"golang.org/x/sync/errgroup"
 )
 
 type ParsedELF struct {
@@ -27,11 +30,14 @@ type ParsedELF struct {
 type LoadOptions struct {
 	ParsedELF *ParsedELF
 	Watcher   MapWatcher
+	PinMaps   string
+	PinProgs  string
 }
 
 type Loader interface {
 	Parse(ctx context.Context, reader io.ReaderAt) (*ParsedELF, error)
 	Load(ctx context.Context, opts *LoadOptions) error
+	WatchMaps(ctx context.Context, watchedMaps map[string]WatchedMap, coll map[string]*ebpf.Map, watcher MapWatcher) error
 }
 
 type WatchedMap struct {
@@ -88,6 +94,12 @@ func (l *loader) Parse(ctx context.Context, progReader io.ReaderAt) (*ParsedELF,
 		return nil, err
 	}
 
+	for _, prog := range spec.Programs {
+		if prog.Type == ebpf.UnspecifiedProgram {
+			contextutils.LoggerFrom(ctx).Debug("Program %s does not specify a type", prog.Name)
+		}
+	}
+
 	watchedMaps := make(map[string]WatchedMap)
 	for name, mapSpec := range spec.Maps {
 		if !isTrackedMap(mapSpec) {
@@ -150,9 +162,26 @@ func (l *loader) Load(ctx context.Context, opts *LoadOptions) error {
 		return ctx.Err()
 	}
 
+	if opts.PinMaps != "" {
+		// Specify that we'd like to pin the referenced maps, or open them if already existing.
+		for _, m := range opts.ParsedELF.Spec.Maps {
+			// Do not pin/load read-only data
+			if strings.HasSuffix(m.Name, ".rodata") {
+				continue
+			}
+
+			// PinByName specifies that we should pin the map by name, or load it if it already exists.
+			m.Pinning = ebpf.PinByName
+		}
+	}
+
 	spec := opts.ParsedELF.Spec
 	// Load our eBPF spec into the kernel
-	coll, err := ebpf.NewCollection(spec)
+	coll, err := ebpf.NewCollectionWithOptions(opts.ParsedELF.Spec, ebpf.CollectionOptions{
+		Maps: ebpf.MapOptions{
+			PinPath: opts.PinMaps,
+		},
+	})
 	if err != nil {
 		return err
 	}
@@ -198,13 +227,29 @@ func (l *loader) Load(ctx context.Context, opts *LoadOptions) error {
 			default:
 				return errors.New("only kprobe programs supported")
 			}
+			if opts.PinProgs != "" {
+				if err := createDir(ctx, opts.PinProgs, 0700); err != nil {
+					return err
+				}
+
+				pinFile := filepath.Join(opts.PinProgs, prog.Name)
+				if err := coll.Programs[name].Pin(pinFile); err != nil {
+					return fmt.Errorf("could not pin program '%s': %v", prog.Name, err)
+				}
+				fmt.Printf("Successfully pinned program '%v'\n", pinFile)
+			}
 		}
 	}
 
-	return l.watchMaps(ctx, opts.ParsedELF.WatchedMaps, coll, opts.Watcher)
+	return l.WatchMaps(ctx, opts.ParsedELF.WatchedMaps, coll.Maps, opts.Watcher)
 }
 
-func (l *loader) watchMaps(ctx context.Context, watchedMaps map[string]WatchedMap, coll *ebpf.Collection, watcher MapWatcher) error {
+func (l *loader) WatchMaps(
+	ctx context.Context,
+	watchedMaps map[string]WatchedMap,
+	maps map[string]*ebpf.Map,
+	watcher MapWatcher,
+) error {
 	contextutils.LoggerFrom(ctx).Info("enter watchMaps()")
 	eg, ctx := errgroup.WithContext(ctx)
 	for name, bpfMap := range watchedMaps {
@@ -221,7 +266,7 @@ func (l *loader) watchMaps(ctx context.Context, watchedMaps map[string]WatchedMa
 			}
 			eg.Go(func() error {
 				watcher.NewRingBuf(name, bpfMap.Labels)
-				return l.startRingBuf(ctx, bpfMap.valueStruct, coll.Maps[name], increment, name, watcher)
+				return l.startRingBuf(ctx, bpfMap.valueStruct, maps[name], increment, name, watcher)
 			})
 		case ebpf.Array:
 			fallthrough
@@ -238,7 +283,7 @@ func (l *loader) watchMaps(ctx context.Context, watchedMaps map[string]WatchedMa
 			eg.Go(func() error {
 				// TODO: output type of instrument in UI?
 				watcher.NewHashMap(name, labelKeys)
-				return l.startHashMap(ctx, bpfMap.mapSpec, coll.Maps[name], instrument, name, watcher)
+				return l.startHashMap(ctx, bpfMap.mapSpec, maps[name], instrument, name, watcher)
 			})
 		default:
 			// TODO: Support more map types
@@ -407,3 +452,17 @@ func (n *noop) Set(
 	labels map[string]string,
 ) {
 }
+
+func createDir(ctx context.Context, path string, perm os.FileMode) error {
+	file, err := os.Stat(path)
+	if os.IsNotExist(err) {
+		contextutils.LoggerFrom(ctx).Info("path does not exist, creating pin directory: %s", path)
+		return os.Mkdir(path, perm)
+	} else if err != nil {
+		return fmt.Errorf("could not create pin directory '%v': %w", path, err)
+	} else if !file.IsDir() {
+		return fmt.Errorf("pin location '%v' exists but is not a directory", path)
+	}
+
+	return nil
+}