Allow configuring an alternative logger (#116)

* allow configuring an alternative logger (#3) * s/Log/Logger/ per PR comment
hashicorp · May 3, 2024 · 901c2c5 · 901c2c5
1 parent 9f136d7
commit 901c2c5
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 13 deletions.
diff --git a/client.go b/client.go
@@ -49,6 +49,7 @@ type QueryParam struct {
 	WantUnicastResponse bool                 // Unicast response desired, as per 5.4 in RFC
 	DisableIPv4         bool                 // Whether to disable usage of IPv4 for MDNS operations. Does not affect discovered addresses.
 	DisableIPv6         bool                 // Whether to disable usage of IPv6 for MDNS operations. Does not affect discovered addresses.
+	Logger              *log.Logger          // Optionally provide a *log.Logger to better manage log output.
 }
 
 // DefaultParams is used to return a default set of QueryParam's
@@ -78,8 +79,11 @@ func Query(params *QueryParam) error {
 // either read or buffer. QueryContext will attempt to stop the query
 // on cancellation.
 func QueryContext(ctx context.Context, params *QueryParam) error {
+	if params.Logger == nil {
+		params.Logger = log.Default()
+	}
 	// Create a new client
-	client, err := newClient(!params.DisableIPv4, !params.DisableIPv6)
+	client, err := newClient(!params.DisableIPv4, !params.DisableIPv6, params.Logger)
 	if err != nil {
 		return err
 	}
@@ -134,11 +138,13 @@ type client struct {
 
 	closed   int32
 	closedCh chan struct{} // TODO(reddaly): This doesn't appear to be used.
+
+	log *log.Logger
 }
 
 // NewClient creates a new mdns Client that can be used to query
 // for records
-func newClient(v4 bool, v6 bool) (*client, error) {
+func newClient(v4 bool, v6 bool, logger *log.Logger) (*client, error) {
 	if !v4 && !v6 {
 		return nil, fmt.Errorf("Must enable at least one of IPv4 and IPv6 querying")
 	}
@@ -154,14 +160,14 @@ func newClient(v4 bool, v6 bool) (*client, error) {
 	if v4 {
 		uconn4, err = net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
 		if err != nil {
-			log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
+			logger.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
 		}
 	}
 
 	if v6 {
 		uconn6, err = net.ListenUDP("udp6", &net.UDPAddr{IP: net.IPv6zero, Port: 0})
 		if err != nil {
-			log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
+			logger.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
 		}
 	}
 
@@ -172,13 +178,13 @@ func newClient(v4 bool, v6 bool) (*client, error) {
 	if v4 {
 		mconn4, err = net.ListenMulticastUDP("udp4", nil, ipv4Addr)
 		if err != nil {
-			log.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
+			logger.Printf("[ERR] mdns: Failed to bind to udp4 port: %v", err)
 		}
 	}
 	if v6 {
 		mconn6, err = net.ListenMulticastUDP("udp6", nil, ipv6Addr)
 		if err != nil {
-			log.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
+			logger.Printf("[ERR] mdns: Failed to bind to udp6 port: %v", err)
 		}
 	}
 
@@ -194,6 +200,7 @@ func newClient(v4 bool, v6 bool) (*client, error) {
 		ipv4UnicastConn:   uconn4,
 		ipv6UnicastConn:   uconn6,
 		closedCh:          make(chan struct{}),
+		log:               logger,
 	}
 	return c, nil
 }
@@ -205,7 +212,7 @@ func (c *client) Close() error {
 		return nil
 	}
 
-	log.Printf("[INFO] mdns: Closing client %v", *c)
+	c.log.Printf("[INFO] mdns: Closing client %v", *c)
 	close(c.closedCh)
 
 	if c.ipv4UnicastConn != nil {
@@ -365,7 +372,7 @@ func (c *client) query(params *QueryParam) error {
 				m.SetQuestion(inp.Name, dns.TypePTR)
 				m.RecursionDesired = false
 				if err := c.sendQuery(m); err != nil {
-					log.Printf("[ERR] mdns: Failed to query instance %s: %v", inp.Name, err)
+					c.log.Printf("[ERR] mdns: Failed to query instance %s: %v", inp.Name, err)
 				}
 			}
 		case <-finish:
@@ -409,12 +416,12 @@ func (c *client) recv(l *net.UDPConn, msgCh chan *msgAddr) {
 		}
 
 		if err != nil {
-			log.Printf("[ERR] mdns: Failed to read packet: %v", err)
+			c.log.Printf("[ERR] mdns: Failed to read packet: %v", err)
 			continue
 		}
 		msg := new(dns.Msg)
 		if err := msg.Unpack(buf[:n]); err != nil {
-			log.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
+			c.log.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
 			continue
 		}
 		select {

diff --git a/server.go b/server.go
@@ -44,6 +44,9 @@ type Config struct {
 	// LogEmptyResponses indicates the server should print an informative message
 	// when there is an mDNS query for which the server has no response.
 	LogEmptyResponses bool
+
+	// Logger can optionally be set to use an alternative logger instead of the default.
+	Logger *log.Logger
 }
 
 // mDNS server is used to listen for mDNS queries and respond if we
@@ -69,6 +72,10 @@ func NewServer(config *Config) (*Server, error) {
 		return nil, fmt.Errorf("no multicast listeners could be started")
 	}
 
+	if config.Logger == nil {
+		config.Logger = log.Default()
+	}
+
 	s := &Server{
 		config:     config,
 		ipv4List:   ipv4List,
@@ -118,7 +125,7 @@ func (s *Server) recv(c *net.UDPConn) {
 			continue
 		}
 		if err := s.parsePacket(buf[:n], from); err != nil {
-			log.Printf("[ERR] mdns: Failed to handle query: %v", err)
+			s.config.Logger.Printf("[ERR] mdns: Failed to handle query: %v", err)
 		}
 	}
 }
@@ -127,7 +134,7 @@ func (s *Server) recv(c *net.UDPConn) {
 func (s *Server) parsePacket(packet []byte, from net.Addr) error {
 	var msg dns.Msg
 	if err := msg.Unpack(packet); err != nil {
-		log.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
+		s.config.Logger.Printf("[ERR] mdns: Failed to unpack packet: %v", err)
 		return err
 	}
 	return s.handleQuery(&msg, from)
@@ -226,7 +233,7 @@ func (s *Server) handleQuery(query *dns.Msg, from net.Addr) error {
 		for i, q := range query.Question {
 			questions[i] = q.Name
 		}
-		log.Printf("no responses for query with questions: %s", strings.Join(questions, ", "))
+		s.config.Logger.Printf("no responses for query with questions: %s", strings.Join(questions, ", "))
 	}
 
 	if mresp := resp(false); mresp != nil {