This commit is contained in:
crusader 2018-09-03 23:14:51 +09:00
parent f5af8e1473
commit c527a84bb2
4 changed files with 36 additions and 13 deletions

View File

@ -30,6 +30,7 @@ func Scan(discoverySession session.DiscoverySession, targetHost *omd.Host) error
lock := semaphore.NewWeighted(Ulimit())
var wg sync.WaitGroup
stopChan := make(chan struct{})
timeout := 500 * time.Millisecond
ports := make(map[int]*omd.Port)
@ -52,7 +53,7 @@ Loop:
wg.Done()
}()
tryConnect(discoverySession, ports, targetHost, port, timeout)
tryConnect(discoverySession, ports, targetHost, port, timeout, stopChan)
}(portNumber)
timer := time.NewTimer(time.Microsecond * 100)
@ -60,6 +61,7 @@ Loop:
select {
case <-timer.C:
case <-discoverySession.StopChan():
close(stopChan)
return nil
}
}
@ -69,7 +71,7 @@ Loop:
return nil
}
func tryConnect(discoverySession session.DiscoverySession, ports map[int]*omd.Port, targetHost *omd.Host, port int, timeout time.Duration) {
func tryConnect(discoverySession session.DiscoverySession, ports map[int]*omd.Port, targetHost *omd.Host, port int, timeout time.Duration, stopChan <-chan struct{}) {
addr := net.JoinHostPort(targetHost.Address, strconv.Itoa(port))
conn, err := net.DialTimeout("tcp", addr, timeout)
dp := discoverySession.DiscoverPort()
@ -77,12 +79,12 @@ func tryConnect(discoverySession session.DiscoverySession, ports map[int]*omd.Po
if err != nil {
if strings.Contains(err.Error(), "too many open files") {
select {
case <-discoverySession.StopChan():
case <-stopChan:
return
default:
}
time.Sleep(timeout)
tryConnect(discoverySession, ports, targetHost, port, timeout)
tryConnect(discoverySession, ports, targetHost, port, timeout, stopChan)
}
return
}

View File

@ -30,6 +30,7 @@ func Scan(discoverySession session.DiscoverySession, targetHost *omd.Host) error
lock := semaphore.NewWeighted(Ulimit())
var wg sync.WaitGroup
stopChan := make(chan struct{})
timeout := 500 * time.Millisecond
ports := make(map[int]*omd.Port)
@ -52,7 +53,7 @@ Loop:
wg.Done()
}()
tryConnect(discoverySession, ports, targetHost, port, timeout)
tryConnect(discoverySession, ports, targetHost, port, timeout, stopChan)
}(portNumber)
timer := time.NewTimer(time.Microsecond * 100)
@ -60,6 +61,7 @@ Loop:
select {
case <-timer.C:
case <-discoverySession.StopChan():
close(stopChan)
return nil
}
}
@ -69,7 +71,7 @@ Loop:
return nil
}
func tryConnect(discoverySession session.DiscoverySession, ports map[int]*omd.Port, targetHost *omd.Host, port int, timeout time.Duration) {
func tryConnect(discoverySession session.DiscoverySession, ports map[int]*omd.Port, targetHost *omd.Host, port int, timeout time.Duration, stopChan <-chan struct{}) {
addr := net.JoinHostPort(targetHost.Address, strconv.Itoa(port))
conn, err := net.DialTimeout("tcp", addr, timeout)
dp := discoverySession.DiscoverPort()
@ -77,12 +79,12 @@ func tryConnect(discoverySession session.DiscoverySession, ports map[int]*omd.Po
if err != nil {
if strings.Contains(err.Error(), "too many open files") {
select {
case <-discoverySession.StopChan():
case <-stopChan:
return
default:
}
time.Sleep(timeout)
tryConnect(discoverySession, ports, targetHost, port, timeout)
tryConnect(discoverySession, ports, targetHost, port, timeout, stopChan)
}
return
}

View File

@ -48,6 +48,12 @@ LOOP:
Zone: discoverySession.Zone(),
DiscoveredDate: omu.NowPtr(),
})
select {
case <-discoverySession.StopChan():
return nil
default:
}
}
return nil

View File

@ -25,6 +25,7 @@ func scanTCP(discoverySession session.DiscoverySession, targetPort *omd.Port) er
matchCtx := osm.NewMatchCtx(hostAddress, portNumber)
connectors := newConnectors()
stopChan := make(chan struct{})
buf := make([]byte, 1024)
var discoveredMatcher osm.Matcher
@ -46,16 +47,16 @@ LOOP:
}
if 0 < n {
discoveredMatcher = hadlePrePacket(matchCtx, _connector, conn, osm.NewPacket(buf, n))
discoveredMatcher = hadlePrePacket(matchCtx, _connector, conn, osm.NewPacket(buf, n), stopChan)
} else {
conn.Close()
discoveredMatcher = hadlePostPacket(matchCtx, _connector, targetPort, nil)
discoveredMatcher = hadlePostPacket(matchCtx, _connector, targetPort, nil, stopChan)
}
if nil != discoveredMatcher {
if "HTTP" == discoveredMatcher.Key() {
hsm := matcher.GetHTTPSubMatchers()
if _discoveredMatcher := hadlePostPacket(matchCtx, _connector, targetPort, hsm); _discoveredMatcher != nil {
if _discoveredMatcher := hadlePostPacket(matchCtx, _connector, targetPort, hsm, stopChan); _discoveredMatcher != nil {
discoveredMatcher = _discoveredMatcher
}
}
@ -86,7 +87,7 @@ LOOP:
return nil
}
func hadlePrePacket(matchCtx *osm.MatchCtx, _connector connector, conn net.Conn, packet *osm.Packet) osm.Matcher {
func hadlePrePacket(matchCtx *osm.MatchCtx, _connector connector, conn net.Conn, packet *osm.Packet, stopChan <-chan struct{}) osm.Matcher {
defer func() {
conn.Close()
}()
@ -138,12 +139,18 @@ LOOP:
if nil != discoveredMatcher {
return discoveredMatcher
}
select {
case <-stopChan:
return nil
default:
}
}
return nil
}
func hadlePostPacket(matchCtx *osm.MatchCtx, _connector connector, targetPort *omd.Port, limitedMatchers []osm.Matcher) osm.Matcher {
func hadlePostPacket(matchCtx *osm.MatchCtx, _connector connector, targetPort *omd.Port, limitedMatchers []osm.Matcher, stopChan <-chan struct{}) osm.Matcher {
matchers := matcher.GetTCPMatchers(false)
if nil != limitedMatchers {
matchers = limitedMatchers
@ -205,6 +212,12 @@ LOOP:
if nil != discoveredMatcher {
return discoveredMatcher
}
select {
case <-stopChan:
return nil
default:
}
}
return nil