probe/discovery/discoverer.go

package discovery

import (
	"log"
	"sync"
	"time"

	omd "git.loafle.net/overflow/model/discovery"
	omu "git.loafle.net/overflow/model/util"
	"git.loafle.net/overflow_scanner/probe/discovery/protocol/mdns"
	"git.loafle.net/overflow_scanner/probe/discovery/protocol/snmp"
	"git.loafle.net/overflow_scanner/probe/discovery/protocol/upnp"
	"git.loafle.net/overflow_scanner/probe/discovery/session"
	"git.loafle.net/overflow_scanner/probe/discovery/target/host"
	"git.loafle.net/overflow_scanner/probe/discovery/target/port"
	"git.loafle.net/overflow_scanner/probe/discovery/target/service"
	"git.loafle.net/overflow_scanner/probe/discovery/types"
)

type Discoverer interface {
	DiscoverHost(requesterID string, zone *omd.Zone, dh *omd.DiscoverHost)
	DiscoverPort(requesterID string, host *omd.Host, dp *omd.DiscoverPort)
	DiscoverService(requesterID string, port *omd.Port, ds *omd.DiscoverService)
	Message() <-chan types.DiscoveryMessage

	Shutdown()
}

func Instance() Discoverer {
	_once.Do(func() {
		i := &ofDiscoverer{}
		i.start()

		_instance = i
	})
	return _instance
}

var _instance Discoverer
var _once sync.Once

type ofDiscoverer struct {
	stopChan chan struct{}
	stopWg   sync.WaitGroup

	requestQueue chan types.DiscoveryRequest
	messageChan  chan types.DiscoveryMessage
}

func (d *ofDiscoverer) DiscoverHost(requesterID string, zone *omd.Zone, dh *omd.DiscoverHost) {
	d.enqueue(retainDiscoveryRequest(requesterID, types.DiscoveryRequestTypeHost, zone, dh))
}

func (d *ofDiscoverer) DiscoverPort(requesterID string, host *omd.Host, dp *omd.DiscoverPort) {
	d.enqueue(retainDiscoveryRequest(requesterID, types.DiscoveryRequestTypePort, host, dp))
}

func (d *ofDiscoverer) DiscoverService(requesterID string, port *omd.Port, ds *omd.DiscoverService) {
	d.enqueue(retainDiscoveryRequest(requesterID, types.DiscoveryRequestTypeService, port, ds))
}

func (d *ofDiscoverer) Message() <-chan types.DiscoveryMessage {
	return d.messageChan
}

func (d *ofDiscoverer) SendMessage(discoveryRequest types.DiscoveryRequest, messageType types.DiscoveryMessageType, data interface{}, err error) {
	d.messageChan <- retainDiscoveryMessage(discoveryRequest, messageType, data, err)
}

func (d *ofDiscoverer) Shutdown() {
	if d.stopChan == nil {
		return
	}

	close(d.requestQueue)
	close(d.messageChan)

	close(d.stopChan)
	d.stopWg.Wait()

	d.stopChan = nil
}

func (d *ofDiscoverer) start() {
	d.stopChan = make(chan struct{})

	d.requestQueue = make(chan types.DiscoveryRequest, 10)
	d.messageChan = make(chan types.DiscoveryMessage, 256)

	d.stopWg.Add(1)
	go d.handleRequest()
}

func (d *ofDiscoverer) enqueue(req *ofDiscoveryRequest) {
	select {
	case d.requestQueue <- req:
		d.SendMessage(req, types.DiscoveryMessageTypeQueueing, omu.Now(), nil)
		go func() {
			select {
			case <-req.dequeue:
			case <-time.After(10 * time.Second):
				req.timeout = true
				d.SendMessage(req, types.DiscoveryMessageTypeTimeout, omu.Now(), nil)
				req.release()
			}
		}()
	default:
		d.SendMessage(req, types.DiscoveryMessageTypeFailedQueueing, omu.Now(), nil)
		req.release()
	}
}

func (d *ofDiscoverer) handleRequest() {
	defer func() {
		d.stopWg.Done()
	}()

LOOP:
	for {
		select {
		case req, ok := <-d.requestQueue:
			if !ok {
				return
			}
			if req.(*ofDiscoveryRequest).timeout {
				continue LOOP
			}
			req.(*ofDiscoveryRequest).dequeue <- true

			d.SendMessage(req, types.DiscoveryMessageTypeStart, omu.Now(), nil)
			d.discover(req)
			d.SendMessage(req, types.DiscoveryMessageTypeStop, omu.Now(), nil)

			req.(*ofDiscoveryRequest).release()
		case <-d.stopChan:
			return
		}
	}
}

func (d *ofDiscoverer) discover(req types.DiscoveryRequest) {
	if types.DiscoveryRequestTypeNone == req.RequestType() {
		return
	}

	s := session.RetainDiscoverySession()
	defer func() {
		session.ReleaseDiscoverySession(s)
	}()

	if err := s.InitWithRequest(req); nil != err {
		d.SendMessage(req, types.DiscoveryMessageTypeError, nil, err)
		return
	}

	d.complexDiscover(s)
	d.hierarchyDiscover(s)
}

func (d *ofDiscoverer) complexDiscover(s session.DiscoverySession) {
	var wg sync.WaitGroup

	discoveredChan := make(chan interface{})
	s.SetDiscoveryDelegator(discoveredChan)
	defer func() {
		s.SetDiscoveryDelegator(nil)
		close(discoveredChan)
	}()
	go func() {
		for {
			select {
			case target, ok := <-discoveredChan:
				if !ok {
					return
				}
				switch target.(type) {
				case *omd.Host:
					d.SendMessage(s.DiscoveryRequest(), types.DiscoveryMessageTypeHost, target, nil)
				case *omd.Port:
					d.SendMessage(s.DiscoveryRequest(), types.DiscoveryMessageTypePort, target, nil)
				case *omd.Service:
					d.SendMessage(s.DiscoveryRequest(), types.DiscoveryMessageTypeService, target, nil)
				default:
				}
			}
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		err := upnp.Scan(s)
		if nil != err {
			log.Printf("UPnP %v", err)
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		err := mdns.Scan(s)
		if nil != err {
			log.Printf("mDNS %v", err)
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		err := snmp.Scan(s)
		if nil != err {
			log.Printf("SNMP %v", err)
		}
	}()

	wg.Wait()
}

func (d *ofDiscoverer) hierarchyDiscover(s session.DiscoverySession) {
	var wg sync.WaitGroup

	discoveredChan := make(chan interface{})
	s.SetDiscoveryDelegator(discoveredChan)
	defer func() {
		s.SetDiscoveryDelegator(nil)
		close(discoveredChan)
	}()
	go func() {
		for {
			select {
			case target, ok := <-discoveredChan:
				if !ok {
					return
				}
				switch target.(type) {
				case *omd.Host:
					d.SendMessage(s.DiscoveryRequest(), types.DiscoveryMessageTypeHost, target, nil)
					if nil != s.DiscoverPort() {
						wg.Add(1)
						go func() {
							defer wg.Done()
							port.Scan(s, target.(*omd.Host))
						}()
					}
				case *omd.Port:
					d.SendMessage(s.DiscoveryRequest(), types.DiscoveryMessageTypePort, target, nil)
					if nil != s.DiscoverService() {
						wg.Add(1)
						go func() {
							defer wg.Done()
							service.Scan(s, target.(*omd.Port))
						}()
					}
				case *omd.Service:
					d.SendMessage(s.DiscoveryRequest(), types.DiscoveryMessageTypeService, target, nil)
				default:
				}
			}
		}
	}()

	// post complexDiscover
	if nil != s.DiscoverPort() {
		discoveredHosts := s.DiscoveredAllHosts(true)
		for _, h := range discoveredHosts {
			wg.Add(1)
			go func(h *omd.Host) {
				defer wg.Done()
				port.Scan(s, h)
			}(h)
		}
	}
	wg.Wait()

	if nil != s.DiscoverService() {
		discoveredPorts := s.DiscoveredAllPorts()
		for _, hostPorts := range discoveredPorts {
			for _, ports := range hostPorts {
				for _, p := range ports {
					wg.Add(1)
					go func(p *omd.Port) {
						defer wg.Done()
						service.Scan(s, p)
					}(p)
				}
			}
		}
	}
	wg.Wait()

	// exec hierarchy Discovery
	if nil != s.DiscoverHost() {
		wg.Add(1)
		go func() {
			defer wg.Done()
			host.Scan(s)
		}()
	} else if nil != s.DiscoverPort() {
		if nil != s.Host() {
			wg.Add(1)
			go func() {
				defer wg.Done()
				port.Scan(s, s.Host())
			}()
		}
	} else if nil != s.DiscoverService() {
		if nil != s.Port() {
			wg.Add(1)
			go func() {
				defer wg.Done()
				service.Scan(s, s.Port())
			}()
		}
	}

	wg.Wait()
}