package registry import ( "bytes" "compress/zlib" "context" "encoding/hex" "encoding/json" "fmt" "io/ioutil" "net" "strconv" "strings" "sync" "time" "github.com/google/uuid" "github.com/micro/go-micro/v2/logger" "github.com/micro/go-micro/v2/util/mdns" ) const ( // every service is written to the global domain so * domain queries work, e.g. // calling mdns.List(registry.ListDomain("*")) will list the services across all // domains globalDomain = "global" ) type mdnsTxt struct { Service string Version string Endpoints []*Endpoint Metadata map[string]string } type mdnsEntry struct { id string node *mdns.Server } // services are a key/value map, with the service name as a key and the value being a // slice of mdns entries, representing the nodes with a single _services entry to be // used for listing type services map[string][]*mdnsEntry // mdsRegistry is a multicast dns registry type mdnsRegistry struct { opts Options // the top level domains, these can be overriden using options defaultDomain string globalDomain string sync.Mutex domains map[string]services mtx sync.RWMutex // watchers watchers map[string]*mdnsWatcher // listener listener chan *mdns.ServiceEntry } type mdnsWatcher struct { id string wo WatchOptions ch chan *mdns.ServiceEntry exit chan struct{} // the mdns domain domain string // the registry registry *mdnsRegistry } func encode(txt *mdnsTxt) ([]string, error) { b, err := json.Marshal(txt) if err != nil { return nil, err } var buf bytes.Buffer defer buf.Reset() w := zlib.NewWriter(&buf) if _, err := w.Write(b); err != nil { w.Close() return nil, err } w.Close() encoded := hex.EncodeToString(buf.Bytes()) // individual txt limit if len(encoded) <= 255 { return []string{encoded}, nil } // split encoded string var record []string for len(encoded) > 255 { record = append(record, encoded[:255]) encoded = encoded[255:] } record = append(record, encoded) return record, nil } func decode(record []string) (*mdnsTxt, error) { encoded := strings.Join(record, "") hr, err := hex.DecodeString(encoded) if err != nil { return nil, err } br := bytes.NewReader(hr) zr, err := zlib.NewReader(br) if err != nil { return nil, err } defer zr.Close() rbuf, err := ioutil.ReadAll(zr) if err != nil { return nil, err } var txt *mdnsTxt if err := json.Unmarshal(rbuf, &txt); err != nil { return nil, err } return txt, nil } func newRegistry(opts ...Option) Registry { options := Options{ Context: context.Background(), Timeout: time.Millisecond * 100, } for _, o := range opts { o(&options) } // set the domain defaultDomain := DefaultDomain if d, ok := options.Context.Value("mdns.domain").(string); ok { defaultDomain = d } return &mdnsRegistry{ defaultDomain: defaultDomain, globalDomain: globalDomain, opts: options, domains: make(map[string]services), watchers: make(map[string]*mdnsWatcher), } } func (m *mdnsRegistry) Init(opts ...Option) error { for _, o := range opts { o(&m.opts) } return nil } func (m *mdnsRegistry) Options() Options { return m.opts } // createServiceMDNSEntry will create a new wildcard mdns entry for the service in the // given domain. This wildcard mdns entry is used when listing services. func createServiceMDNSEntry(name, domain string) (*mdnsEntry, error) { ip := net.ParseIP("0.0.0.0") s, err := mdns.NewMDNSService(name, "_services", domain+".", "", 9999, []net.IP{ip}, nil) if err != nil { return nil, err } srv, err := mdns.NewServer(&mdns.Config{Zone: &mdns.DNSSDService{MDNSService: s}, LocalhostChecking: true}) if err != nil { return nil, err } return &mdnsEntry{id: "*", node: srv}, nil } func (m *mdnsRegistry) getMdnsEntries(domain, serviceName string) ([]*mdnsEntry, error) { entries, ok := m.domains[domain][serviceName] if ok { return entries, nil } // create the wildcard entry used for list queries in this domain entry, err := createServiceMDNSEntry(serviceName, domain) if err != nil { return nil, err } return []*mdnsEntry{entry}, nil } func registerService(service *Service, entries []*mdnsEntry, options RegisterOptions) ([]*mdnsEntry, error) { var lastError error for _, node := range service.Nodes { var seen bool for _, entry := range entries { if node.Id == entry.id { seen = true break } } // this node has already been registered, continue if seen { continue } txt, err := encode(&mdnsTxt{ Service: service.Name, Version: service.Version, Endpoints: service.Endpoints, Metadata: node.Metadata, }) if err != nil { lastError = err continue } host, pt, err := net.SplitHostPort(node.Address) if err != nil { lastError = err continue } port, _ := strconv.Atoi(pt) if logger.V(logger.DebugLevel, logger.DefaultLogger) { logger.Debugf("[mdns] registry create new service with ip: %s for: %s", net.ParseIP(host).String(), host) } // we got here, new node s, err := mdns.NewMDNSService( node.Id, service.Name, options.Domain+".", "", port, []net.IP{net.ParseIP(host)}, txt, ) if err != nil { lastError = err continue } srv, err := mdns.NewServer(&mdns.Config{Zone: s, LocalhostChecking: true}) if err != nil { lastError = err continue } entries = append(entries, &mdnsEntry{id: node.Id, node: srv}) } return entries, lastError } func createGlobalDomainService(service *Service, options RegisterOptions) *Service { srv := *service srv.Nodes = nil for _, n := range service.Nodes { node := n // set the original domain in node metadata if node.Metadata == nil { node.Metadata = map[string]string{"domain": options.Domain} } else { node.Metadata["domain"] = options.Domain } srv.Nodes = append(srv.Nodes, node) } return &srv } func (m *mdnsRegistry) Register(service *Service, opts ...RegisterOption) error { m.Lock() // parse the options var options RegisterOptions for _, o := range opts { o(&options) } if len(options.Domain) == 0 { options.Domain = m.defaultDomain } // create the domain in the memory store if it doesn't yet exist if _, ok := m.domains[options.Domain]; !ok { m.domains[options.Domain] = make(services) } entries, err := m.getMdnsEntries(options.Domain, service.Name) if err != nil { m.Unlock() return err } entries, gerr := registerService(service, entries, options) // save the mdns entry m.domains[options.Domain][service.Name] = entries m.Unlock() // register in the global Domain so it can be queried as one if options.Domain != m.globalDomain { srv := createGlobalDomainService(service, options) if err := m.Register(srv, append(opts, RegisterDomain(m.globalDomain))...); err != nil { gerr = err } } return gerr } func (m *mdnsRegistry) Deregister(service *Service, opts ...DeregisterOption) error { // parse the options var options DeregisterOptions for _, o := range opts { o(&options) } if len(options.Domain) == 0 { options.Domain = m.defaultDomain } // register in the global Domain var err error if options.Domain != m.globalDomain { defer func() { err = m.Deregister(service, append(opts, DeregisterDomain(m.globalDomain))...) }() } // we want to unlock before we call deregister on the global domain, so it's important this unlock // is applied after the defer m.Deregister is called above m.Lock() defer m.Unlock() // the service wasn't registered, we can safely exist if _, ok := m.domains[options.Domain]; !ok { return err } // loop existing entries, check if any match, shutdown those that do var newEntries []*mdnsEntry for _, entry := range m.domains[options.Domain][service.Name] { var remove bool for _, node := range service.Nodes { if node.Id == entry.id { entry.node.Shutdown() remove = true break } } // keep it? if !remove { newEntries = append(newEntries, entry) } } // we have no new entries, we can exit if len(newEntries) == 0 { return nil } // we have more than one entry remaining, we can exit if len(newEntries) > 1 { m.domains[options.Domain][service.Name] = newEntries return err } // our remaining entry is not a wildcard, we can exit if len(newEntries) == 1 && newEntries[0].id != "*" { m.domains[options.Domain][service.Name] = newEntries return err } // last entry is the wildcard for list queries. Remove it. newEntries[0].node.Shutdown() delete(m.domains[options.Domain], service.Name) // check to see if we can delete the domain entry if len(m.domains[options.Domain]) == 0 { delete(m.domains, options.Domain) } return err } func (m *mdnsRegistry) GetService(service string, opts ...GetOption) ([]*Service, error) { // parse the options var options GetOptions for _, o := range opts { o(&options) } if len(options.Domain) == 0 { options.Domain = m.defaultDomain } if options.Domain == WildcardDomain { options.Domain = m.globalDomain } serviceMap := make(map[string]*Service) entries := make(chan *mdns.ServiceEntry, 10) done := make(chan bool) p := mdns.DefaultParams(service) // set context with timeout var cancel context.CancelFunc p.Context, cancel = context.WithTimeout(context.Background(), m.opts.Timeout) defer cancel() // set entries channel p.Entries = entries // set the domain p.Domain = options.Domain go func() { for { select { case e := <-entries: // list record so skip if e.Name == "_services" { continue } if e.TTL == 0 { continue } txt, err := decode(e.InfoFields) if err != nil { continue } if txt.Service != service { continue } s, ok := serviceMap[txt.Version] if !ok { s = &Service{ Name: txt.Service, Version: txt.Version, Endpoints: txt.Endpoints, } } addr := "" // prefer ipv4 addrs if len(e.AddrV4) > 0 { addr = e.AddrV4.String() // else use ipv6 } else if len(e.AddrV6) > 0 { addr = "[" + e.AddrV6.String() + "]" } else { if logger.V(logger.InfoLevel, logger.DefaultLogger) { logger.Infof("[mdns]: invalid endpoint received: %v", e) } continue } s.Nodes = append(s.Nodes, &Node{ Id: strings.TrimSuffix(e.Name, "."+p.Service+"."+p.Domain+"."), Address: fmt.Sprintf("%s:%d", addr, e.Port), Metadata: txt.Metadata, }) serviceMap[txt.Version] = s case <-p.Context.Done(): close(done) return } } }() // execute the query if err := mdns.Query(p); err != nil { return nil, err } // wait for completion <-done // create list and return services := make([]*Service, 0, len(serviceMap)) for _, service := range serviceMap { services = append(services, service) } return services, nil } func (m *mdnsRegistry) ListServices(opts ...ListOption) ([]*Service, error) { // parse the options var options ListOptions for _, o := range opts { o(&options) } if len(options.Domain) == 0 { options.Domain = m.defaultDomain } if options.Domain == WildcardDomain { options.Domain = m.globalDomain } serviceMap := make(map[string]bool) entries := make(chan *mdns.ServiceEntry, 10) done := make(chan bool) p := mdns.DefaultParams("_services") // set context with timeout var cancel context.CancelFunc p.Context, cancel = context.WithTimeout(context.Background(), m.opts.Timeout) defer cancel() // set entries channel p.Entries = entries // set domain p.Domain = options.Domain var services []*Service go func() { for { select { case e := <-entries: 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 services = append(services, &Service{Name: name}) } case <-p.Context.Done(): close(done) return } } }() // execute query if err := mdns.Query(p); err != nil { return nil, err } // wait till done <-done return services, nil } func (m *mdnsRegistry) Watch(opts ...WatchOption) (Watcher, error) { var wo WatchOptions for _, o := range opts { o(&wo) } if len(wo.Domain) == 0 { wo.Domain = m.defaultDomain } if wo.Domain == WildcardDomain { wo.Domain = m.globalDomain } md := &mdnsWatcher{ id: uuid.New().String(), wo: wo, ch: make(chan *mdns.ServiceEntry, 32), exit: make(chan struct{}), domain: wo.Domain, registry: m, } m.mtx.Lock() defer m.mtx.Unlock() // save the watcher m.watchers[md.id] = md // check of the listener exists if m.listener != nil { return md, nil } // start the listener go func() { // go to infinity for { m.mtx.Lock() // just return if there are no watchers if len(m.watchers) == 0 { m.listener = nil m.mtx.Unlock() return } // check existing listener if m.listener != nil { m.mtx.Unlock() return } // reset the listener exit := make(chan struct{}) ch := make(chan *mdns.ServiceEntry, 32) m.listener = ch m.mtx.Unlock() // send messages to the watchers go func() { send := func(w *mdnsWatcher, e *mdns.ServiceEntry) { select { case w.ch <- e: default: } } for { select { case <-exit: return case e, ok := <-ch: if !ok { return } m.mtx.RLock() // send service entry to all watchers for _, w := range m.watchers { send(w, e) } m.mtx.RUnlock() } } }() // start listening, blocking call mdns.Listen(ch, exit) // mdns.Listen has unblocked // kill the saved listener m.mtx.Lock() m.listener = nil close(ch) m.mtx.Unlock() } }() return md, nil } func (m *mdnsRegistry) String() string { return "mdns" } func (m *mdnsWatcher) Next() (*Result, error) { for { select { case e := <-m.ch: txt, err := decode(e.InfoFields) if err != nil { continue } if len(txt.Service) == 0 || len(txt.Version) == 0 { continue } // Filter watch options // wo.Service: Only keep services we care about if len(m.wo.Service) > 0 && txt.Service != m.wo.Service { continue } var action string if e.TTL == 0 { action = "delete" } else { action = "create" } service := &Service{ Name: txt.Service, Version: txt.Version, Endpoints: txt.Endpoints, Metadata: txt.Metadata, } // skip anything without the domain we care about suffix := fmt.Sprintf(".%s.%s.", service.Name, m.domain) if !strings.HasSuffix(e.Name, suffix) { continue } var addr string if len(e.AddrV4) > 0 { addr = e.AddrV4.String() } else if len(e.AddrV6) > 0 { addr = "[" + e.AddrV6.String() + "]" } else { addr = e.Addr.String() } service.Nodes = append(service.Nodes, &Node{ Id: strings.TrimSuffix(e.Name, suffix), Address: fmt.Sprintf("%s:%d", addr, e.Port), Metadata: txt.Metadata, }) return &Result{ Action: action, Service: service, }, nil case <-m.exit: return nil, ErrWatcherStopped } } } func (m *mdnsWatcher) Stop() { select { case <-m.exit: return default: close(m.exit) // remove self from the registry m.registry.mtx.Lock() delete(m.registry.watchers, m.id) m.registry.mtx.Unlock() } } // NewRegistry returns a new default registry which is mdns func NewRegistry(opts ...Option) Registry { return newRegistry(opts...) }