overflow_discovery/discovery/discovery.go

package discovery

import (
	"fmt"
	"sync"

	"git.loafle.net/commons_go/logging"
	"git.loafle.net/commons_go/rpc"
	"git.loafle.net/commons_go/util/net/cidr"
	"git.loafle.net/overflow/overflow_discovery/api/module/discovery/model"
)

var RPCServlet rpc.Servlet

var discoverer *discovery

func DiscoveryInit() {
	discoverer = &discovery{}
	discoverer.start()
}

func DiscoveryDestroy() {
	discoverer.stop()
	discoverer = nil
}

func DiscoverZone(dz *model.DiscoveryZone) {
	discoverer.discoverZone(dz)
}

func DiscoverHost(zone *model.Zone, dh *model.DiscoveryHost) {
	discoverer.discoverHost(zone, dh)
}

func DiscoverPort(host *model.Host, dp *model.DiscoveryPort) {
	discoverer.discoverPort(host, dp)
}

func DiscoverService(port *model.Port, ds *model.DiscoveryService) {
	discoverer.discoverService(port, ds)
}

func Stop() {

}

type discovery struct {
	sendChan chan interface{}
	errChan  chan error

	stopChan chan struct{}
	stopWg   sync.WaitGroup
}

func (d *discovery) start() {
	if d.stopChan != nil {
		panic("Discovery: discovery is already running. Stop it before starting it again")
	}

	d.stopChan = make(chan struct{})
	d.sendChan = make(chan interface{})
	d.errChan = make(chan error)
	logging.Logger().Info(fmt.Sprintf("Discovery: discovery is started"))
}

func (d *discovery) stop() {
	if d.stopChan == nil {
		panic("Discovery: discovery must be started before stopping it")
	}
	close(d.stopChan)
	d.stopWg.Wait()
	d.stopChan = nil

	logging.Logger().Info(fmt.Sprintf("Discovery: discovery is stopped"))
}

func (d *discovery) discoverZone(dz *model.DiscoveryZone) {
	go taskScan(d,
		func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) {
			scanZone(dz, resultChan, errChan, doneChan, stopChan)
		},
		func(result interface{}) {
			z := result.(*model.Zone)
			logging.Logger().Info(fmt.Sprintf("zone: %v", z))
			d.sendResult("DiscoveryService.DiscoveredZone", z)
			if nil != dz.DiscoveryHost {
				cr, _ := cidr.NewCIDRRanger(z.Network)
				dh := &model.DiscoveryHost{
					FirstScanRange: cr.First().String(),
					LastScanRange:  cr.Last().String(),
					DiscoveryPort:  dz.DiscoveryHost.DiscoveryPort,
				}
				d.discoverHost(z, dh)
			}
		},
	)
}

func (d *discovery) discoverHost(zone *model.Zone, dh *model.DiscoveryHost) {
	go taskScan(d,
		func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) {
			scanHost(zone, dh, resultChan, errChan, doneChan, stopChan)
		},
		func(result interface{}) {
			h := result.(*model.Host)
			zone.AddHost(h)
			logging.Logger().Info(fmt.Sprintf("host: %v", h))
			d.sendResult("DiscoveryService.DiscoveredHost", h)
			if nil != dh.DiscoveryPort {
				d.discoverPort(h, dh.DiscoveryPort)
			}
		},
	)
}

func (d *discovery) discoverPort(host *model.Host, dp *model.DiscoveryPort) {
	go taskScan(d,
		func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) {
			scanPort(host, dp, resultChan, errChan, doneChan, stopChan)
		},
		func(result interface{}) {
			p := result.(*model.Port)
			logging.Logger().Info(fmt.Sprintf("port: %v", p))
			d.sendResult("DiscoveryService.DiscoveredPort", p)
			if nil != dp.DiscoveryService {
				d.discoverService(p, dp.DiscoveryService)
			}
		},
	)
}

func (d *discovery) discoverService(port *model.Port, ds *model.DiscoveryService) {
	go taskScan(d,
		func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}) {
			scanService(port, ds, resultChan, errChan, doneChan, stopChan)
		},
		func(result interface{}) {
			s := result.(*model.Service)
			d.sendResult("DiscoveryService.DiscoveredService", s)
			logging.Logger().Info(fmt.Sprintf("service: %s(%s)[%s:%d]", s.ServiceName, s.CryptoType, port.Host.IP, port.PortNumber))
		},
	)
}

func (d *discovery) sendResult(method string, args ...interface{}) {
	go RPCServlet.Send(method, args...)
	// go notify.Notifier.Notify(method, args...)
}

func (d *discovery) sendError() {

}

func taskScan(d *discovery, task func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}, stopChan chan struct{}), onResult func(result interface{})) {
	d.stopWg.Add(1)
	resultChan := make(chan interface{})
	errChan := make(chan error)
	stopChan := make(chan struct{})
	doneChan := make(chan struct{})

	defer func() {
		close(resultChan)
		close(errChan)
		close(stopChan)
		close(doneChan)
		d.stopWg.Done()
	}()

	go task(resultChan, errChan, doneChan, stopChan)

	for {
		select {
		case r := <-resultChan:
			onResult(r)
		case err := <-errChan:
			logging.Logger().Info(fmt.Sprintf("task err: %v", err))
		case <-doneChan:
			logging.Logger().Debug(fmt.Sprintf("Discovery: task is complete"))
			return
		case <-d.stopChan:
			close(stopChan)
			<-doneChan
			return
		}
	}
}