From ca0e8562f3ff133582cd56de205c02a2abd86fcb Mon Sep 17 00:00:00 2001
From: Jeffrey Picard <jeff@jeffreypicard.com>
Date: Wed, 24 Nov 2021 18:36:19 -0500
Subject: [PATCH] Add disable-federation flag, default false

---
 server/args.go       | 65 +++++++++++++++++++++++---------------------
 server/federation.go | 10 +++++--
 2 files changed, 42 insertions(+), 33 deletions(-)

diff --git a/server/args.go b/server/args.go
index 2834e5e..00b24c9 100644
--- a/server/args.go
+++ b/server/args.go
@@ -37,20 +37,21 @@ type Args struct {
 }
 
 const (
-	DefaultHost            = "0.0.0.0"
-	DefaultPort            = "50051"
-	DefaultEsHost          = "http://localhost"
-	DefaultEsIndex         = "claims"
-	DefaultEsPort          = "9200"
-	DefaultPrometheusPort  = "2112"
-	DefaultRefreshDelta    = 5
-	DefaultCacheTTL        = 5
-	DefaultPeerFile        = "peers.txt"
-	DefaultCountry         = "US"
-	DefaultLoadPeers       = true
-	DefaultStartPrometheus = true
-	DefaultStartUDP        = true
-	DefaultWritePeers      = true
+	DefaultHost              = "0.0.0.0"
+	DefaultPort              = "50051"
+	DefaultEsHost            = "http://localhost"
+	DefaultEsIndex           = "claims"
+	DefaultEsPort            = "9200"
+	DefaultPrometheusPort    = "2112"
+	DefaultRefreshDelta      = 5
+	DefaultCacheTTL          = 5
+	DefaultPeerFile          = "peers.txt"
+	DefaultCountry           = "US"
+	DefaultLoadPeers         = true
+	DefaultStartPrometheus   = true
+	DefaultStartUDP          = true
+	DefaultWritePeers        = true
+	DefaultDisableFederation = false
 )
 
 // GetEnvironment takes the environment variables as an array of strings
@@ -100,6 +101,7 @@ func ParseArgs(searchRequest *pb.SearchRequest) *Args {
 	startPrometheus := parser.Flag("", "start-prometheus", &argparse.Options{Required: false, Help: "Start prometheus server", Default: DefaultStartPrometheus})
 	startUdp := parser.Flag("", "start-udp", &argparse.Options{Required: false, Help: "Start UDP ping server", Default: DefaultStartUDP})
 	writePeers := parser.Flag("", "write-peers", &argparse.Options{Required: false, Help: "Write peer to disk as we learn about them", Default: DefaultWritePeers})
+	disableFederation := parser.Flag("", "disable-federation", &argparse.Options{Required: false, Help: "Disable server federation", Default: DefaultDisableFederation})
 
 	text := parser.String("", "text", &argparse.Options{Required: false, Help: "text query"})
 	name := parser.String("", "name", &argparse.Options{Required: false, Help: "name"})
@@ -118,23 +120,24 @@ func ParseArgs(searchRequest *pb.SearchRequest) *Args {
 	}
 
 	args := &Args{
-		CmdType:         SearchCmd,
-		Host:            *host,
-		Port:            *port,
-		EsHost:          *esHost,
-		EsPort:          *esPort,
-		PrometheusPort:  *prometheusPort,
-		EsIndex:         *esIndex,
-		RefreshDelta:    *refreshDelta,
-		CacheTTL:        *cacheTTL,
-		PeerFile:        *peerFile,
-		Country:         *country,
-		DisableEs:       *disableEs,
-		Debug:           *debug,
-		LoadPeers:       *loadPeers,
-		StartPrometheus: *startPrometheus,
-		StartUDP:        *startUdp,
-		WritePeers:      *writePeers,
+		CmdType:           SearchCmd,
+		Host:              *host,
+		Port:              *port,
+		EsHost:            *esHost,
+		EsPort:            *esPort,
+		PrometheusPort:    *prometheusPort,
+		EsIndex:           *esIndex,
+		RefreshDelta:      *refreshDelta,
+		CacheTTL:          *cacheTTL,
+		PeerFile:          *peerFile,
+		Country:           *country,
+		DisableEs:         *disableEs,
+		Debug:             *debug,
+		LoadPeers:         *loadPeers,
+		StartPrometheus:   *startPrometheus,
+		StartUDP:          *startUdp,
+		WritePeers:        *writePeers,
+		DisableFederation: *disableFederation,
 	}
 
 	if esHost, ok := environment["ELASTIC_HOST"]; ok {
diff --git a/server/federation.go b/server/federation.go
index e232ab0..abfabf1 100644
--- a/server/federation.go
+++ b/server/federation.go
@@ -267,7 +267,10 @@ func (s *Server) writePeers() {
 
 // notifyPeer takes a peer to notify and a new peer we just learned about
 // and calls AddPeer on the first.
-func notifyPeer(peerToNotify *FederatedServer, newPeer *FederatedServer) error {
+func (s *Server) notifyPeer(peerToNotify *FederatedServer, newPeer *FederatedServer) error {
+	if s.Args.DisableFederation {
+		return nil
+	}
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 
@@ -303,7 +306,7 @@ func (s *Server) notifyPeerSubs(newServer *FederatedServer) {
 	s.PeerSubsMut.RLock()
 	for key, peer := range s.PeerSubs {
 		log.Printf("Notifying peer %s of new node %+v\n", key, newServer)
-		err := notifyPeer(peer, newServer)
+		err := s.notifyPeer(peer, newServer)
 		if err != nil {
 			log.Println("Failed to send data to ", key)
 			log.Println(err)
@@ -327,6 +330,9 @@ func (s *Server) notifyPeerSubs(newServer *FederatedServer) {
 // if they're online, and adds them to our list of peer. If we're not currently
 // subscribed to a peer, it will also subscribe to it.
 func (s *Server) addPeer(msg *pb.ServerMessage, ping bool, subscribe bool) error {
+	if s.Args.DisableFederation {
+		return nil
+	}
 	// First thing we get our external ip if we don't have it, otherwise we
 	// could end up subscribed to our self, which is silly.
 	nilIP := net.IP{}