使用channel优化代码逻辑，复用函数

Author: dushixiang

4dnat.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"net"
 	"os"
+	"strings"
 	"sync"
 	"time"
 )
@@ -67,82 +68,44 @@ func copyIO(src, dest net.Conn, wg *sync.WaitGroup) {
 func listener(listenPort0, listenPort1 string) {
 	ln0 := listen(listenPort0)
 	ln1 := listen(listenPort1)
-	fmt.Printf("[#] 4dnat listen port on: [%s] [%s]\n", listenPort0, listenPort1)
+	fmt.Printf("[#] 4dnat listen port on: [%s %s],[%s %s]\n", ln0.Addr().Network(), ln0.Addr().String(), ln1.Addr().Network(), ln1.Addr().String())
 
 	for true {
 		cc := make(chan net.Conn, 2)
 
-		go loopAccept(cc, listenPort0, ln0)
-		go loopAccept(cc, listenPort1, ln1)
+		go accept(cc, ln0)
+		go accept(cc, ln1)
 
 		conn0, conn1 := <-cc, <-cc
 		mutualCopyIO(conn0, conn1)
 	}
-
-}
-
-func loopAccept(cc chan net.Conn, listenPort string, ln net.Listener) {
-	for true {
-		fmt.Printf("[#] 4dnat waiting for client to connect port [%s]\n", listenPort)
-		c, err := accept(ln)
-		if err != nil {
-			continue
-		}
-		cc <- c
-		break
-	}
 }
 
 func forward(listenPort string, targetAddress string) {
 	ln := listen(listenPort)
-	fmt.Printf("[#] 4dnat listen on: [%s] forward to: [%s]\n", listenPort, targetAddress)
-	for {
-		conn0, err := accept(ln)
-		if err != nil {
-			continue
-		}
-		go func() {
-			fmt.Printf("[#] try to connect: [%s] \n", targetAddress)
-			// after server accept will be connect the target address,if failed will be retry.
-			for true {
-				conn1, err := dial(targetAddress)
-				if err != nil {
-					fmt.Printf("[#] retry to connect: [%s] after [%d] second\n", targetAddress, RetryInterval)
-					time.Sleep(time.Duration(RetryInterval) * time.Second)
-					continue
-				}
-
-				go mutualCopyIO(conn0, conn1)
-				break
-			}
-		}()
+	fmt.Printf("[#] 4dnat listen on: [%s %s] forward to: [%s]\n", ln.Addr().Network(), ln.Addr().String(), targetAddress)
+	for true {
+		cc := make(chan net.Conn, 2)
+
+		go accept(cc, ln)
+		go dial(cc, targetAddress)
+
+		conn0, conn1 := <-cc, <-cc
+		go mutualCopyIO(conn0, conn1)
 	}
 }
 
 func agent(targetAddress0 string, targetAddress1 string) {
-
 	fmt.Printf("[#] 4dnat agent with: [%s] [%s]\n", targetAddress0, targetAddress1)
-
-	var conn0 net.Conn = nil
 	for {
-		if conn0 == nil {
-			conn, err0 := dial(targetAddress0)
-			if err0 != nil {
-				time.Sleep(time.Duration(RetryInterval) * time.Second)
-				continue
-			}
-			conn0 = conn
-		}
+		cc := make(chan net.Conn, 2)
 
-		conn1, err1 := dial(targetAddress1)
-		if err1 != nil {
-			time.Sleep(time.Duration(RetryInterval) * time.Second)
-			continue
-		}
+		go dial(cc, targetAddress0)
+		go dial(cc, targetAddress1)
 
+		conn0, conn1 := <-cc, <-cc
 		mutualCopyIO(conn0, conn1)
 	}
-
 }
 
 func mutualCopyIO(conn0, conn1 net.Conn) {
@@ -153,28 +116,40 @@ func mutualCopyIO(conn0, conn1 net.Conn) {
 	wg.Wait()
 }
 
-func dial(targetAddress string) (net.Conn, error) {
-	conn, err := net.Dial("tcp", targetAddress)
-	if err != nil {
-		fmt.Printf("[x] connect [%s] error [%s].\n", targetAddress, err.Error())
-		return conn, err
+func dial(cc chan net.Conn, targetAddress string) {
+	for true {
+		conn, err := net.Dial("tcp", targetAddress)
+		if err != nil {
+			fmt.Printf("[x] connect [%s] error [%s].\n", targetAddress, err.Error())
+			fmt.Printf("[#] retry to connect: [%s] after [%d] second\n", targetAddress, RetryInterval)
+			time.Sleep(time.Duration(RetryInterval) * time.Second)
+			continue
+		}
+		fmt.Printf("[+] [%s]->[%s] connected to target.\n", conn.LocalAddr().String(), targetAddress)
+		cc <- conn
+		break
 	}
-	fmt.Printf("[+] [%s]->[%s] connected to target.\n", conn.LocalAddr().String(), targetAddress)
-	return conn, err
 }
 
-func accept(ln net.Listener) (net.Conn, error) {
-	conn, err := ln.Accept()
-	if err != nil {
-		fmt.Printf("[x] accept error [%s].\n", err.Error())
-		return nil, err
+func accept(cc chan net.Conn, ln net.Listener) {
+	for true {
+		fmt.Printf("[#] 4dnat waiting for client to connect [%s %s]\n", ln.Addr().Network(), ln.Addr().String())
+		c, err := ln.Accept()
+		if err != nil {
+			fmt.Printf("[x] accept error [%s].\n", err.Error())
+			continue
+		}
+		fmt.Printf("[+] [%s]<-[%s] new client connected.\n", c.LocalAddr().String(), c.RemoteAddr().String())
+		cc <- c
+		break
 	}
-	fmt.Printf("[+] [%s]<-[%s] new client connected.\n", conn.LocalAddr().String(), conn.RemoteAddr().String())
-	return conn, nil
 }
 
 func listen(listenPort string) net.Listener {
-	ln, err := net.Listen("tcp", ":"+listenPort)
+	if !strings.Contains(listenPort, ":") {
+		listenPort = "0.0.0.0:" + listenPort
+	}
+	ln, err := net.Listen("tcp", listenPort)
 	if err != nil {
 		fmt.Printf("[x] listen error [%s].\n", err.Error())
 		os.Exit(0)