From abc5aa684724fdb587c8f98fab2b160ec09e2bdd Mon Sep 17 00:00:00 2001
From: dushixiang <798148596@qq.com>
Date: Thu, 10 Sep 2020 19:07:43 +0800
Subject: [PATCH] =?UTF-8?q?1.=20=E4=BF=AE=E6=94=B9switch=20case=E4=B8=BAgo?=
 =?UTF-8?q?=E9=A3=8E=E6=A0=BC=202.=20=E4=BD=BF=E7=94=A8log=E4=BB=A3?=
 =?UTF-8?q?=E6=9B=BFfmt=203.=20=E8=BD=AC=E5=8F=91=E6=A8=A1=E5=BC=8F?=
 =?UTF-8?q?=E4=B8=8B=EF=BC=8C=E4=BF=AE=E6=94=B9=E4=B8=BA=E6=8E=A5=E6=94=B6?=
 =?UTF-8?q?=E5=88=B0=E8=BF=9E=E6=8E=A5=E4=B9=8B=E5=90=8E=E5=86=8D=E5=8E=BB?=
 =?UTF-8?q?=E8=BF=9E=E6=8E=A5=E7=9B=AE=E6=A0=87=E5=9C=B0=E5=9D=80?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 4dnat.go | 49 ++++++++++++++++++++++++-------------------------
 1 file changed, 24 insertions(+), 25 deletions(-)

diff --git a/4dnat.go b/4dnat.go
index 9de088b..c8413e4 100644
--- a/4dnat.go
+++ b/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,12 +119,12 @@ 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
 	}
@@ -133,13 +132,13 @@ func dial(cc chan net.Conn, targetAddress string) {
 
 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