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) }