From c669a2b155b7f30c15fe41872449899a2e7a3a5d Mon Sep 17 00:00:00 2001
From: Asim Aslam <asim@aslam.me>
Date: Mon, 9 Sep 2019 05:11:25 -0700
Subject: [PATCH] Use .micro domain for mdns

---
 registry/mdns/mdns.go     | 12 +++++++++++
 registry/mdns_registry.go | 44 ++++++++++++++++++++++++++++++++-------
 registry/mdns_watcher.go  |  9 +++++---
 3 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/registry/mdns/mdns.go b/registry/mdns/mdns.go
index 6739c694..f84e8961 100644
--- a/registry/mdns/mdns.go
+++ b/registry/mdns/mdns.go
@@ -2,6 +2,8 @@
 package mdns
 
 import (
+	"context"
+
 	"github.com/micro/go-micro/registry"
 )
 
@@ -9,3 +11,13 @@ import (
 func NewRegistry(opts ...registry.Option) registry.Registry {
 	return registry.NewRegistry(opts...)
 }
+
+// Domain sets the mdnsDomain
+func Domain(d string) registry.Option {
+	return func(o *registry.Options) {
+		if o.Context == nil {
+			o.Context = context.Background()
+		}
+		o.Context = context.WithValue(o.Context, "mdns.domain", d)
+	}
+}
diff --git a/registry/mdns_registry.go b/registry/mdns_registry.go
index 8e68948a..039710a6 100644
--- a/registry/mdns_registry.go
+++ b/registry/mdns_registry.go
@@ -14,6 +14,11 @@ import (
 	hash "github.com/mitchellh/hashstructure"
 )
 
+var (
+	// use a .micro domain rather than .local
+	mdnsDomain = "micro"
+)
+
 type mdnsTxt struct {
 	Service   string
 	Version   string
@@ -29,6 +34,8 @@ type mdnsEntry struct {
 
 type mdnsRegistry struct {
 	opts Options
+	// the mdns domain
+	domain string
 
 	sync.Mutex
 	services map[string][]*mdnsEntry
@@ -36,11 +43,25 @@ type mdnsRegistry struct {
 
 func newRegistry(opts ...Option) Registry {
 	options := Options{
+		Context: context.Background(),
 		Timeout: time.Millisecond * 100,
 	}
 
+	for _, o := range opts {
+		o(&options)
+	}
+
+	// set the domain
+	domain := mdnsDomain
+
+	d, ok := options.Context.Value("mdns.domain").(string)
+	if ok {
+		domain = d
+	}
+
 	return &mdnsRegistry{
 		opts:     options,
+		domain:   domain,
 		services: make(map[string][]*mdnsEntry),
 	}
 }
@@ -66,7 +87,7 @@ func (m *mdnsRegistry) Register(service *Service, opts ...RegisterOption) error
 		s, err := mdns.NewMDNSService(
 			service.Name,
 			"_services",
-			"",
+			m.domain+".",
 			"",
 			9999,
 			[]net.IP{net.ParseIP("0.0.0.0")},
@@ -141,7 +162,7 @@ func (m *mdnsRegistry) Register(service *Service, opts ...RegisterOption) error
 		s, err := mdns.NewMDNSService(
 			node.Id,
 			service.Name,
-			"",
+			m.domain+".",
 			"",
 			port,
 			[]net.IP{net.ParseIP(host)},
@@ -214,6 +235,8 @@ func (m *mdnsRegistry) GetService(service string) ([]*Service, error) {
 	p.Context, _ = context.WithTimeout(context.Background(), m.opts.Timeout)
 	// set entries channel
 	p.Entries = entries
+	// set the domain
+	p.Domain = m.domain
 
 	go func() {
 		for {
@@ -223,7 +246,9 @@ func (m *mdnsRegistry) GetService(service string) ([]*Service, error) {
 				if p.Service == "_services" {
 					continue
 				}
-
+				if p.Domain != m.domain {
+					continue
+				}
 				if e.TTL == 0 {
 					continue
 				}
@@ -288,6 +313,8 @@ func (m *mdnsRegistry) ListServices() ([]*Service, error) {
 	p.Context, _ = context.WithTimeout(context.Background(), m.opts.Timeout)
 	// set entries channel
 	p.Entries = entries
+	// set domain
+	p.Domain = m.domain
 
 	var services []*Service
 
@@ -298,7 +325,9 @@ func (m *mdnsRegistry) ListServices() ([]*Service, error) {
 				if e.TTL == 0 {
 					continue
 				}
-
+				if !strings.HasSuffix(e.Name, p.Domain+".") {
+					continue
+				}
 				name := strings.TrimSuffix(e.Name, "."+p.Service+"."+p.Domain+".")
 				if !serviceMap[name] {
 					serviceMap[name] = true
@@ -329,9 +358,10 @@ func (m *mdnsRegistry) Watch(opts ...WatchOption) (Watcher, error) {
 	}
 
 	md := &mdnsWatcher{
-		wo:   wo,
-		ch:   make(chan *mdns.ServiceEntry, 32),
-		exit: make(chan struct{}),
+		wo:     wo,
+		ch:     make(chan *mdns.ServiceEntry, 32),
+		exit:   make(chan struct{}),
+		domain: m.domain,
 	}
 
 	go func() {
diff --git a/registry/mdns_watcher.go b/registry/mdns_watcher.go
index e1c326ff..ce13866f 100644
--- a/registry/mdns_watcher.go
+++ b/registry/mdns_watcher.go
@@ -11,6 +11,8 @@ type mdnsWatcher struct {
 	wo   WatchOptions
 	ch   chan *mdns.ServiceEntry
 	exit chan struct{}
+	// the mdns domain
+	domain string
 }
 
 func (m *mdnsWatcher) Next() (*Result, error) {
@@ -46,13 +48,14 @@ func (m *mdnsWatcher) Next() (*Result, error) {
 				Endpoints: txt.Endpoints,
 			}
 
-			// TODO: don't hardcode .local.
-			if !strings.HasSuffix(e.Name, "."+service.Name+".local.") {
+			// skip anything without the domain we care about
+			suffix := fmt.Sprintf(".%s.%s.", service.Name, m.domain)
+			if !strings.HasSuffix(e.Name, suffix) {
 				continue
 			}
 
 			service.Nodes = append(service.Nodes, &Node{
-				Id:       strings.TrimSuffix(e.Name, "."+service.Name+".local."),
+				Id:       strings.TrimSuffix(e.Name, suffix),
 				Address:  fmt.Sprintf("%s:%d", e.AddrV4.String(), e.Port),
 				Metadata: txt.Metadata,
 			})