1. 修改switch case为go风格

2. 使用log代替fmt
3. 转发模式下，修改为接收到连接之后再去连接目标地址

4dnat.go

@@ -1,8 +1,8 @@
 package main
 
 import (
-	"fmt"
 	"io"
+	"log"
 	"net"
 	"os"
 	"strings"
@@ -12,28 +12,26 @@ import (
 
 const RetryInterval int = 5
 
+func init() {
+	log.SetPrefix("[4dnat] ")
+}
+
 func main() {
 	printBanner()
 	args := os.Args
 
 	if len(args) < 4 {
 		printHelp()
-		os.Exit(-1)
+		os.Exit(0)
 	}
 
 	switch args[1] {
-	case "-listen":
-	case "-l":
+	case "-listen", "-l":
 		listener(args[2], args[3])
-		break
-	case "-forward":
-	case "-f":
+	case "-forward", "-f":
 		forward(args[2], args[3])
-		break
-	case "-agent":
-	case "-a":
+	case "-agent", "-a":
 		agent(args[2], args[3])
-		break
 	default:
 		printHelp()
 	}
@@ -59,16 +57,16 @@ func printBanner() {
 
 func copyIO(src, dest net.Conn, wg *sync.WaitGroup) {
 	defer src.Close()
-	fmt.Printf("[#] [%s]->[%s] ==> [%s]->[%s]\n", src.RemoteAddr().String(), src.LocalAddr().String(), dest.LocalAddr().String(), dest.RemoteAddr().String())
+	log.Printf("[#] [%s->%s] ==> [%s->%s]\n", src.RemoteAddr().String(), src.LocalAddr().String(), dest.LocalAddr().String(), dest.RemoteAddr().String())
 	_, _ = io.Copy(src, dest)
-	fmt.Printf("[-] [%s]->[%s] closed.\n", src.RemoteAddr().String(), src.LocalAddr().String())
+	log.Printf("[-] [%s->%s] closed.\n", src.RemoteAddr().String(), src.LocalAddr().String())
 	wg.Done()
 }
 
 func listener(listenPort0, listenPort1 string) {
 	ln0 := listen(listenPort0)
 	ln1 := listen(listenPort1)
-	fmt.Printf("[#] 4dnat listen port on: [%s %s],[%s %s]\n", ln0.Addr().Network(), ln0.Addr().String(), ln1.Addr().Network(), ln1.Addr().String())
+	log.Printf("[#] 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)
@@ -83,20 +81,21 @@ func listener(listenPort0, listenPort1 string) {
 
 func forward(listenPort string, targetAddress string) {
 	ln := listen(listenPort)
-	fmt.Printf("[#] 4dnat listen on: [%s %s] forward to: [%s]\n", ln.Addr().Network(), ln.Addr().String(), targetAddress)
+	log.Printf("[#] 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 := <-cc
 
-		conn0, conn1 := <-cc, <-cc
+		go dial(cc, targetAddress)
+		conn1 := <-cc
 		go mutualCopyIO(conn0, conn1)
 	}
 }
 
 func agent(targetAddress0 string, targetAddress1 string) {
-	fmt.Printf("[#] 4dnat agent with: [%s] [%s]\n", targetAddress0, targetAddress1)
+	log.Printf("[#] agent with: [%s %s]\n", targetAddress0, targetAddress1)
 	for {
 		cc := make(chan net.Conn, 2)
 
@@ -120,26 +119,26 @@ 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)
+			log.Printf("[x] connect [%s] error [%s].\n", targetAddress, err.Error())
+			log.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)
+		log.Printf("[+] [%s->%s] connected to target.\n", conn.LocalAddr().String(), targetAddress)
 		cc <- conn
 		break
 	}
 }
 
 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())
+		log.Printf("[#] 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())
+			log.Printf("[x] accept error [%s].\n", err.Error())
 			continue
 		}
-		fmt.Printf("[+] [%s]<-[%s] new client connected.\n", c.LocalAddr().String(), c.RemoteAddr().String())
+		log.Printf("[+] [%s<-%s] new client connected.\n", c.LocalAddr().String(), c.RemoteAddr().String())
 		cc <- c
 		break
 	}
@@ -151,7 +150,7 @@ func listen(listenPort string) net.Listener {
 	}
 	ln, err := net.Listen("tcp", listenPort)
 	if err != nil {
-		fmt.Printf("[x] listen error [%s].\n", err.Error())
+		log.Printf("[x] listen error [%s].\n", err.Error())
 		os.Exit(0)
 	}
 	return ln