probe/discovery/session/discovery-session.go

package session

import (
	"encoding/json"
	"fmt"
	"sync"

	omd "git.loafle.net/overflow/model/discovery"
	"git.loafle.net/overflow_scanner/probe/discovery/types"
)

type DiscoverySession interface {
	Zone() *omd.Zone
	Host() *omd.Host
	Port() *omd.Port

	DiscoverHost() *omd.DiscoverHost
	DiscoverPort() *omd.DiscoverPort
	DiscoverService() *omd.DiscoverService

	InitWithDiscoverHost(request types.DiscoveryRequest, zone *omd.Zone, discoverHost *omd.DiscoverHost)
	InitWithDiscoverPort(request types.DiscoveryRequest, host *omd.Host, discoverPort *omd.DiscoverPort)
	InitWithDiscoverService(request types.DiscoveryRequest, port *omd.Port, discoverService *omd.DiscoverService)
	AddHost(host *omd.Host) *omd.Host
	AddPort(port *omd.Port) *omd.Port
	AddService(service *omd.Service) *omd.Service
}

type ofDiscoverySession struct {
	discoveryRequest types.DiscoveryRequest

	zone *omd.Zone
	host *omd.Host
	port *omd.Port

	discoverHost    *omd.DiscoverHost
	discoverPort    *omd.DiscoverPort
	discoverService *omd.DiscoverService

	hosts    map[string]*omd.Host
	ports    map[*omd.Host]map[json.Number]map[string]*omd.Port
	services map[*omd.Port]map[string]map[string]*omd.Service
}

func (ds *ofDiscoverySession) init() {
	ds.discoveryRequest = nil
	ds.zone = nil
	ds.host = nil
	ds.port = nil
	ds.discoverHost = nil
	ds.discoverPort = nil
	ds.discoverService = nil

	ds.hosts = make(map[string]*omd.Host)
	ds.ports = make(map[*omd.Host]map[json.Number]map[string]*omd.Port)
	ds.services = make(map[*omd.Port]map[string]map[string]*omd.Service)
}

func (ds *ofDiscoverySession) initWithRequest(request types.DiscoveryRequest) {
	ds.init()

	ds.discoveryRequest = request
}

func (ds *ofDiscoverySession) InitWithDiscoverHost(request types.DiscoveryRequest, zone *omd.Zone, discoverHost *omd.DiscoverHost) {
	ds.initWithRequest(request)

	ds.setZone(zone)
	ds.setDiscoverHost(discoverHost)
}

func (ds *ofDiscoverySession) InitWithDiscoverPort(request types.DiscoveryRequest, host *omd.Host, discoverPort *omd.DiscoverPort) {
	ds.initWithRequest(request)

	ds.setHost(host)
	ds.setDiscoverPort(discoverPort)
}

func (ds *ofDiscoverySession) InitWithDiscoverService(request types.DiscoveryRequest, port *omd.Port, discoverService *omd.DiscoverService) {
	ds.initWithRequest(request)

	ds.setPort(port)
	ds.setDiscoverService(discoverService)
}

func (ds *ofDiscoverySession) Zone() *omd.Zone {
	return ds.zone
}
func (ds *ofDiscoverySession) Host() *omd.Host {
	return ds.host
}
func (ds *ofDiscoverySession) Port() *omd.Port {
	return ds.port
}

func (ds *ofDiscoverySession) DiscoverHost() *omd.DiscoverHost {
	return ds.discoverHost
}
func (ds *ofDiscoverySession) DiscoverPort() *omd.DiscoverPort {
	return ds.discoverPort
}
func (ds *ofDiscoverySession) DiscoverService() *omd.DiscoverService {
	return ds.discoverService
}

func (ds *ofDiscoverySession) AddHost(host *omd.Host) *omd.Host {
	h := ds.findHost(host, true)

	if "" == h.Mac && "" != host.Mac {
		h.Mac = host.Mac
	}

	h.Meta = ds.appendMeta(h.Meta, host.Meta)

	ds.discoveryRequest.SendMessage(types.DiscoveryMessageTypeHost, h, nil)

	return h
}

func (ds *ofDiscoverySession) AddPort(port *omd.Port) *omd.Port {
	p := ds.findPort(port, true)

	p.Meta = ds.appendMeta(p.Meta, port.Meta)

	return p
}

func (ds *ofDiscoverySession) AddService(service *omd.Service) *omd.Service {
	s := ds.findService(service, true)

	s.Meta = ds.appendMeta(s.Meta, service.Meta)

	return s
}

func (ds *ofDiscoverySession) findHost(host *omd.Host, add bool) *omd.Host {
	h, ok := ds.hosts[host.Address]
	if !ok {
		if !add {
			return nil
		}
		ds.hosts[host.Address] = host
		h = host
	}

	return h
}

func (ds *ofDiscoverySession) findPort(port *omd.Port, add bool) *omd.Port {
	h := ds.findHost(port.Host, false)
	if nil == h {
		return nil
	}

	hostPorts, ok := ds.ports[h]
	if !ok {
		if !add {
			return nil
		}

		ds.ports[h] = make(map[json.Number]map[string]*omd.Port)
		hostPorts = ds.ports[h]
	}

	ports, ok := hostPorts[port.PortNumber]
	if !ok {
		if !add {
			return nil
		}

		hostPorts[port.PortNumber] = make(map[string]*omd.Port)
		ports = hostPorts[port.PortNumber]
	}

	p, ok := ports[port.MetaPortType.Key]
	if !ok {
		if !add {
			return nil
		}

		ports[port.MetaPortType.Key] = port
		p = ports[port.MetaPortType.Key]
	}

	return p
}

func (ds *ofDiscoverySession) findService(service *omd.Service, add bool) *omd.Service {
	p := ds.findPort(service.Port, false)
	if nil == p {
		return nil
	}

	portServices, ok := ds.services[p]
	if !ok {
		if !add {
			return nil
		}

		ds.services[p] = make(map[string]map[string]*omd.Service)
		portServices = ds.services[p]
	}

	services, ok := portServices[service.Key]
	if !ok {
		if !add {
			return nil
		}

		portServices[service.Key] = make(map[string]*omd.Service)
		services = portServices[service.Key]
	}

	s, ok := services[service.MetaCryptoType.Key]
	if !ok {
		if !add {
			return nil
		}

		services[service.MetaCryptoType.Key] = service
		s = services[service.MetaCryptoType.Key]
	}

	return s
}

func (ds *ofDiscoverySession) appendMeta(oriMeta map[string]string, newMeta map[string]string) map[string]string {
	if nil == newMeta {
		return oriMeta
	}
	if nil == oriMeta {
		return newMeta
	}

LOOP:
	for k, v := range oriMeta {
		_v, _ok := oriMeta[k]
		if !_ok {
			oriMeta[k] = v
			continue LOOP
		}
		if v == _v {
			continue LOOP
		}
		oriMeta[k] = fmt.Sprintf("%s|||%s", _v, v)
	}

	return oriMeta
}

func (ds *ofDiscoverySession) setZone(zone *omd.Zone) {
	if nil == zone {
		return
	}
	ds.zone = zone
}

func (ds *ofDiscoverySession) setHost(host *omd.Host) {
	if nil == host {
		return
	}
	ds.setZone(host.Zone)
	ds.host = host
}

func (ds *ofDiscoverySession) setPort(port *omd.Port) {
	if nil == port {
		return
	}
	ds.setHost(port.Host)
	ds.port = port
}

func (ds *ofDiscoverySession) setDiscoverHost(discoverHost *omd.DiscoverHost) {
	if nil == discoverHost {
		return
	}

	ds.discoverHost = discoverHost
	ds.setDiscoverPort(discoverHost.DiscoverPort)
}

func (ds *ofDiscoverySession) setDiscoverPort(discoverPort *omd.DiscoverPort) {
	if nil == discoverPort {
		return
	}

	ds.discoverPort = discoverPort
	ds.setDiscoverService(discoverPort.DiscoverService)
}

func (ds *ofDiscoverySession) setDiscoverService(discoverService *omd.DiscoverService) {
	if nil == discoverService {
		return
	}

	ds.discoverService = discoverService
}

var discoverySessionPool sync.Pool

func RetainDiscoverySession() *ofDiscoverySession {
	v := discoverySessionPool.Get()
	var ds *ofDiscoverySession
	if v == nil {
		ds = &ofDiscoverySession{}
	} else {
		ds = v.(*ofDiscoverySession)
	}

	return ds
}

func ReleaseDiscoverySession(ds *ofDiscoverySession) {
	ds.discoveryRequest = nil
	ds.zone = nil
	ds.host = nil
	ds.port = nil
	ds.discoverHost = nil
	ds.discoverPort = nil
	ds.discoverService = nil

	ds.hosts = nil
	ds.ports = nil
	ds.services = nil

	discoverySessionPool.Put(ds)
}