improve performance by reworking the port scan (#4)

This results in decreasing the fast scan time from 0.309s to 0.020s for localhost on my machine. Expose the underlying errors in the public api, to allow users to handle these themselves.
JustinTimperio · Mar 16, 2021 · de19ce1 · de19ce1
1 parent 2e4b0c2
commit de19ce1
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 98 deletions.
diff --git a/gomap.go b/gomap.go
@@ -20,61 +20,53 @@ type IPScanResult struct {
 }
 
 // RangeScanResult contains multiple IPScanResults
-type RangeScanResult struct {
-	all []IPScanResult
-}
+type RangeScanResult []*IPScanResult
 
 // ScanIP scans a single IP for open ports
-func ScanIP(hostname string, fastscan bool) IPScanResult {
-	ipScan, err := scanIPPorts(hostname, "tcp", fastscan)
-	if err != nil {
-		fmt.Println(err)
-	}
-	return ipScan
+func ScanIP(hostname string, fastscan bool) (*IPScanResult, error) {
+	return scanIPPorts(hostname, "tcp", fastscan)
 }
 
 // ScanRange scans every address on a CIDR for open ports
-func ScanRange(fastscan bool) RangeScanResult {
-	rangeScan, err := scanIPRange("tcp", fastscan)
-	if err != nil {
-		fmt.Println(err)
-	}
-	return rangeScan
+func ScanRange(fastscan bool) (RangeScanResult, error) {
+	return scanIPRange("tcp", fastscan)
 }
 
-// Provides a string that prints the results of a single ScanIp
+// String with the results of a single scanned IP
 func (results *IPScanResult) String() string {
 	b := bytes.NewBuffer(nil)
 	ip := results.ip[len(results.ip)-1]
-	b.WriteString(fmt.Sprintf("\nHost: %s (%s)\n", results.hostname, ip.String()))
-	active := false
 
+	fmt.Fprintf(b, "\nHost: %s (%s)\n", results.hostname, ip)
+
+	active := false
 	for _, r := range results.results {
 		if r.State {
 			active = true
 			break
 		}
 	}
 	if active {
-		b.WriteString(fmt.Sprintf("\t|     %s	%s\n", "Port", "Service"))
-		b.WriteString(fmt.Sprintf("\t|     %s	%s\n", "----", "-------"))
+		fmt.Fprintf(b, "\t|     %s	%s\n", "Port", "Service")
+		fmt.Fprintf(b, "\t|     %s	%s\n", "----", "-------")
 		for _, v := range results.results {
 			if v.State {
-				b.WriteString(fmt.Sprintf("\t|---- %d	%s\n", v.Port, v.Service))
+				fmt.Fprintf(b, "\t|---- %d	%s\n", v.Port, v.Service)
 			}
 		}
 	} else if results.hostname != "Unknown" {
-		b.WriteString(fmt.Sprintf("\t|---- %s\n", "No Open Ports Found"))
+		fmt.Fprintf(b, "\t|---- %s\n", "No Open Ports Found")
 	}
 	return b.String()
 }
 
-// Provides a string that prints the results of a multiple ScanIP's
-func (results *RangeScanResult) String() string {
+// String with the results of multiple scanned IP's
+func (results RangeScanResult) String() string {
 	b := bytes.NewBuffer(nil)
-	for _, r := range results.all {
+	for _, r := range results {
 		ip := r.ip[len(r.ip)-1]
-		b.WriteString(fmt.Sprintf("\nHost: %s (%s)\n", r.hostname, ip.String()))
+
+		fmt.Fprintf(b, "\nHost: %s (%s)\n", r.hostname, ip)
 		active := false
 
 		for _, r := range r.results {
@@ -84,17 +76,17 @@ func (results *RangeScanResult) String() string {
 			}
 		}
 		if active {
-			b.WriteString(fmt.Sprintf("\t|     %s	%s\n", "Port", "Service"))
-			b.WriteString(fmt.Sprintf("\t|     %s	%s\n", "----", "-------"))
+			fmt.Fprintf(b, "\t|     %s	%s\n", "Port", "Service")
+			fmt.Fprintf(b, "\t|     %s	%s\n", "----", "-------")
 			for _, v := range r.results {
 				if v.State {
-					b.WriteString(fmt.Sprintf("\t|---- %d	%s\n", v.Port, v.Service))
+					fmt.Fprintf(b, "\t|---- %d	%s\n", v.Port, v.Service)
 				}
 			}
 		} else if r.hostname != "Unknown" {
-			b.WriteString(fmt.Sprintf("\t|---- %s\n", "No Open Ports Found"))
+			fmt.Fprintf(b, "\t|---- %s\n", "No Open Ports Found")
 		}
-
 	}
+
 	return b.String()
 }
diff --git a/gomap_scan.go b/gomap_scan.go
@@ -14,35 +14,29 @@ import (
 // I am fairly happy with this code since its just iterating
 // over scanIPPorts. Most issues are deeper in the code.
 func scanIPRange(proto string, fastscan bool) (RangeScanResult, error) {
-	var results []IPScanResult
 	iprange := getLocalRange()
 	hosts := createHostRange(iprange)
 
+	var results RangeScanResult
 	for _, h := range hosts {
 		scan, err := scanIPPorts(h, proto, fastscan)
 		if err != nil {
 			continue
 		}
 		results = append(results, scan)
 	}
-	rangeScan := RangeScanResult{all: results}
-	return rangeScan, nil
+
+	return results, nil
 }
 
 // scanIPPorts scans a list of ports on <hostname> <protocol>
-func scanIPPorts(hostname string, proto string, fastscan bool) (IPScanResult, error) {
-	var (
-		results []portResult
-		scanned IPScanResult
-		tasks   int
-		start   = 0
-		end     = 50000
-	)
+func scanIPPorts(hostname string, proto string, fastscan bool) (*IPScanResult, error) {
+	var results []portResult
 
 	// checks if device is online
 	addr, err := net.LookupIP(hostname)
 	if err != nil {
-		return scanned, err
+		return nil, err
 	}
 
 	// This gets the device name. ('/etc/hostname')
@@ -53,80 +47,63 @@ func scanIPPorts(hostname string, proto string, fastscan bool) (IPScanResult, er
 	hname, err := net.LookupAddr(hostname)
 	if fastscan {
 		if err != nil {
-			return scanned, err
-		}
-		tasks = len(commonlist)
-
-	} else {
-		if err != nil {
-			hname = append(hname, "Unknown")
+			return nil, err
 		}
-		tasks = len(detailedlist)
+	} else if err != nil {
+		hname = append(hname, "Unknown")
 	}
 
-	// Opens pool of connections to crawl ports
-	// After 3 failed semaphore implementations I settled on a less
-	// generic piece of code that launches routines in batches then waits
-	// I am open to new solutions to this brick of code
-	resultChannel := make(chan portResult, tasks)
-	if fastscan {
-		for i := start; i <= end; i++ {
-			if service, ok := commonlist[i]; ok {
-				if i%100 == 0 {
-					fmt.Printf("\033[2K\rHost: %s | Ports Scanned %d/%d", hostname, len(resultChannel), tasks)
-				}
-				// Spaces out goroutines slightly
-				if i%25 == 0 {
-					time.Sleep(50 * time.Millisecond)
-				}
-				go scanPort(resultChannel, proto, hostname, service, i, fastscan)
-			}
+	in := make(chan int)
+	go func() {
+		for i := 0; i <= 65535; i++ {
+			in <- i
 		}
 
+		close(in)
+	}()
+
+	var list map[int]string
+	if fastscan {
+		list = commonlist
 	} else {
-		for i := start; i <= end; i++ {
-			if service, ok := detailedlist[i]; ok {
-				if i%100 == 0 {
-					fmt.Printf("\033[2K\rHost: %s | Ports Scanned %d/%d", hostname, len(resultChannel), tasks)
-				}
-				// Spaces out goroutines slightly
-				if i%25 == 0 {
-					time.Sleep(50 * time.Millisecond)
-				}
-				go scanPort(resultChannel, proto, hostname, service, i, fastscan)
+		list = detailedlist
+	}
+
+	tasks := len(list)
+
+	resultChannel := make(chan portResult, tasks)
+	worker := func() {
+		for port := range in {
+			if service, ok := list[port]; ok {
+				scanPort(resultChannel, proto, hostname, service, port)
 			}
 		}
 	}
 
-	// This waits for all routines to finish.
-	// Overall this has been more performant than wait-groups
-	// and it allows for an active counter to display progress
-	for {
-		if len(resultChannel) == tasks {
-			close(resultChannel)
-			break
-		} else {
-			fmt.Printf("\033[2K\rHost: %s | Ports Scanned %d/%d", hostname, len(resultChannel), tasks)
-			time.Sleep(100 * time.Millisecond)
-		}
+	for i := 0; i < 10; i++ {
+		go worker()
 	}
 
 	// Combines all results from resultChannel and
 	// returns a IPScanResult struct
 	for result := range resultChannel {
 		results = append(results, result)
+		fmt.Printf("\033[2K\rHost: %s | Ports Scanned %d/%d", hostname, len(results), tasks)
+
+		if len(results) == tasks {
+			close(resultChannel)
+		}
 	}
 
-	scanned = IPScanResult{
+	return &IPScanResult{
 		hostname: hname[0],
 		ip:       addr,
 		results:  results,
-	}
-	return scanned, nil
+	}, nil
 }
 
 // scanPort scans a single ip port combo
-func scanPort(resultChannel chan<- portResult, protocol, hostname, service string, port int, fastscan bool) {
+func scanPort(resultChannel chan<- portResult, protocol, hostname, service string, port int) {
 	// This detection method only works on some types of services
 	// but is a reasonable solution for this application
 	result := portResult{Port: port, Service: service}
@@ -141,7 +118,6 @@ func scanPort(resultChannel chan<- portResult, protocol, hostname, service strin
 	defer conn.Close()
 	result.State = true
 	resultChannel <- result
-	return
 }
 
 // createHostRange converts a input ip addr string to a slice of ips on the cidr

diff --git a/gomap_test.go b/gomap_test.go
@@ -9,9 +9,8 @@ import (
 
 func TestMain(m *testing.M) {
 	fastscan := true
-	results := gomap.ScanRange(fastscan)
-	// results := gomap.ScanIP("192.168.1.120", fastscan)
-
-	fmt.Printf(results.String())
+	results, _ := gomap.ScanRange(fastscan)
+	// results, _ := gomap.ScanIP("127.0.0.1", fastscan)
 
+	fmt.Println(results)
 }