124 lines
4.2 KiB
Go
124 lines
4.2 KiB
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"
|
|
)
|
|
|
|
const (
|
|
configEnvironment = "PROXY_POOL_CONFIG"
|
|
controlPlaneAddressEnvironment = "PROXY_POOL_CONTROL_PLANE_ADDRESS"
|
|
checkerIDEnvironment = "PROXY_POOL_CHECKER_ID"
|
|
instanceIDEnvironment = "PROXY_POOL_CHECKER_INSTANCE_ID"
|
|
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")
|
|
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
|
|
}
|
|
}
|
|
}
|
|
supportedLevels, err := parseLevels(*levels)
|
|
if ctx == nil || run == nil || err != nil || !validValue(*configPath) || !validValue(*controlPlaneAddress) ||
|
|
!validValue(*checkerID) || !validValue(*instanceID) || *maxInFlight <= 0 {
|
|
_, _ = fmt.Fprintf(stderr,
|
|
"proxy-checker: -config, -control-plane, -checker-id, -instance-id and positive -max-in-flight are required; "+
|
|
"environment fallbacks: %s, %s, %s, %s, %s\n",
|
|
configEnvironment, controlPlaneAddressEnvironment, checkerIDEnvironment, instanceIDEnvironment, maxInFlightEnvironment,
|
|
)
|
|
return 2
|
|
}
|
|
err = run(ctx, bootstrap.Options{
|
|
ConfigPath: *configPath, Resolver: config.OSResolver{}, ControlPlaneAddress: *controlPlaneAddress,
|
|
CheckerID: *checkerID, InstanceID: *instanceID, MaxInFlight: *maxInFlight, SupportedLevels: supportedLevels,
|
|
})
|
|
if err == nil || (errors.Is(err, context.Canceled) && ctx.Err() != nil) {
|
|
return 0
|
|
}
|
|
_, _ = fmt.Fprintf(stderr, "proxy-checker: %v\n", err)
|
|
return 1
|
|
}
|
|
|
|
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
|
|
}
|