probe/discovery/discoverer.go

package discovery

import (
	"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/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() {
		_instance = newDiscoverer()
	})
	return _instance
}

var _instance Discoverer
var _once sync.Once

func newDiscoverer() Discoverer {
	i := &ofDiscoverer{}
	i.start()

	return i
}

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()
		upnp.Scan(s)
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		mdns.Scan(s)
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		snmp.Scan(s)
	}()

	wg.Wait()
}

func (d *ofDiscoverer) hierarchyDiscover(s session.DiscoverySession) {

}