This commit is contained in:
2020-02-23 13:37:29 +08:00
parent 81be6741b4
commit 8c4fbf0bfa

View File

@@ -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
}