diff --git a/commons/pcap/pcap.go b/commons/pcap/pcap.go index 839d9ce..e865d5b 100644 --- a/commons/pcap/pcap.go +++ b/commons/pcap/pcap.go @@ -33,18 +33,20 @@ func RetainScanner(zone *model.Zone) (PCapScanner, error) { } func ReleaseScanner(zone *model.Zone) { - mtx.Lock() - defer mtx.Unlock() - if ps, ok := instances[zone.Network]; ok { - go func() { - time.Sleep(2 * time.Second) + go func() { + time.Sleep(2 * time.Second) + + mtx.Lock() + defer mtx.Unlock() + + if ps, ok := instances[zone.Network]; ok { if ps.release() { ps.stop() delete(instances, zone.Network) } - }() - } + } + }() } func ReleaseScannerAll() { diff --git a/discovery/discovery.go b/discovery/discovery.go index d47ec3d..7ef4c99 100644 --- a/discovery/discovery.go +++ b/discovery/discovery.go @@ -73,9 +73,9 @@ func (d *discovery) stop() { } func (d *discovery) discoverZone(dz *model.DiscoveryZone) { - taskScan(d, - func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}) { - scanZone(dz, resultChan, errChan, doneChan) + 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) @@ -94,9 +94,9 @@ func (d *discovery) discoverZone(dz *model.DiscoveryZone) { } func (d *discovery) discoverHost(zone *model.Zone, dh *model.DiscoveryHost) { - taskScan(d, - func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}) { - scanHost(zone, dh, resultChan, errChan, doneChan) + 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) @@ -109,9 +109,9 @@ func (d *discovery) discoverHost(zone *model.Zone, dh *model.DiscoveryHost) { } func (d *discovery) discoverPort(host *model.Host, dp *model.DiscoveryPort) { - taskScan(d, - func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}) { - scanPort(host, dp, resultChan, errChan, doneChan) + 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) @@ -124,9 +124,9 @@ func (d *discovery) discoverPort(host *model.Host, dp *model.DiscoveryPort) { } func (d *discovery) discoverService(port *model.Port, ds *model.DiscoveryService) { - taskScan(d, - func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}) { - scanService(port, ds, resultChan, errChan, doneChan) + 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) @@ -143,20 +143,22 @@ func (d *discovery) sendError() { } -func taskScan(d *discovery, task func(resultChan chan interface{}, errChan chan error, doneChan chan struct{}), onResult func(result interface{})) { +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) + go task(resultChan, errChan, doneChan, stopChan) for { select { @@ -168,6 +170,8 @@ func taskScan(d *discovery, task func(resultChan chan interface{}, errChan chan logging.Logger().Debug(fmt.Sprintf("Discovery: task is complete")) return case <-d.stopChan: + close(stopChan) + <-doneChan return } } diff --git a/discovery/host.go b/discovery/host.go index 086aa25..fbc0dc8 100644 --- a/discovery/host.go +++ b/discovery/host.go @@ -9,22 +9,24 @@ import ( "git.loafle.net/overflow/overflow_discovery/discovery/ipv6" ) -func scanHost(zone *model.Zone, dh *model.DiscoveryHost, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}) { +func scanHost(zone *model.Zone, dh *model.DiscoveryHost, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}, stopChan chan struct{}) { + defer func() { + doneChan <- struct{}{} + }() + _, ipNet, err := net.ParseCIDR(zone.Network) if nil != err { errChan <- err - doneChan <- struct{}{} return } switch len(ipNet.IP) { case net.IPv4len: - ipv4.ScanHost(zone, dh, resultChan, errChan, doneChan) + ipv4.ScanHost(zone, dh, resultChan, errChan, stopChan) case net.IPv6len: - ipv6.ScanHost(zone, dh, resultChan, errChan, doneChan) + ipv6.ScanHost(zone, dh, resultChan, errChan, stopChan) default: errChan <- fmt.Errorf("Discovery: Not supported ip length") - doneChan <- struct{}{} return } } diff --git a/discovery/ipv4/host.go b/discovery/ipv4/host.go index 642099b..9f5911e 100644 --- a/discovery/ipv4/host.go +++ b/discovery/ipv4/host.go @@ -14,11 +14,7 @@ import ( "git.loafle.net/overflow/overflow_discovery/commons/pcap" ) -func ScanHost(zone *model.Zone, dh *model.DiscoveryHost, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}) { - defer func() { - doneChan <- struct{}{} - }() - +func ScanHost(zone *model.Zone, dh *model.DiscoveryHost, resultChan chan interface{}, errChan chan error, stopChan chan struct{}) { ps, err := pcap.RetainScanner(zone) if nil != err { errChan <- fmt.Errorf("Discovery: Cannot retain pcap instance %v", err) @@ -58,11 +54,13 @@ func ScanHost(zone *model.Zone, dh *model.DiscoveryHost, resultChan chan interfa if h := handlePacketARP(zone, cr, hosts, packet); nil != h { logging.Logger().Debug(fmt.Sprintf("Discovery: Host[%v] is founded", h)) } + case <-stopChan: + return } } }() - if err := sendARP(ps, zone, hostRanges); nil != err { + if err := sendARP(ps, zone, hostRanges, stopChan); nil != err { errChan <- err return } @@ -91,7 +89,7 @@ func handlePacketARP(zone *model.Zone, cr cidr.CIDRRanger, hosts map[string]*mod return h } -func sendARP(ps pcap.PCapScanner, zone *model.Zone, hostRanges []net.IP) error { +func sendARP(ps pcap.PCapScanner, zone *model.Zone, hostRanges []net.IP, stopChan chan struct{}) error { hwAddr, err := net.ParseMAC(zone.Mac) if nil != err { return err @@ -113,6 +111,13 @@ func sendARP(ps pcap.PCapScanner, zone *model.Zone, hostRanges []net.IP) error { if err := ps.WritePacketData(buf.Bytes()); err != nil { return err } + + select { + case <-stopChan: + return nil + default: + } + time.Sleep(time.Microsecond * 100) } return nil diff --git a/discovery/ipv4/port.go b/discovery/ipv4/port.go index 11f7999..ba53dd8 100644 --- a/discovery/ipv4/port.go +++ b/discovery/ipv4/port.go @@ -2,6 +2,6 @@ package ipv4 import "git.loafle.net/overflow/overflow_discovery/api/module/discovery/model" -func ScanPort(host *model.Host, dp *model.DiscoveryPort, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}) { +func ScanPort(host *model.Host, dp *model.DiscoveryPort, resultChan chan interface{}, errChan chan error, stopChan chan struct{}) { } diff --git a/discovery/ipv6/host.go b/discovery/ipv6/host.go index 357cd02..9a9905a 100644 --- a/discovery/ipv6/host.go +++ b/discovery/ipv6/host.go @@ -2,6 +2,6 @@ package ipv6 import "git.loafle.net/overflow/overflow_discovery/api/module/discovery/model" -func ScanHost(zone *model.Zone, dh *model.DiscoveryHost, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}) { +func ScanHost(zone *model.Zone, dh *model.DiscoveryHost, resultChan chan interface{}, errChan chan error, stopChan chan struct{}) { } diff --git a/discovery/ipv6/port.go b/discovery/ipv6/port.go index e4332c1..fd22582 100644 --- a/discovery/ipv6/port.go +++ b/discovery/ipv6/port.go @@ -2,6 +2,6 @@ package ipv6 import "git.loafle.net/overflow/overflow_discovery/api/module/discovery/model" -func ScanPort(host *model.Host, dp *model.DiscoveryPort, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}) { +func ScanPort(host *model.Host, dp *model.DiscoveryPort, resultChan chan interface{}, errChan chan error, stopChan chan struct{}) { } diff --git a/discovery/port.go b/discovery/port.go index ef0459f..ba4af2a 100644 --- a/discovery/port.go +++ b/discovery/port.go @@ -9,22 +9,24 @@ import ( "git.loafle.net/overflow/overflow_discovery/discovery/ipv6" ) -func scanPort(host *model.Host, dp *model.DiscoveryPort, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}) { +func scanPort(host *model.Host, dp *model.DiscoveryPort, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}, stopChan chan struct{}) { + defer func() { + doneChan <- struct{}{} + }() + _, ipNet, err := net.ParseCIDR(host.Zone.Network) if nil != err { errChan <- err - doneChan <- struct{}{} return } switch len(ipNet.IP) { case net.IPv4len: - ipv4.ScanPort(host, dp, resultChan, errChan, doneChan) + ipv4.ScanPort(host, dp, resultChan, errChan, stopChan) case net.IPv6len: - ipv6.ScanPort(host, dp, resultChan, errChan, doneChan) + ipv6.ScanPort(host, dp, resultChan, errChan, stopChan) default: errChan <- fmt.Errorf("Discovery: Not supported ip length") - doneChan <- struct{}{} return } } diff --git a/discovery/service.go b/discovery/service.go index 8f9b3e0..6df7389 100644 --- a/discovery/service.go +++ b/discovery/service.go @@ -2,6 +2,9 @@ package discovery import "git.loafle.net/overflow/overflow_discovery/api/module/discovery/model" -func scanService(port *model.Port, ds *model.DiscoveryService, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}) { - doneChan <- struct{}{} +func scanService(port *model.Port, ds *model.DiscoveryService, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}, stopChan chan struct{}) { + defer func() { + doneChan <- struct{}{} + }() + } diff --git a/discovery/zone.go b/discovery/zone.go index d325247..0677b6f 100644 --- a/discovery/zone.go +++ b/discovery/zone.go @@ -8,7 +8,11 @@ import ( "git.loafle.net/overflow/overflow_discovery/api/module/discovery/model" ) -func scanZone(dz *model.DiscoveryZone, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}) { +func scanZone(dz *model.DiscoveryZone, resultChan chan interface{}, errChan chan error, doneChan chan<- struct{}, stopChan chan struct{}) { + defer func() { + doneChan <- struct{}{} + }() + var err error var ifaces []net.Interface var addrs []net.Addr @@ -58,9 +62,6 @@ func scanZone(dz *model.DiscoveryZone, resultChan chan interface{}, errChan chan resultChan <- z } } - - doneChan <- struct{}{} - } func checkExclude(ep []string, iface string) bool {