proxy-pool/cmd/proxy-checker/main.go

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/checker/bootstrap"
	"proxy-pool/internal/config"
	"proxy-pool/internal/platform/logging"
)

const (
	configEnvironment              = "PROXY_POOL_CONFIG"
	controlPlaneAddressEnvironment = "PROXY_POOL_CONTROL_PLANE_ADDRESS"
	checkerIDEnvironment           = "PROXY_POOL_CHECKER_ID"
	instanceIDEnvironment          = "PROXY_POOL_CHECKER_INSTANCE_ID"
	autoIdentityEnvironment        = "PROXY_POOL_AUTO_IDENTITY"
	maxInFlightEnvironment         = "PROXY_POOL_CHECKER_MAX_IN_FLIGHT"
)

type environmentLookup func(string) string
type checkerRun func(context.Context, bootstrap.Options) error

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	os.Exit(execute(ctx, os.Args[1:], os.Getenv, bootstrap.Run, os.Stderr))
}

func execute(ctx context.Context, args []string, getenv environmentLookup, run checkerRun, stderr io.Writer) int {
	flags := flag.NewFlagSet("proxy-checker", flag.ContinueOnError)
	flags.SetOutput(stderr)
	configPath := flags.String("config", "", "configuration file path")
	controlPlaneAddress := flags.String("control-plane", "", "remote Controller control-plane address")
	checkerID := flags.String("checker-id", "", "unique Checker identifier")
	instanceID := flags.String("instance-id", "", "unique Checker process instance identifier")
	autoIdentity := flags.Bool("auto-identity", false, "derive Checker and missing instance IDs from the mTLS SPIFFE certificate")
	maxInFlight := flags.Int("max-in-flight", 0, "maximum concurrent tasks")
	levels := flags.String("levels", "basic,egress,target", "supported levels: basic,egress,target")
	if err := flags.Parse(args); err != nil {
		if errors.Is(err, flag.ErrHelp) {
			return 0
		}
		return 2
	}
	if flags.NArg() != 0 {
		_, _ = fmt.Fprintln(stderr, "proxy-checker: unexpected positional arguments")
		return 2
	}
	if getenv != nil {
		setIfEmpty(configPath, getenv(configEnvironment))
		setIfEmpty(controlPlaneAddress, getenv(controlPlaneAddressEnvironment))
		setIfEmpty(checkerID, getenv(checkerIDEnvironment))
		setIfEmpty(instanceID, getenv(instanceIDEnvironment))
		if *maxInFlight == 0 {
			if value, err := strconv.Atoi(getenv(maxInFlightEnvironment)); err == nil {
				*maxInFlight = value
			}
		}
		if !flagWasSet(flags, "auto-identity") {
			if value, err := strconv.ParseBool(getenv(autoIdentityEnvironment)); err == nil {
				*autoIdentity = value
			}
		}
	}
	supportedLevels, err := parseLevels(*levels)
	missingIdentity := !*autoIdentity && (!validValue(*checkerID) || !validValue(*instanceID))
	invalidProvidedIdentity := (*checkerID != "" && !validValue(*checkerID)) || (*instanceID != "" && !validValue(*instanceID))
	if ctx == nil || run == nil || err != nil || !validValue(*configPath) || !validValue(*controlPlaneAddress) ||
		missingIdentity || invalidProvidedIdentity || *maxInFlight <= 0 {
		_, _ = fmt.Fprintf(stderr,
			"proxy-checker: -config, -control-plane and positive -max-in-flight are required; -checker-id/-instance-id are required unless -auto-identity is enabled; "+
				"environment fallbacks: %s, %s, %s, %s, %s, %s\n",
			configEnvironment, controlPlaneAddressEnvironment, checkerIDEnvironment, instanceIDEnvironment, maxInFlightEnvironment, autoIdentityEnvironment,
		)
		return 2
	}
	err = run(ctx, bootstrap.Options{
		ConfigPath: *configPath, Resolver: config.OSResolver{}, ControlPlaneAddress: *controlPlaneAddress,
		CheckerID: *checkerID, InstanceID: *instanceID, AutoIdentity: *autoIdentity, MaxInFlight: *maxInFlight, SupportedLevels: supportedLevels,
	})
	if err == nil || (errors.Is(err, context.Canceled) && ctx.Err() != nil) {
		return 0
	}
	logging.WriteProcessError(stderr, "proxy-checker", err)
	return 1
}

func flagWasSet(flags *flag.FlagSet, name string) bool {
	set := false
	flags.Visit(func(item *flag.Flag) {
		if item.Name == name {
			set = true
		}
	})
	return set
}

func parseLevels(value string) ([]controlplanev1.CheckLevel, error) {
	levels := make([]controlplanev1.CheckLevel, 0, 3)
	seen := make(map[controlplanev1.CheckLevel]struct{})
	for _, item := range strings.Split(value, ",") {
		var level controlplanev1.CheckLevel
		switch strings.ToLower(strings.TrimSpace(item)) {
		case "basic":
			level = controlplanev1.CheckLevel_CHECK_LEVEL_BASIC
		case "egress":
			level = controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS
		case "target":
			level = controlplanev1.CheckLevel_CHECK_LEVEL_TARGET
		default:
			return nil, errors.New("invalid checker level")
		}
		if _, duplicate := seen[level]; duplicate {
			return nil, errors.New("duplicate checker level")
		}
		seen[level] = struct{}{}
		levels = append(levels, level)
	}
	if len(levels) == 0 {
		return nil, errors.New("empty checker level set")
	}
	return levels, nil
}

func setIfEmpty(target *string, value string) {
	if target != nil && *target == "" {
		*target = value
	}
}

func validValue(value string) bool {
	return value != "" && strings.TrimSpace(value) == value
}