From fa310c5c0f34fa3b83f79b68058d437d067b8305 Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Sun, 11 Feb 2024 19:58:54 -0500
Subject: [PATCH] Fix single domain address

---
 main.go    |  2 +-
 scanner.go | 18 ++++++++++--------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/main.go b/main.go
index eb8af96..83cb241 100644
--- a/main.go
+++ b/main.go
@@ -108,7 +108,7 @@ func main() {
 	for i := 0; i < thread; i++ {
 		go func() {
 			for ip := range hostChan {
-				s.Scan(ip, outCh, true)
+				ip = s.Scan(ip, outCh, true)
 				if ip.Infinity { // only one ip
 					for i := 0; i < thread - 1; i++ {
 						go s.Scan(ip, outCh, i%2 == 1)
diff --git a/scanner.go b/scanner.go
index 40cda96..c3eed29 100644
--- a/scanner.go
+++ b/scanner.go
@@ -17,7 +17,7 @@ type Scanner struct {
 	low  net.IP
 }
 
-func (s *Scanner) Scan(host Host, out chan<- string, increment bool) {
+func (s *Scanner) Scan(host Host, out chan<- string, increment bool) Host {
 	if host.Infinity && host.IP != nil {
 		s.mu.Lock()
 		if s.high == nil {
@@ -34,18 +34,19 @@ func (s *Scanner) Scan(host Host, out chan<- string, increment bool) {
 		}
 		s.mu.Unlock()
 	}
-	ScanTLS(host, out, increment)
+	host = ScanTLS(host, out, increment)
 	if host.Infinity && host.IP != nil {
 		go s.Scan(host, out, increment)
 	}
+	return host
 }
 
-func ScanTLS(host Host, out chan<- string, increment bool) {
+func ScanTLS(host Host, out chan<- string, increment bool) Host {
 	if host.IP == nil {
 		ips, err := net.LookupIP(host.Origin)
 		if err != nil {
 			slog.Debug("Failed to lookup", "origin", host.Origin, "err", err)
-			return
+			return host
 		}
 		var arr []net.IP
 		for _, ip := range ips {
@@ -55,7 +56,7 @@ func ScanTLS(host Host, out chan<- string, increment bool) {
 		}
 		if len(arr) == 0 {
 			slog.Debug("No IP found", "origin", host.Origin)
-			return
+			return host
 		}
 		host.IP = arr[0]
 	}
@@ -63,13 +64,13 @@ func ScanTLS(host Host, out chan<- string, increment bool) {
 	conn, err := net.DialTimeout("tcp", hostPort, time.Duration(timeout)*time.Second)
 	if err != nil {
 		slog.Debug("Cannot dial", "target", hostPort)
-		return
+		return host
 	}
 	defer conn.Close()
 	err = conn.SetDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
 	if err != nil {
 		slog.Error("Error setting deadline", "err", err)
-		return
+		return host
 	}
 	tlsCfg := &tls.Config{
 		InsecureSkipVerify: true,
@@ -83,7 +84,7 @@ func ScanTLS(host Host, out chan<- string, increment bool) {
 	err = c.Handshake()
 	if err != nil {
 		slog.Debug("TLS handshake failed", "target", hostPort)
-		return
+		return host
 	}
 	state := c.ConnectionState()
 	alpn := state.NegotiatedProtocol
@@ -101,6 +102,7 @@ func ScanTLS(host Host, out chan<- string, increment bool) {
 	log("Connected to target", "feasible", feasible, "ip", host.IP.String(),
 		"origin", host.Origin,
 		"tls", tls.VersionName(state.Version), "alpn", alpn, "cert-domain", domain, "cert-issuer", issuers)
+	return host
 }
 
 func nextIP(ip net.IP, increment bool) net.IP {