diff --git a/cmd/start.go b/cmd/start.go index 44b16f5..b481bcd 100644 --- a/cmd/start.go +++ b/cmd/start.go @@ -18,6 +18,7 @@ const ( flagGRPCAddress = "grpc" flagOperator = "operator" flagSentry = "sentry" + flagSentryLabel = "label" flagMaxReadSize = "max-read-size" ) @@ -71,9 +72,10 @@ func startCmd() *cobra.Command { // if we're running in kubernetes, we can auto-discover sentries operator, _ := cmd.Flags().GetBool(flagOperator) sentries, _ := cmd.Flags().GetStringArray(flagSentry) + labels, _ := cmd.Flags().GetStringArray(flagSentryLabel) maxReadSize, _ := cmd.Flags().GetInt(flagMaxReadSize) - watcher, err := NewSentryWatcher(ctx, logger, all, hc, operator, sentries, maxReadSize) + watcher, err := NewSentryWatcher(ctx, labels, logger, all, hc, operator, sentries, maxReadSize) if err != nil { return err } @@ -88,6 +90,7 @@ func startCmd() *cobra.Command { cmd.Flags().StringArrayP(flagListen, "l", nil, "Privval listen addresses for the proxy (e.g. tcp://0.0.0.0:1234)") cmd.Flags().StringArrayP(flagSentry, "s", nil, "Privval connect addresses for the proxy") + cmd.Flags().StringArrayP(flagSentryLabel, "L", nil, "the label of the sentry to connect to") cmd.Flags().BoolP(flagOperator, "o", true, "Use this when running in kubernetes with the Cosmos Operator to auto-discover sentries") cmd.Flags().StringP(flagGRPCAddress, "g", "", "GRPC address for the proxy") cmd.Flags().BoolP(flagAll, "a", false, "Connect to sentries on all nodes") diff --git a/cmd/watcher.go b/cmd/watcher.go index d6f588a..c52a18f 100644 --- a/cmd/watcher.go +++ b/cmd/watcher.go @@ -6,6 +6,7 @@ import ( "fmt" "net" "os" + "strings" "time" cometlog "github.com/cometbft/cometbft/libs/log" @@ -26,6 +27,7 @@ type SentryWatcher struct { all bool client *kubernetes.Clientset hc signer.HorcruxConnection + labels string log cometlog.Logger node string operator bool @@ -38,6 +40,7 @@ type SentryWatcher struct { func NewSentryWatcher( ctx context.Context, + labels []string, logger cometlog.Logger, all bool, // should we connect to sentries on all nodes, or just this node? hc signer.HorcruxConnection, @@ -83,11 +86,22 @@ func NewSentryWatcher( persistentSentries[i] = signer.NewReconnRemoteSigner(sentry, logger, hc, dialer, maxReadSize) } + uniqueLabelMap := make(map[string]bool) + labels = append(labels, labelCosmosSentry) + finalLabels := []string{} + for _, label := range labels { + if _, exists := uniqueLabelMap[label]; !exists { + uniqueLabelMap[label] = true + finalLabels = append(finalLabels, label) + } + } + return &SentryWatcher{ all: all, client: clientset, done: make(chan struct{}), hc: hc, + labels: strings.Join(finalLabels, ","), log: logger, node: thisNode, operator: operator, @@ -150,7 +164,7 @@ func (w *SentryWatcher) reconcileSentries( configNodes := make([]string, 0) services, err := w.client.CoreV1().Services("").List(ctx, metav1.ListOptions{ - LabelSelector: labelCosmosSentry, + LabelSelector: w.labels, }) if err != nil {