golib/v2/features/mdns/mdns.go

package mdns

import (
	"context"
	"errors"
	"net"
	"strings"
	"time"

	"golang.org/x/net/ipv4"
	"golib/v2/pkg/mdns"
)

const (
	DefaultTimout = 3 * time.Second
)

type Server struct {
	Name    []string
	Address *net.UDPAddr

	server *mdns.Conn
}

func (s *Server) Close() error {
	return s.server.Close()
}

func (s *Server) ListenAndServe() error {
	conn, err := net.ListenUDP(mdns.NetType, s.Address)
	if err != nil {
		return err
	}
	cfg := &mdns.Config{
		LocalNames: s.Name,
	}
	s.server, err = mdns.Server(ipv4.NewPacketConn(conn), cfg)
	if err != nil {
		return err
	}
	select {}
}

func ListenAndServe(name string) error {
	return ListenAndServeNames([]string{name})
}

func ListenAndServeNames(name []string) error {
	server := &Server{
		Name:    name,
		Address: mdns.Address,
	}
	return server.ListenAndServe()
}

type Handler func(name string, addr net.IP)

type Client struct {
	Name    []string
	Address *net.UDPAddr
	Timout  time.Duration
	Handle  Handler

	server *mdns.Conn
}

func (c *Client) initServer() error {
	if c.server != nil {
		return nil
	}
	conn, err := net.ListenUDP(mdns.NetType, c.Address)
	if err != nil {
		return err
	}
	if c.Timout <= 0 {
		c.Timout = DefaultTimout
	}
	c.server, err = mdns.Server(ipv4.NewPacketConn(conn), &mdns.Config{QueryInterval: c.Timout})
	if err != nil {
		return err
	}
	return nil
}

func (c *Client) Lookup() ([]net.IP, error) {
	if err := c.initServer(); err != nil {
		return nil, err
	}
	ips := make([]net.IP, 0, len(c.Name))
	for _, name := range c.Name {
		ctx, cancel := context.WithTimeout(context.Background(), c.Timout)
		var src net.Addr
		_, src, err := c.server.Query(ctx, name)
		cancel()
		if err != nil {
			continue // 忽略错误
		}
		ips = append(ips, src.(*net.IPAddr).IP)
	}
	return ips, nil
}

func (c *Client) LookupWithName() (map[string]net.IP, error) {
	if err := c.initServer(); err != nil {
		return nil, err
	}
	ips := make(map[string]net.IP)
	for _, name := range c.Name {
		ctx, cancel := context.WithTimeout(context.Background(), c.Timout)
		answer, src, err := c.server.Query(ctx, name)
		cancel()
		if err != nil {
			return nil, err
		}
		ips[mdns.UnFqdn(answer.Name.String())] = src.(*net.IPAddr).IP
	}
	return ips, nil
}

func (c *Client) ListenAndServe() error {
	if err := c.initServer(); err != nil {
		return err
	}
	timer := time.NewTimer(c.Timout)
	for {
		select {
		case <-timer.C:
			for _, name := range c.Name {
				go func(name string) {
					ctx, cancel := context.WithTimeout(context.Background(), c.Timout)
					answer, src, err := c.server.Query(ctx, name)
					cancel()
					if err != nil {
						return
					}
					if c.Handle != nil {
						c.Handle(mdns.UnFqdn(answer.Name.String()), src.(*net.IPAddr).IP)
					}
				}(name)
			}
			timer.Reset(c.Timout)
		}
	}
}

func Lookup(name string) (net.IP, error) {
	ips, err := Lookups([]string{name})
	if err != nil {
		return nil, err
	}
	if len(ips) > 0 {
		return ips[0], nil
	}
	return nil, errors.New("not found")
}

func Lookups(name []string) ([]net.IP, error) {
	client := &Client{
		Name:    name,
		Address: mdns.Address,
	}
	return client.Lookup()
}

func Fqdn(name string) string {
	return UnFqdn(name) + ".local"
}

func UnFqdn(name string) string {
	return strings.TrimSuffix(name, ".local")
}