outbound.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"maps"
	"os"
	"strings"
	"time"

	"github.com/containernetworking/cni/pkg/skel"
	"github.com/containernetworking/cni/pkg/types"
	current "github.com/containernetworking/cni/pkg/types/100"
	"github.com/containernetworking/cni/pkg/version"
	"github.com/containernetworking/plugins/pkg/utils"
	bv "github.com/containernetworking/plugins/pkg/utils/buildversion"
	"github.com/ncode/cni-outbound/pkg/iptables"
)

type LogConfig struct {
	Enable    bool   `json:"enable"`
	Directory string `json:"directory"`
}

type PluginConf struct {
	types.NetConf

	MainChainName string                  `json:"mainChainName"`
	DefaultAction string                  `json:"defaultAction"`
	OutboundRules []iptables.OutboundRule `json:"outboundRules"`
	Logging       LogConfig               `json:"logging"`
	Metadata      map[string]string       `json:"metadata"`
	DryRun        bool                    `json:"dryRun"`
	LogDrops      bool                    `json:"logDrops"`
}

func getLogAttrs() slog.Attr {
	var attrs []any

	if metadata != nil {
		for k, v := range metadata {
			attrs = append(attrs, slog.String(k, v))
		}
	}

	return slog.Group("metadata", attrs...)
}

var (
	logger             = slog.New(slog.NewTextHandler(io.Discard, nil))
	newIPTablesManager = func(conf *PluginConf) (iptables.Manager, error) {
		return iptables.NewIPTablesManager(conf.MainChainName, conf.DefaultAction, conf.DryRun, conf.LogDrops)
	}
	metadata = map[string]string{}
)

func setupLogging(conf *PluginConf) error {
	if !conf.Logging.Enable {
		return nil
	}

	var logWriter *os.File
	if conf.Logging.Directory == "" {
		conf.Logging.Directory = "/var/log/cni"
	}

	currentDate := time.Now().Format("2006-01-02")
	logFileName := fmt.Sprintf("%s/outbound-%s.log", strings.TrimSuffix(conf.Logging.Directory, "/"), currentDate)

	file, err := os.OpenFile(logFileName, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return fmt.Errorf("failed to open log file: %v", err)
	}
	logWriter = file

	opts := slog.HandlerOptions{
		AddSource: true,
		Level:     slog.LevelInfo,
	}
	handler := slog.NewJSONHandler(logWriter, &opts)
	logger = slog.New(handler)
	return nil
}

func generateChainName(netName, containerID string) string {
	return utils.MustFormatChainNameWithPrefix(netName, containerID, "OUT-")
}

func parseArgs(args, containerID string) ([]iptables.OutboundRule, map[string]string, error) {
	logger.Log(context.Background(), slog.LevelInfo,
		"Parsing CNI arguments",
		getLogAttrs(),
		slog.String("details", args),
	)

	metadata := make(map[string]string)
	var additionalRules []iptables.OutboundRule

	if args == "" {
		logger.Log(context.Background(), slog.LevelInfo,
			"No additional args provided",
			getLogAttrs(),
		)
		return nil, metadata, nil
	}

	// Skip CNI env vars that are already available in skel.CmdArgs
	skipKeys := map[string]bool{
		"CNI_COMMAND":     true,
		"CNI_CONTAINERID": true,
		"CNI_PATH":        true,
		"IgnoreUnknown":   true,
	}

	kvs := strings.Split(args, ";")
	for _, kv := range kvs {
		parts := strings.SplitN(kv, "=", 2)
		if len(parts) != 2 {
			continue
		}

		key, value := parts[0], parts[1]
		if ok := skipKeys[key]; ok {
			continue
		} else if key == "outbound.additional_rules" {
			logger.Log(context.Background(), slog.LevelInfo,
				"Found outbound.additional_rules",
				getLogAttrs(),
				slog.String("rules", value),
			)

			if err := json.Unmarshal([]byte(value), &additionalRules); err != nil {
				logger.Log(context.Background(), slog.LevelError,
					"Failed to parse additional rules",
					getLogAttrs(),
					slog.Any("error", err),
				)
				return nil, nil, fmt.Errorf("failed to parse additional rules from CNI args: %v", err)
			}
		} else {
			metadata[key] = value
			logger.Log(context.Background(), slog.LevelInfo,
				"Found metadata",
				getLogAttrs(),
				slog.String("key", key),
			)
		}
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Parsed args",
		getLogAttrs(),
		slog.Int("ruleCount", len(additionalRules)),
		slog.Int("metadataCount", len(metadata)),
	)

	return additionalRules, metadata, nil
}

func parseConfig(stdin []byte, args, containerID string) (*PluginConf, error) {
	conf := PluginConf{}

	// Parse additional rules and metadata from args
	additionalRules, argsMetadata, err := parseArgs(args, containerID)
	if err != nil {
		return nil, err
	}

	if err := json.Unmarshal(stdin, &conf); err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to parse network configuration",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return nil, fmt.Errorf("failed to parse network configuration: %v", err)
	}

	if err := setupLogging(&conf); err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to setup logging",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return nil, fmt.Errorf("failed to setup logging: %v", err)
	}

	if conf.DryRun {
		logger.Log(context.Background(), slog.LevelInfo,
			"Dry run mode enabled - traffic will be logged but not blocked",
			getLogAttrs(),
		)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Parsing configuration",
		getLogAttrs(),
	)

	// Parse prevResult if it exists
	if conf.RawPrevResult != nil {
		if err := version.ParsePrevResult(&conf.NetConf); err != nil {
			logger.Log(context.Background(), slog.LevelError,
				"Could not parse prevResult",
				getLogAttrs(),
				slog.Any("error", err),
			)
			return nil, fmt.Errorf("could not parse prevResult: %v", err)
		}

		result, err := current.NewResultFromResult(conf.PrevResult)
		if err != nil {
			logger.Log(context.Background(), slog.LevelError,
				"Failed to convert prevResult to current.Result",
				getLogAttrs(),
				slog.Any("error", err),
			)
			return nil, fmt.Errorf("failed to convert prevResult to current.Result: %v", err)
		}

		if len(result.Interfaces) == 0 {
			return nil, fmt.Errorf("invalid prevResult structure: missing interfaces")
		}

		if len(result.IPs) == 0 {
			return nil, fmt.Errorf("invalid prevResult structure: missing ips")
		}

		conf.PrevResult = result
	}

	// Initialize metadata map only if we have metadata to add
	if len(argsMetadata) > 0 && conf.Metadata == nil {
		conf.Metadata = make(map[string]string)
	}

	// Add rules from args if any exist
	if len(additionalRules) > 0 {
		logger.Log(context.Background(), slog.LevelInfo,
			"Appending additional rules",
			getLogAttrs(),
			slog.Int("ruleCount", len(additionalRules)),
		)
		conf.OutboundRules = append(conf.OutboundRules, additionalRules...)
	}

	// Merge metadata from args only if we have metadata
	if len(argsMetadata) > 0 {
		maps.Copy(conf.Metadata, argsMetadata)
		maps.Copy(metadata, conf.Metadata)
	}

	// Set defaults if needed
	if conf.MainChainName == "" {
		logger.Log(context.Background(), slog.LevelInfo,
			"Using default MainChainName: CNI-OUTBOUND",
			getLogAttrs(),
		)
		conf.MainChainName = "CNI-OUTBOUND"
	}

	if conf.DefaultAction == "" {
		logger.Log(context.Background(), slog.LevelInfo,
			"Using default DefaultAction: DROP",
			getLogAttrs(),
		)
		conf.DefaultAction = "DROP"
	}

	return &conf, nil
}

func cmdAdd(args *skel.CmdArgs) error {
	metadata["component"] = "CNI-Outbound"
	metadata["containerID"] = args.ContainerID

	conf, err := parseConfig(args.StdinData, args.Args, args.ContainerID)
	if err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to parse config",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return err
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"CNI ADD called",
		getLogAttrs(),
	)

	logger.Log(context.Background(), slog.LevelInfo,
		"Creating IPTablesManager",
		getLogAttrs(),
	)

	iptManager, err := newIPTablesManager(conf)
	if err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to create IPTablesManager",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to create IPTablesManager: %v", err)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Ensuring main chain exists",
		getLogAttrs(),
	)

	if err := iptManager.EnsureMainChainExists(); err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to ensure main chain exists",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to ensure main chain exists: %v", err)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Creating container chain",
		getLogAttrs(),
	)

	containerChain := generateChainName(conf.Name, args.ContainerID)
	if err := iptManager.CreateContainerChain(containerChain); err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to create container chain",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to create container chain: %v", err)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Adding rules to container chain",
	)

	for _, rule := range conf.OutboundRules {
		logger.Log(context.Background(), slog.LevelInfo,
			"Adding rule",
			getLogAttrs(),
			slog.Any("rule", rule),
		)
		if err := iptManager.AddRule(containerChain, rule); err != nil {
			logger.Log(context.Background(), slog.LevelError,
				"Failed to add rule to container chain",
				getLogAttrs(),
				slog.Any("error", err),
				slog.Any("rule", rule),
			)
			return fmt.Errorf("failed to add rule to container chain: %v", err)
		}
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Adding jump rule to main chain",
		getLogAttrs(),
	)

	var result *current.Result
	if conf.PrevResult == nil {
		// If there's no previous result, return an error
		return fmt.Errorf("no prevResult found")
	}

	result, err = current.NewResultFromResult(conf.PrevResult)
	if err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to parse prevResult",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to parse prevResult: %v", err)
	}

	containerIP := result.IPs[0].Address.IP.String()
	logger.Log(context.Background(), slog.LevelInfo,
		"Container IP obtained",
		getLogAttrs(),
		slog.String("ip", containerIP),
	)

	if err := iptManager.AddJumpRule(containerIP, containerChain); err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to add jump rule to main chain",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to add jump rule to main chain: %v", err)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"CNI ADD completed successfully",
		getLogAttrs(),
	)
	return types.PrintResult(result, conf.CNIVersion)
}

func cmdDel(args *skel.CmdArgs) error {
	metadata["component"] = "CNI-Outbound"
	metadata["containerID"] = args.ContainerID

	conf, err := parseConfig(args.StdinData, args.Args, args.ContainerID)
	if err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to parse config",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return err
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"CNI DEL called",
		getLogAttrs(),
	)

	logger.Log(context.Background(), slog.LevelInfo,
		"Creating IPTablesManager",
		getLogAttrs(),
	)

	iptManager, err := newIPTablesManager(conf)
	if err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to create IPTablesManager",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to create IPTablesManager: %v", err)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Removing container chain",
		getLogAttrs(),
	)

	containerChain := generateChainName(conf.Name, args.ContainerID)
	if err := iptManager.RemoveJumpRuleByTargetChain(containerChain); err != nil {
		logger.Log(context.Background(), slog.LevelWarn,
			"Failed to remove jump rule from main chain",
			getLogAttrs(),
			slog.Any("error", err),
		)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Clearing and deleting container chain",
		getLogAttrs(),
	)

	if err := iptManager.ClearAndDeleteChain(containerChain); err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to clear and delete container chain",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to clear and delete container chain: %v", err)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"CNI DEL completed successfully",
		getLogAttrs(),
	)
	return nil
}

func cmdCheck(args *skel.CmdArgs) error {
	metadata["component"] = "CNI-Outbound"
	metadata["containerID"] = args.ContainerID

	conf, err := parseConfig(args.StdinData, args.Args, args.ContainerID)
	if err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to parse config",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return err
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"CNI CHECK called",
		getLogAttrs(),
	)

	logger.Log(context.Background(), slog.LevelInfo,
		"Creating IPTablesManager",
		getLogAttrs(),
	)

	iptManager, err := newIPTablesManager(conf)
	if err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to create IPTablesManager",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to create IPTablesManager: %v", err)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Checking if main chain exists",
		getLogAttrs(),
	)

	exists, err := iptManager.ChainExists(conf.MainChainName)
	if err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to check if main chain exists",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to check if main chain exists: %v", err)
	}
	if !exists {
		logger.Log(context.Background(), slog.LevelError,
			"Main chain does not exist",
			getLogAttrs(),
			slog.String("chain", conf.MainChainName),
		)
		return fmt.Errorf("main chain %s does not exist", conf.MainChainName)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Checking container chain",
		getLogAttrs(),
	)

	containerChain := generateChainName(conf.Name, args.ContainerID)
	exists, err = iptManager.ChainExists(containerChain)
	if err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Failed to check if container chain exists",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("failed to check if container chain exists: %v", err)
	}
	if !exists {
		logger.Log(context.Background(), slog.LevelError,
			"Container chain does not exist",
			getLogAttrs(),
			slog.String("chain", containerChain),
		)
		return fmt.Errorf("container chain %s does not exist", containerChain)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"Verifying rules in container chain",
		getLogAttrs(),
	)

	if err := iptManager.VerifyRules(containerChain, conf.OutboundRules); err != nil {
		logger.Log(context.Background(), slog.LevelError,
			"Rule verification failed",
			getLogAttrs(),
			slog.Any("error", err),
		)
		return fmt.Errorf("rule verification failed: %v", err)
	}

	logger.Log(context.Background(), slog.LevelInfo,
		"CNI CHECK completed successfully",
		getLogAttrs(),
	)
	return nil
}

func main() {
	skel.PluginMain(cmdAdd, cmdCheck, cmdDel, version.All, bv.BuildString("outbound"))
}