optimize: 优化协程和chan退出关闭逻辑，避免出现脱管的协程和chan

XinRoom · Mar 22, 2023 · 23d2a4d · 23d2a4d
1 parent b5fbf96
commit 23d2a4d
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 33 deletions.
diff --git a/cmd/go-portScan.go b/cmd/go-portScan.go
@@ -179,37 +179,29 @@ func run(c *cli.Context) error {
 	}
 
 	go func() {
-		for {
-			select {
-			case ret := <-retChan:
-				if ret.Port == 0 {
-					single <- struct{}{}
-					return
-				}
-				if maxOpenPort > 0 {
-					ipPortNumRW.Lock()
-					if _, ok := ipPortNumMap[ret.Ip.String()]; ok {
-						ipPortNumMap[ret.Ip.String()] += 1
-					}
-					ipPortNumRW.Unlock()
+		for ret := range retChan {
+			if maxOpenPort > 0 {
+				ipPortNumRW.Lock()
+				if _, ok := ipPortNumMap[ret.Ip.String()]; ok {
+					ipPortNumMap[ret.Ip.String()] += 1
 				}
-				myLog.Println(ret.String())
-				if csvWrite != nil {
-					line := []string{ret.Ip.String(), strconv.Itoa(int(ret.Port)), ret.Service, "", "", "", ""}
-					if ret.HttpInfo != nil {
-						line[3] = ret.HttpInfo.Title
-						line[4] = strconv.Itoa(ret.HttpInfo.StatusCode)
-						line[5] = ret.HttpInfo.Server
-						line[6] = ret.HttpInfo.TlsCN
-					}
-					csvWrite.Write(line)
-					csvWrite.Flush()
-					csvFile.Sync()
+				ipPortNumRW.Unlock()
+			}
+			myLog.Println(ret.String())
+			if csvWrite != nil {
+				line := []string{ret.Ip.String(), strconv.Itoa(int(ret.Port)), ret.Service, "", "", "", ""}
+				if ret.HttpInfo != nil {
+					line[3] = ret.HttpInfo.Title
+					line[4] = strconv.Itoa(ret.HttpInfo.StatusCode)
+					line[5] = ret.HttpInfo.Server
+					line[6] = ret.HttpInfo.TlsCN
 				}
-			default:
-				time.Sleep(time.Millisecond * 10)
+				csvWrite.Write(line)
+				csvWrite.Flush()
+				csvFile.Sync()
 			}
 		}
+		single <- struct{}{}
 	}()
 
 	// Initialize the Scanner

diff --git a/core/port/syn/syn.go b/core/port/syn/syn.go
@@ -10,6 +10,7 @@ import (
 	"github.com/google/gopacket/layers"
 	"github.com/google/gopacket/pcap"
 	limiter "golang.org/x/time/rate"
+	"io"
 	"math/rand"
 	"net"
 	"sync"
@@ -105,9 +106,6 @@ func NewSynScanner(firstIp net.IP, retChan chan port.OpenIpPort, option port.Opt
 				ss.portProbeWg.Add(1)
 				ss.retChan <- t
 				ss.portProbeWg.Done()
-				if t.Port == 0 {
-					break
-				}
 			}
 		}()
 	}
@@ -141,7 +139,7 @@ func NewSynScanner(firstIp net.IP, retChan chan port.OpenIpPort, option port.Opt
 // Scan scans the dst IP address and port of this scanner.
 func (ss *SynScanner) Scan(dstIp net.IP, dst uint16) (err error) {
 	if ss.isDone {
-		return errors.New("scanner is closed")
+		return io.EOF
 	}
 
 	// 与recv协同，当队列缓冲区到达80%时降半速，90%将为1/s
@@ -250,7 +248,6 @@ func (ss *SynScanner) Wait() {
 // Close cleans up the handle and chan.
 func (ss *SynScanner) Close() {
 	ss.isDone = true
-	ss.openPortChan <- port.OpenIpPort{}
 	if ss.handle != nil {
 		ss.handle.Close()
 	}
@@ -262,6 +259,8 @@ func (ss *SynScanner) Close() {
 	}
 	ss.watchMacCacheT = nil
 	ss.watchIpStatusT = nil
+	close(ss.openPortChan)
+	close(ss.retChan)
 }
 
 // WaitLimiter Waiting for the speed limit
@@ -419,6 +418,9 @@ func (ss *SynScanner) recv() {
 		// Read in the next packet.
 		data, _, err = ss.handle.ReadPacketData()
 		if err != nil {
+			if err == io.EOF {
+				return
+			}
 			continue
 		}
 

diff --git a/core/port/tcp/tcp.go b/core/port/tcp/tcp.go
@@ -89,7 +89,7 @@ func (ts *TcpScanner) Wait() {
 // Close chan
 func (ts *TcpScanner) Close() {
 	ts.isDone = true
-	ts.retChan <- port.OpenIpPort{}
+	close(ts.retChan)
 }
 
 // WaitLimiter Waiting for the speed limit