Updated mDNS

pkg/log/log.go

@@ -37,7 +37,11 @@ func (*Log) New(gopi.Config) error {
 func (this *Log) Print(args ...interface{}) {
 	this.Lock()
 	defer this.Unlock()
-	log.Print(args...)
+	if this.t != nil {
+		this.t.Log(args...)
+	} else {
+		log.Print(args...)
+	}
 }
 
 func (this *Log) IsDebug() bool {
@@ -52,21 +56,33 @@ func (this *Log) Debug(args ...interface{}) {
 	if this.IsDebug() {
 		this.Lock()
 		defer this.Unlock()
-		log.Print(args...)
+		if this.t != nil {
+			this.t.Log(args...)
+		} else {
+			log.Print(args...)
+		}
 	}
 }
 
 func (this *Log) Printf(fmt string, args ...interface{}) {
 	this.Lock()
 	defer this.Unlock()
-	log.Printf(fmt, args...)
+	if this.t != nil {
+		this.t.Logf(fmt, args...)
+	} else {
+		log.Printf(fmt, args...)
+	}
 }
 
 func (this *Log) Debugf(fmt string, args ...interface{}) {
 	if this.IsDebug() {
 		this.Lock()
 		defer this.Unlock()
-		log.Printf(fmt, args...)
+		if this.t != nil {
+			this.t.Logf(fmt, args...)
+		} else {
+			log.Printf(fmt, args...)
+		}
 	}
 }
 
@@ -75,6 +91,8 @@ func (this *Log) T() *testing.T {
 }
 
 func (this *Log) SetT(t *testing.T) {
+	this.Lock()
+	defer this.Unlock()
 	this.t = t
 	*this.debug = true
 }

pkg/mdns/discovery.go

@@ -22,24 +22,30 @@ type Discovery struct {
 
 const (
 	queryServices = "_services._dns-sd._udp"
-	queryRepeat   = 3
-	queryBackoff  = time.Millisecond * 200
+	queryRepeat   = 2
+	queryBackoff  = time.Millisecond * 50
 )
 
 ///////////////////////////////////////////////////////////////////////////////
 // RUN
 
 func (this *Discovery) Run(ctx context.Context) error {
-	// Subscribe to messages
+	// Subscribe to DNS messages
 	ch := this.Publisher.Subscribe()
 	defer this.Publisher.Unsubscribe(ch)
 
 FOR_LOOP:
 	for {
 		select {
 		case evt := <-ch:
-			if msg, ok := evt.(*dnsevent); ok {
-				NewServices(msg.Msg(), this.Listener.Domain())
+			if msg, ok := evt.(*msgevent); ok {
+				if services := NewServices(msg.Msg, this.Listener.Domain()).Services(); len(services) > 0 {
+					for _, service := range services {
+						if err := this.Publisher.Emit(service, true); err != nil {
+							this.Print(err)
+						}
+					}
+				}
 			}
 		case <-ctx.Done():
 			break FOR_LOOP
@@ -56,20 +62,99 @@ FOR_LOOP:
 ///////////////////////////////////////////////////////////////////////////////
 // PUBLIC METHODS
 
-func (this *Discovery) EnumerateServices(ctx context.Context) error {
+func (this *Discovery) Lookup(ctx context.Context, srv string) ([]gopi.ServiceRecord, error) {
 	this.WaitGroup.Add(1)
 	defer this.WaitGroup.Done()
 
+	// Collect services in goroutine
+	var wg sync.WaitGroup
+	ch := this.Publisher.Subscribe()
+	records := make([]*service, 0, 10)
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		defer this.Publisher.Unsubscribe(ch)
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case evt := <-ch:
+				if service, ok := evt.(*service); ok {
+					if service.Service() == srv {
+						records = append(records, service)
+					}
+				}
+			}
+		}
+	}()
+
+	// Query for lookup on all interfaces
+	zone := this.Listener.Domain()
+	if err := this.query(ctx, msgQueryLookup(srv, zone), 0); err != nil {
+		return nil, err
+	}
+
+	// Wait for end of collection of names
+	wg.Wait()
+
+	// Collect services
+	result := make([]gopi.ServiceRecord, 0, len(records))
+	for _, record := range records {
+		result = append(result, record)
+	}
+
+	// Return result
+	return result, nil
+}
+
+func (this *Discovery) EnumerateServices(ctx context.Context) ([]string, error) {
+	this.WaitGroup.Add(1)
+	defer this.WaitGroup.Done()
+
+	// Collect names in goroutine
+	var wg sync.WaitGroup
+	names := make(map[string]bool)
+	ch := this.Publisher.Subscribe()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		defer this.Publisher.Unsubscribe(ch)
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case evt := <-ch:
+				if srv, ok := evt.(*service); ok {
+					if srv.Service() == fqn(queryServices) && srv.ttl != 0 {
+						key := srv.Name()
+						names[key] = true
+					}
+				}
+			}
+		}
+	}()
+
 	// Query for services on all interfaces
 	zone := this.Listener.Domain()
 	if err := this.query(ctx, msgQueryServices(zone), 0); err != nil {
-		return err
+		return nil, err
 	}
 
-	// Wait for completion
-	<-ctx.Done()
+	// Wait for end of collection of names
+	wg.Wait()
 
-	return ctx.Err()
+	// Collect names
+	result := make([]string, 0, len(names))
+	for name := range names {
+		result = append(result, name)
+	}
+
+	return result, nil
+}
+
+func (this *Discovery) Serve(context.Context, []gopi.ServiceRecord) error {
+	return gopi.ErrNotImplemented
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -96,9 +181,16 @@ func (this *Discovery) query(ctx context.Context, msg *dns.Msg, iface int) error
 	}
 }
 
-func msgQueryServices(domain string) *dns.Msg {
+func msgQueryServices(zone string) *dns.Msg {
+	msg := new(dns.Msg)
+	msg.SetQuestion(fqn(queryServices)+fqn(zone), dns.TypePTR)
+	msg.RecursionDesired = false
+	return msg
+}
+
+func msgQueryLookup(srv, zone string) *dns.Msg {
 	msg := new(dns.Msg)
-	msg.SetQuestion(fqn(queryServices)+fqn(domain), dns.TypePTR)
+	msg.SetQuestion(fqn(srv)+fqn(zone), dns.TypePTR)
 	msg.RecursionDesired = false
 	return msg
 }

pkg/mdns/discovery_test.go

@@ -1,6 +1,7 @@
 package mdns_test
 
 import (
+	"sync"
 	"testing"
 	"time"
 
@@ -28,11 +29,38 @@ func Test_Discovery_001(t *testing.T) {
 		defer cancel()
 
 		// Cancel after one second
-		if err := app.Discovery.EnumerateServices(ctx); err != nil {
-			t.Error("EnumerateServices: ", err)
+		if services, err := app.Discovery.EnumerateServices(ctx); err != nil {
+			t.Error("EnumerateServices:", err)
 		} else {
-			t.Log("EnumerateServices: OK")
+			t.Log("EnumerateServices:", services)
 		}
+	})
+}
 
+func Test_Discovery_002(t *testing.T) {
+	tool.Test(t, nil, new(DiscoveryApp), func(app *DiscoveryApp) {
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+
+		// Cancel after one second
+		if services, err := app.Discovery.EnumerateServices(ctx); err != nil {
+			t.Error("EnumerateServices:", err)
+		} else {
+			var wg sync.WaitGroup
+			for _, service := range services {
+				wg.Add(1)
+				go func(service string) {
+					defer wg.Done()
+					ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+					defer cancel()
+					if r, err := app.Discovery.Lookup(ctx, service); err != nil {
+						t.Error(err)
+					} else {
+						t.Log("Lookup:", service, r)
+					}
+				}(service)
+			}
+			wg.Wait()
+		}
 	})
 }

pkg/mdns/event.go

@@ -10,39 +10,33 @@ import (
 ///////////////////////////////////////////////////////////////////////////////
 // TYPES
 
-type dnsevent struct {
-	msg *dns.Msg
+type msgevent struct {
+	*dns.Msg
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // NEW
 
-func NewDNSEvent(msg *dns.Msg) gopi.Event {
-	this := new(dnsevent)
-	this.msg = msg
-	return this
+func NewMsgEvent(msg *dns.Msg) gopi.Event {
+	return &msgevent{msg}
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // PUBLIC PROPERTIES
 
-func (this *dnsevent) Name() string {
-	if this.msg == nil {
+func (this *msgevent) Name() string {
+	if this.Msg == nil {
 		return ""
 	} else {
-		return this.msg.MsgHdr.String()
+		return this.Msg.MsgHdr.String()
 	}
 }
 
-func (this *dnsevent) Msg() *dns.Msg {
-	return this.msg
-}
-
 ///////////////////////////////////////////////////////////////////////////////
 // STRINGIFY
 
-func (this *dnsevent) String() string {
-	str := "<dns.event"
+func (this *msgevent) String() string {
+	str := "<dns.msg"
 	str += fmt.Sprintf(" name=%q", this.Name())
 	return str + ">"
 }

pkg/mdns/init.go

@@ -1,5 +1,13 @@
 package mdns
 
-func init() {
+import (
+	"reflect"
+
+	gopi "github.com/djthorpe/gopi/v3"
+	graph "github.com/djthorpe/gopi/v3/pkg/graph"
+)
 
+func init() {
+	// Register *mdns.Discovery -> gopi.ServiceDiscovery
+	graph.RegisterUnit(reflect.TypeOf(&Discovery{}), reflect.TypeOf((*gopi.ServiceDiscovery)(nil)))
 }

pkg/mdns/listener.go

@@ -224,7 +224,7 @@ func (this *Listener) run4(ctx context.Context, conn *ipv4.PacketConn) {
 				continue
 			} else if msg, err := parseDnsPacket(buf[:n], cm.IfIndex, from); err != nil {
 				this.Print("DNS Error:", err)
-			} else if err := this.Publisher.Emit(NewDNSEvent(msg), true); err != nil {
+			} else if err := this.Publisher.Emit(NewMsgEvent(msg), true); err != nil {
 				this.Print("Emit Error:", err)
 			}
 		}
@@ -246,7 +246,7 @@ func (this *Listener) run6(ctx context.Context, conn *ipv6.PacketConn) {
 				continue
 			} else if msg, err := parseDnsPacket(buf[:n], cm.IfIndex, from); err != nil {
 				this.Print("DNS Error:", err)
-			} else if err := this.Publisher.Emit(NewDNSEvent(msg), true); err != nil {
+			} else if err := this.Publisher.Emit(NewMsgEvent(msg), true); err != nil {
 				this.Print("Emit Error:", err)
 			}
 		}