From 14b027a3c200f9ddad89afa1f4b1e70b7e9ad4a2 Mon Sep 17 00:00:00 2001
From: Kyle Havlovitz <kylehav@gmail.com>
Date: Tue, 29 Aug 2017 18:27:47 -0700
Subject: [PATCH] Add segment addr field to tags for LAN flood joiner

---
 agent/consul/flood.go        |  4 ++--
 agent/consul/server.go       |  2 +-
 agent/consul/server_serf.go  |  1 +
 agent/metadata/server.go     |  7 +++++++
 agent/router/serf_flooder.go | 11 ++++++++++-
 5 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/agent/consul/flood.go b/agent/consul/flood.go
index 0e233eab9c..1287d0d7ad 100644
--- a/agent/consul/flood.go
+++ b/agent/consul/flood.go
@@ -24,7 +24,7 @@ func (s *Server) FloodNotify() {
 // Flood is a long-running goroutine that floods servers from the LAN to the
 // given global Serf instance, such as the WAN. This will exit once either of
 // the Serf instances are shut down.
-func (s *Server) Flood(portFn router.FloodPortFn, global *serf.Serf) {
+func (s *Server) Flood(addrFn router.FloodAddrFn, portFn router.FloodPortFn, global *serf.Serf) {
 	s.floodLock.Lock()
 	floodCh := make(chan struct{})
 	s.floodCh = append(s.floodCh, floodCh)
@@ -61,6 +61,6 @@ func (s *Server) Flood(portFn router.FloodPortFn, global *serf.Serf) {
 		}
 
 	FLOOD:
-		router.FloodJoins(s.logger, portFn, s.config.Datacenter, s.serfLAN, global)
+		router.FloodJoins(s.logger, addrFn, portFn, s.config.Datacenter, s.serfLAN, global)
 	}
 }
diff --git a/agent/consul/server.go b/agent/consul/server.go
index 4bbb745aef..25091c557b 100644
--- a/agent/consul/server.go
+++ b/agent/consul/server.go
@@ -410,7 +410,7 @@ func NewServerLogger(config *Config, logger *log.Logger, tokens *token.Store) (*
 		}
 		return 0, false
 	}
-	go s.Flood(portFn, s.serfWAN)
+	go s.Flood(nil, portFn, s.serfWAN)
 
 	// Start monitoring leadership. This must happen after Serf is set up
 	// since it can fire events when leadership is obtained.
diff --git a/agent/consul/server_serf.go b/agent/consul/server_serf.go
index 6f2805cc99..af8dbec297 100644
--- a/agent/consul/server_serf.go
+++ b/agent/consul/server_serf.go
@@ -43,6 +43,7 @@ func (s *Server) setupSerf(conf *serf.Config, ch chan serf.Event, path string, w
 	conf.Tags["segment"] = segment
 	if segment == "" {
 		for _, s := range s.config.Segments {
+			conf.Tags["segment_addr_"+s.Name] = s.Advertise
 			conf.Tags["segment_port_"+s.Name] = fmt.Sprintf("%d", s.Port)
 		}
 	}
diff --git a/agent/metadata/server.go b/agent/metadata/server.go
index 7ff6c69ce2..5a539a97c5 100644
--- a/agent/metadata/server.go
+++ b/agent/metadata/server.go
@@ -33,6 +33,7 @@ type Server struct {
 	Datacenter   string
 	Segment      string
 	Port         int
+	SegmentAddrs map[string]string
 	SegmentPorts map[string]int
 	WanJoinPort  int
 	Bootstrap    bool
@@ -96,8 +97,13 @@ func IsConsulServer(m serf.Member) (bool, *Server) {
 		return false, nil
 	}
 
+	segment_addrs := make(map[string]string)
 	segment_ports := make(map[string]int)
 	for name, value := range m.Tags {
+		if strings.HasPrefix(name, "segment_addr_") {
+			segment_addrs[strings.TrimPrefix(name, "segment_addr_")] = value
+		}
+
 		if strings.HasPrefix(name, "segment_port_") {
 			segment_port, err := strconv.Atoi(value)
 			if err != nil {
@@ -146,6 +152,7 @@ func IsConsulServer(m serf.Member) (bool, *Server) {
 		Datacenter:   datacenter,
 		Segment:      segment,
 		Port:         port,
+		SegmentAddrs: segment_addrs,
 		SegmentPorts: segment_ports,
 		WanJoinPort:  wan_join_port,
 		Bootstrap:    bootstrap,
diff --git a/agent/router/serf_flooder.go b/agent/router/serf_flooder.go
index fee912dd65..ceea2f14fd 100644
--- a/agent/router/serf_flooder.go
+++ b/agent/router/serf_flooder.go
@@ -10,6 +10,10 @@ import (
 	"github.com/hashicorp/serf/serf"
 )
 
+// FloodAddrFn gets the address to use for a given server when flood-joining. This
+// will return false if it doesn't have one.
+type FloodAddrFn func(*metadata.Server) (string, bool)
+
 // FloodPortFn gets the port to use for a given server when flood-joining. This
 // will return false if it doesn't have one.
 type FloodPortFn func(*metadata.Server) (int, bool)
@@ -19,7 +23,7 @@ type FloodPortFn func(*metadata.Server) (int, bool)
 // local area are of the form <node> and those in the global area are of the
 // form <node>.<dc> as is done for WAN and general network areas in Consul
 // Enterprise.
-func FloodJoins(logger *log.Logger, portFn FloodPortFn,
+func FloodJoins(logger *log.Logger, addrFn FloodAddrFn, portFn FloodPortFn,
 	localDatacenter string, localSerf *serf.Serf, globalSerf *serf.Serf) {
 
 	// Names in the global Serf have the datacenter suffixed.
@@ -64,6 +68,11 @@ func FloodJoins(logger *log.Logger, portFn FloodPortFn,
 			logger.Printf("[DEBUG] consul: Failed to flood-join %q (bad address %q): %v",
 				server.Name, server.Addr.String(), err)
 		}
+		if addrFn != nil {
+			if a, ok := addrFn(server); ok {
+				addr = a
+			}
+		}
 
 		// Let the callback see if it can get the port number, otherwise
 		// leave it blank to behave as if we just supplied an address.