probe/discovery/protocol/mdns/mdns.go

package mdns

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	omd "git.loafle.net/overflow/model/discovery"
	omm "git.loafle.net/overflow/model/meta"
	omu "git.loafle.net/overflow/model/util"
	"git.loafle.net/overflow_scanner/probe/discovery/session"
	"github.com/grandcat/zeroconf"
)

func Scan(discoverySession session.DiscoverySession) error {
	serviceEntries, err := browse("_services._dns-sd._udp", "local")
	if nil != err {
		log.Print("Cannot find service ", err)
		return nil
	}

	metaIPTypeEnum := omm.ToMetaIPTypeEnum(discoverySession.Zone().MetaIPType)

SERVICE_LOOP:
	for _, serviceEntry := range serviceEntries {
		name := removeDomainName(serviceEntry.Instance, serviceEntry.Domain)
		entries, _err := browse(name, serviceEntry.Domain)
		if nil != _err {
			log.Print("Cannot find entry ", _err)
			continue SERVICE_LOOP
		}

	ENTRY_LOOP:
		for _, entry := range entries {
			log.Print("serviceEntry ", entry)

			name := entry.Instance   // HP\ LaserJet\ P1505n
			service := entry.Service // _pdl-datastream._tcp
			port := entry.Port       // 9100
			meta := toMeta(entry.Text)
			hostName := removeDomainName(entry.HostName, entry.Domain) // NPIFACA9B
			serviceName, portType, cryptoType := parseService(service)

			var metaPortType *omm.MetaPortType
			switch portType {
			case "tcp":
				metaPortType = omm.ToMetaPortType(omm.MetaPortTypeEnumTCP)
			case "udp":
				metaPortType = omm.ToMetaPortType(omm.MetaPortTypeEnumUDP)
			}

			metaCryptoType := omm.ToMetaCryptoType(omm.MetaCryptoTypeEnumUNKNOWN)
			switch cryptoType {
			case "NONE":
				metaCryptoType = omm.ToMetaCryptoType(omm.MetaCryptoTypeEnumNONE)
			case "TLS":
				metaCryptoType = omm.ToMetaCryptoType(omm.MetaCryptoTypeEnumTLS)
			}

			switch metaIPTypeEnum {
			case omm.MetaIPTypeEnumV4:
				for _, ipv4 := range entry.AddrIPv4 {
					h := discoverySession.AddHost(&omd.Host{
						MetaIPType:     omm.ToMetaIPType(metaIPTypeEnum),
						Name:           hostName,
						Address:        ipv4.String(),
						Meta:           meta,
						Zone:           discoverySession.Zone(),
						DiscoveredDate: omu.NowPtr(),
					})

					if 1 > port {
						continue ENTRY_LOOP
					}

					p := discoverySession.AddPort(&omd.Port{
						MetaPortType: metaPortType,
						PortNumber:   json.Number(strconv.Itoa(port)),
						Meta:         meta,
						Host:         h,
					})

					discoverySession.AddService(&omd.Service{
						MetaCryptoType: metaCryptoType,
						Key:            serviceName,
						Name:           name,
						Port:           p,
					})
				}

			case omm.MetaIPTypeEnumV6:
				for _, ipv6 := range entry.AddrIPv6 {
					h := discoverySession.AddHost(&omd.Host{
						MetaIPType:     omm.ToMetaIPType(metaIPTypeEnum),
						Name:           hostName,
						Address:        ipv6.String(),
						Meta:           meta,
						Zone:           discoverySession.Zone(),
						DiscoveredDate: omu.NowPtr(),
					})

					if 1 > port {
						continue ENTRY_LOOP
					}

					p := discoverySession.AddPort(&omd.Port{
						MetaPortType: metaPortType,
						PortNumber:   json.Number(strconv.Itoa(port)),
						Meta:         meta,
						Host:         h,
					})

					discoverySession.AddService(&omd.Service{
						MetaCryptoType: metaCryptoType,
						Key:            serviceName,
						Name:           name,
						Port:           p,
					})
				}
			}

		}
	}

	return nil
}

func removeDomainName(instance string, domain string) (service string) {
	_instance := strings.TrimRight(instance, ".")
	return strings.TrimRight(_instance, fmt.Sprintf(".%s", domain))
}

func toMeta(text []string) map[string]string {
	if nil == text || 0 == len(text) {
		return nil
	}

	meta := make(map[string]string)
	for _, v := range text {
		ss := strings.SplitN(v, "=", 2)
		if 2 != len(ss) {
			continue
		}
		if nil != ss {
			meta[ss[0]] = ss[1]
		}
	}
	return meta
}

func parseService(service string) (name string, portType string, cryptoType string) {
	ss := strings.SplitN(service, ".", 2)
	if nil == ss {
		return "UNKNOWN", "UNKNOWN", "UNKNOWN"
	}

	_name := strings.TrimLeft(ss[0], "_")
	_portType := strings.TrimLeft(ss[1], "_")

	name = _name
	portType = _portType
	cryptoType = "NONE"
	return
}

func browse(service string, domain string) ([]*zeroconf.ServiceEntry, error) {
	resolver, err := zeroconf.NewResolver(nil)
	if err != nil {
		return nil, fmt.Errorf("Failed to initialize resolver %v", err)
	}

	entryChan := make(chan *zeroconf.ServiceEntry)
	timerStopped := make(chan struct{})

	serviceEntries := make([]*zeroconf.ServiceEntry, 0)

	go func(_entryChan <-chan *zeroconf.ServiceEntry) {
		var delay atomic.Value
		delay.Store(false)
		ticker := time.NewTicker(time.Second * 1)

		for {
			select {
			case entry, ok := <-_entryChan:
				if !ok {
					return
				}
				delay.Store(true)
				serviceEntries = append(serviceEntries, entry)
			case <-ticker.C:
				if false == delay.Load().(bool) {
					ticker.Stop()
					timerStopped <- struct{}{}
					return
				}
				delay.Store(false)
			}
		}
	}(entryChan)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second*time.Duration(30))
	defer cancel()
	err = resolver.Browse(ctx, service, domain, entryChan)
	if err != nil {
		return nil, fmt.Errorf("Failed to browse %v", err)
	}

	select {
	case <-ctx.Done():
	case <-timerStopped:
	}

	return serviceEntries, nil
}