From 8c4fbf0bfa1e64fb657b4b35c0c66b87bf83737a Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sun, 23 Feb 2020 13:37:29 +0800
Subject: [PATCH] reflect

---
 simple-tcp-proxy.go | 100 ++++++++++++++++++++++++++++++--------------
 1 file changed, 68 insertions(+), 32 deletions(-)

diff --git a/simple-tcp-proxy.go b/simple-tcp-proxy.go
index 32f9f3d..127977c 100644
--- a/simple-tcp-proxy.go
+++ b/simple-tcp-proxy.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -10,6 +11,14 @@ import (
 	"strings"
 )
 
+type ConnErrorType int32
+
+const (
+	ConnClosed       ConnErrorType = 0
+	ConnReset        ConnErrorType = 1
+	ConnUnknownError ConnErrorType = 99999
+)
+
 type TcpListen struct {
 	Listen  string `json:"listen"`
 	Backend string `json:"backend"`
@@ -27,25 +36,25 @@ type TcpListenConfig struct {
 //     }
 //   ]
 // }
-// https://blog.csdn.net/fyxichen/article/details/51505542
+// Modified from: https://blog.csdn.net/fyxichen/article/details/51505542
 func main() {
 	tcpListenConfig, err := readConfigFile()
-
 	if err != nil {
 		fmt.Printf("[ERROR] Error in load config file: %v\n", err)
+		printHelpMessage()
 		return
 	}
-
 	if len(tcpListenConfig.TcpListens) == 0 {
 		fmt.Println("[ERROR] Not find any mapping in config")
+		printHelpMessage()
 		return
 	}
 
 	for _, tcpListen := range tcpListenConfig.TcpListens {
 		go serverTcpListen(tcpListen)
 	}
-	ExitChan := make(chan bool, 1)
-	<-ExitChan
+	exitChan := make(chan bool, 1)
+	<-exitChan // never end, wait kill or Ctrl+C
 }
 
 func serverTcpListen(tcpListen TcpListen) {
@@ -75,35 +84,63 @@ func handleTcpRequest(tcpListen TcpListen, sconn net.Conn) {
 		fmt.Printf("[ERROR] Connect to %v failed:%v\n", ip, err)
 		return
 	}
-	ExitChan := make(chan bool, 1)
-	go func(sconn net.Conn, dconn net.Conn, Exit chan bool) {
-		_, err := io.Copy(dconn, sconn)
-		if err != nil {
-			errStr := fmt.Sprintf("%v", err)
-			if strings.Contains(errStr, "use of closed network connection") {
-				fmt.Printf("[INFO] Connection closed: %v\n", ip)
-			} else {
-				fmt.Printf("[WARN] Send data to %v failed:%v\n", ip, err)
-			}
-		}
-		Exit <- true
-	}(sconn, dconn, ExitChan)
-	go func(sconn net.Conn, dconn net.Conn, Exit chan bool) {
-		_, err := io.Copy(sconn, dconn)
-		if err != nil {
-			errStr := fmt.Sprintf("%v", err)
-			if strings.Contains(errStr, "use of closed network connection") {
-				fmt.Printf("[INFO] Connection closed: %v\n", ip)
-			} else {
-				fmt.Printf("[WARN] Receive data from %v failed:%v\n", ip, err)
-			}
-		}
-		Exit <- true
-	}(sconn, dconn, ExitChan)
-	<-ExitChan
+	exitChan := make(chan bool, 1)
+	twoDirectionCopyConn(sconn, dconn, ip, exitChan)
+
+	<-exitChan
 	dconn.Close()
 }
 
+func twoDirectionCopyConn(conn1 net.Conn, conn2 net.Conn, ip string, exitChan chan bool) {
+	go copyConn(conn1, conn2, ip, exitChan)
+	go copyConn(conn2, conn1, ip, exitChan)
+}
+
+func copyConn(sconn net.Conn, dconn net.Conn, ip string, exitChan chan bool) {
+	_, err := io.Copy(dconn, sconn)
+	if err != nil {
+		printConnError(ip, err)
+	}
+	exitChan <- true
+}
+
+func printConnError(ip string, err error) {
+	connError := sortConnError(err)
+	if connError == ConnClosed {
+		fmt.Printf("[INFO] Connection closed: %v\n", ip)
+	} else {
+		fmt.Printf("[WARN] Receive data from %v failed:%v\n", ip, err)
+	}
+}
+
+func sortConnError(err error) ConnErrorType {
+	errStr := fmt.Sprintf("%v", err)
+	if strings.Contains(errStr, "use of closed network connection") {
+		return ConnClosed
+	}
+	if strings.Contains(errStr, "connection reset by peer") {
+		return ConnReset
+	}
+	return ConnUnknownError
+}
+
+func printHelpMessage() {
+	tcpListenConfig := TcpListenConfig{
+		TcpListens: []TcpListen{
+			TcpListen{
+				Listen:  ":8080",
+				Backend: "127.0.0.1:9090",
+			},
+		},
+	}
+	if tcpListenConfigJSON, err := json.Marshal(tcpListenConfig); err == nil {
+		var out bytes.Buffer
+		if err := json.Indent(&out, tcpListenConfigJSON, "", "  "); err == nil {
+			fmt.Printf("[INFO] Sample JSON config [tcp_listen_config.json]:\n%s\n", string(out.Bytes()))
+		}
+	}
+}
+
 func readConfigFile() (*TcpListenConfig, error) {
 	configFile, err := os.Open("tcp_listen_config.json")
 	if err != nil {
@@ -120,6 +157,5 @@ func readConfigFile() (*TcpListenConfig, error) {
 	if err != nil {
 		return nil, err
 	}
-
 	return &tcpListenConfig, nil
 }