feat: switch binary to attach/detach mode, to use in initcontainer

Author: santhoshmprabhu

bpf-prog/azure-block-iptables/cmd/azure-block-iptables/main.go

@@ -4,118 +4,61 @@
 package main
 
 import (
-	"context"
+	"flag"
 	"fmt"
 	"log"
 	"os"
-	"os/signal"
-	"path/filepath"
-	"syscall"
-	"time"
 
 	"github.com/Azure/azure-container-networking/bpf-prog/azure-block-iptables/pkg/bpfprogram"
 	"github.com/cilium/ebpf/rlimit"
-	"github.com/fsnotify/fsnotify"
 )
 
-const (
-	DefaultConfigFile = "/etc/cni/net.d/iptables-allow-list"
-)
+// ProgramVersion is set during build
+var version = "unknown"
 
-// BlockConfig holds configuration for the application
-type BlockConfig struct {
-	ConfigFile      string
+// Config holds configuration for the application
+type Config struct {
+	Mode            string // "attach" or "detach"
 	AttacherFactory bpfprogram.AttacherFactory
 }
 
-// NewDefaultBlockConfig creates a new BlockConfig with default values
-func NewDefaultBlockConfig() *BlockConfig {
-	return &BlockConfig{
-		ConfigFile:      DefaultConfigFile,
-		AttacherFactory: bpfprogram.NewProgram,
-	}
-}
+// parseArgs parses command line arguments and returns the configuration
+func parseArgs() (*Config, error) {
+	var (
+		mode        = flag.String("mode", "", "Operation mode: 'attach' or 'detach' (required)")
+		showVersion = flag.Bool("version", false, "Show version information")
+		showHelp    = flag.Bool("help", false, "Show help information")
+	)
 
-// isFileEmptyOrMissing checks if the config file exists and has content
-// Returns: 1 if empty, 0 if has content, -1 if missing/error
-func isFileEmptyOrMissing(filename string) int {
-	stat, err := os.Stat(filename)
-	if err != nil {
-		if os.IsNotExist(err) {
-			log.Printf("Config file %s does not exist", filename)
-			return -1 // File missing
-		}
-		log.Printf("Error checking file %s: %v", filename, err)
-		return -1 // Treat errors as missing
-	}
+	flag.Parse()
 
-	if stat.Size() == 0 {
-		log.Printf("Config file %s is empty", filename)
-		return 1 // File empty
+	if *showVersion {
+		fmt.Printf("azure-block-iptables version %s\n", version)
+		os.Exit(0)
 	}
 
-	log.Printf("Config file %s has content (size: %d bytes)", filename, stat.Size())
-	return 0 // File exists and has content
-}
-
-// setupFileWatcher sets up a file watcher for the config file
-func setupFileWatcher(configFile string) (*fsnotify.Watcher, error) {
-	watcher, err := fsnotify.NewWatcher()
-	if err != nil {
-		return nil, fmt.Errorf("failed to create file watcher: %w", err)
+	if *showHelp {
+		flag.PrintDefaults()
+		os.Exit(0)
 	}
 
-	// Watch the directory containing the config file
-	dir := filepath.Dir(configFile)
-	err = watcher.Add(dir)
-	if err != nil {
-		watcher.Close()
-		return nil, fmt.Errorf("failed to add watch for directory %s: %w", dir, err)
+	if *mode == "" {
+		return nil, fmt.Errorf("mode is required. Use -mode=attach or -mode=detach")
 	}
 
-	log.Printf("Watching directory %s for changes to %s", dir, configFile)
-	return watcher, nil
-}
-
-func checkFileStatusAndUpdateBPF(configFile string, bp bpfprogram.Attacher) {
-	// Check current state and take action
-	fileState := isFileEmptyOrMissing(configFile)
-	switch fileState {
-	case 1: // File is empty
-		log.Println("File is empty, attaching BPF program")
-		if err := bp.Attach(); err != nil { // No-op if already attached
-			log.Printf("Failed to attach BPF program: %v", err)
-		}
-	case 0: // File has content
-		log.Println("File has content, detaching BPF program")
-		if err := bp.Detach(); err != nil { // No-op if already detached
-			log.Printf("Failed to detach BPF program: %v", err)
-		}
-	case -1: // File is missing
-		log.Println("Config file was deleted, detaching BPF program")
-		if err := bp.Detach(); err != nil { // No-op if already detached
-			log.Printf("Failed to detach BPF program: %v", err)
-		}
-	}
-}
-
-// handleFileEvent processes file system events
-func handleFileEvent(event fsnotify.Event, configFile string, bp bpfprogram.Attacher) {
-	// Check if the event is for our config file
-	if filepath.Base(event.Name) != filepath.Base(configFile) {
-		return
+	if *mode != "attach" && *mode != "detach" {
+		return nil, fmt.Errorf("invalid mode '%s'. Must be 'attach' or 'detach'", *mode)
 	}
 
-	log.Printf("Config file changed: %s (operation: %s)", event.Name, event.Op)
-
-	// Small delay to handle rapid successive events
-	time.Sleep(100 * time.Millisecond)
-	checkFileStatusAndUpdateBPF(configFile, bp)
+	return &Config{
+		Mode:            *mode,
+		AttacherFactory: bpfprogram.NewProgram,
+	}, nil
 }
 
-// run is the main application logic, separated for easier testing
-func run(config *BlockConfig) error {
-	log.Printf("Using config file: %s", config.ConfigFile)
+// attachMode handles the attach operation
+func attachMode(config *Config) error {
+	log.Println("Starting attach mode...")
 
 	// Remove memory limit for eBPF
 	if err := rlimit.RemoveMemlock(); err != nil {
@@ -124,62 +67,45 @@ func run(config *BlockConfig) error {
 
 	// Initialize BPF program attacher using the factory
 	bp := config.AttacherFactory()
-	defer bp.Close()
 
-	// Check initial state of the config file
-	checkFileStatusAndUpdateBPF(config.ConfigFile, bp)
-
-	// Setup file watcher
-	watcher, err := setupFileWatcher(config.ConfigFile)
-	if err != nil {
-		return fmt.Errorf("failed to setup file watcher: %w", err)
+	// Attach the BPF program
+	if err := bp.Attach(); err != nil {
+		return fmt.Errorf("failed to attach BPF program: %w", err)
 	}
-	defer watcher.Close()
-
-	// Setup signal handling
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-
-	log.Println("Starting file watch loop...")
-
-	// Main event loop
-	for {
-		select {
-		case event, ok := <-watcher.Events:
-			if !ok {
-				log.Println("Watcher events channel closed")
-				return nil
-			}
-			handleFileEvent(event, config.ConfigFile, bp)
-
-		case err, ok := <-watcher.Errors:
-			if !ok {
-				log.Println("Watcher errors channel closed")
-				return nil
-			}
-			log.Printf("Watcher error: %v", err)
-
-		case sig := <-sigChan:
-			log.Printf("Received signal: %v", sig)
-			cancel()
-			return nil
-
-		case <-ctx.Done():
-			log.Println("Context cancelled, exiting")
-			return nil
-		}
+
+	log.Println("BPF program attached successfully")
+	return nil
+}
+
+// detachMode handles the detach operation
+func detachMode(config *Config) error {
+	log.Println("Starting detach mode...")
+
+	// Initialize BPF program attacher using the factory
+	bp := config.AttacherFactory()
+	bp.Detach()
+	log.Println("BPF program detached successfully")
+	return nil
+}
+
+// run is the main application logic
+func run(config *Config) error {
+	switch config.Mode {
+	case "attach":
+		return attachMode(config)
+	case "detach":
+		return detachMode(config)
+	default:
+		return fmt.Errorf("unsupported mode: %s", config.Mode)
 	}
 }
 
 func main() {
-	config := NewDefaultBlockConfig()
-
-	// Parse command line arguments
-	if len(os.Args) > 1 {
-		config.ConfigFile = os.Args[1]
+	config, err := parseArgs()
+	if err != nil {
+		log.Printf("Error parsing arguments: %v", err)
+		flag.PrintDefaults()
+		os.Exit(1)
 	}
 
 	if err := run(config); err != nil {

bpf-prog/azure-block-iptables/pkg/bpfprogram/interface.go

@@ -3,10 +3,10 @@ package bpfprogram
 // Attacher defines the interface for BPF program attachment operations.
 // This interface allows for dependency injection and easier testing with mock implementations.
 type Attacher interface {
-	// Attach attaches the BPF program to LSM hooks
+	// Attach attaches the BPF program to LSM hooks and pins the links and maps
 	Attach() error
 
-	// Detach detaches the BPF program from LSM hooks
+	// Unpins the links and maps (causes detachment)
 	Detach() error
 
 	// IsAttached returns true if the BPF program is currently attached

bpf-prog/azure-block-iptables/pkg/bpfprogram/program.go

@@ -20,6 +20,10 @@ const (
 	BPFMapPinPath = "/sys/fs/bpf/block-iptables"
 	// EventCounterMapName is the name used for pinning the event counter map
 	EventCounterMapName = "iptables_block_event_counter"
+	// IptablesLegacyBlockProgramName is the name used for pinning the legacy iptables block program
+	IptablesLegacyBlockProgramName = "iptables_legacy_block"
+	// IptablesNftablesBlockProgramName is the name used for pinning the nftables block program
+	IptablesNftablesBlockProgramName = "iptables_nftables_block"
 	// NetNSPath is the path to the host network namespace
 	NetNSPath = "/proc/self/ns/net"
 )
@@ -75,6 +79,33 @@ func (p *Program) unpinEventCounterMap() error {
 	return nil
 }
 
+// unpinLinks unpins the links to BPF programs from the filesystem
+func (p *Program) unpinLinks() error {
+	var errs []error
+
+	// Unpin the legacy iptables block program
+	legacyPinPath := filepath.Join(BPFMapPinPath, IptablesLegacyBlockProgramName)
+	if err := os.Remove(legacyPinPath); err != nil && !os.IsNotExist(err) {
+		errs = append(errs, errors.Wrapf(err, "failed to remove pinned legacy program %s", legacyPinPath))
+	} else {
+		log.Printf("Legacy iptables block program unpinned from %s", legacyPinPath)
+	}
+
+	// Unpin the nftables block program
+	nftablesPinPath := filepath.Join(BPFMapPinPath, IptablesNftablesBlockProgramName)
+	if err := os.Remove(nftablesPinPath); err != nil && !os.IsNotExist(err) {
+		errs = append(errs, errors.Wrapf(err, "failed to remove pinned nftables program %s", nftablesPinPath))
+	} else {
+		log.Printf("Nftables block program unpinned from %s", nftablesPinPath)
+	}
+
+	if len(errs) > 0 {
+		return errors.Errorf("failed to unpin programs: %v", errs)
+	}
+
+	return nil
+}
+
 func getHostNetnsInode() (uint64, error) {
 	var stat syscall.Stat_t
 	err := syscall.Stat(NetNSPath, &stat)
@@ -149,6 +180,9 @@ func (p *Program) Attach() error {
 			p.objs = nil
 			return errors.Wrap(err, "failed to attach iptables_legacy_block LSM")
 		}
+
+		pinPath := filepath.Join(BPFMapPinPath, IptablesLegacyBlockProgramName)
+		l.Pin(pinPath)
 		links = append(links, l)
 	}
 
@@ -167,6 +201,8 @@ func (p *Program) Attach() error {
 			p.objs = nil
 			return errors.Wrap(err, "failed to attach block_nf_netlink LSM")
 		}
+		pinPath := filepath.Join(BPFMapPinPath, IptablesNftablesBlockProgramName)
+		l.Pin(pinPath)
 		links = append(links, l)
 	}
 
@@ -177,38 +213,25 @@ func (p *Program) Attach() error {
 	return nil
 }
 
-// Detach detaches the BPF program from LSM hooks.
 func (p *Program) Detach() error {
-	if !p.attached {
-		log.Println("BPF program already detached")
-		return nil
-	}
+	return p.cleanupPinnedResources()
+}
 
-	log.Println("Detaching BPF program...")
+// cleanupPinnedResources removes pinned resources even when the program is not currently attached
+func (p *Program) cleanupPinnedResources() error {
+	log.Println("Cleaning up pinned resources...")
 
-	// Unpin the event counter map from filesystem
-	if err := p.unpinEventCounterMap(); err != nil {
-		log.Printf("Warning: failed to unpin event counter map: %v", err)
+	// Try to unpin links
+	if err := p.unpinLinks(); err != nil {
+		log.Printf("Warning: failed to unpin links: %v", err)
 	}
 
-	// Close all links
-	for _, l := range p.links {
-		if err := l.Close(); err != nil {
-			log.Printf("Warning: failed to close link: %v", err)
-		}
-	}
-	p.links = nil
-
-	// Close objects
-	if p.objs != nil {
-		if err := p.objs.Close(); err != nil {
-			log.Printf("Warning: failed to close BPF objects: %v", err)
-		}
-		p.objs = nil
+	// Try to unpin the event counter map
+	if err := p.unpinEventCounterMap(); err != nil {
+		log.Printf("Warning: failed to unpin event counter map: %v", err)
 	}
 
-	p.attached = false
-	log.Println("BPF program detached successfully")
+	log.Println("Pinned resources cleanup completed")
 	return nil
 }