package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"strings"
)

const ConfigFileName = "tcp_listen_config.json"

type ConnErrorType int32

const (
	ConnClosed       ConnErrorType = 0
	ConnReset        ConnErrorType = 1
	ConnUnknownError ConnErrorType = 99999
)

type TcpListen struct {
	Listen  string `json:"listen"`
	Backend string `json:"backend"`
}

type TcpListenConfig struct {
	TcpListens []TcpListen `json:"tcpListens"`
}

// {
//   "tcpListens": [
//     {
//       "listen": ":8080",
//       "backend": "11.22.33.44:8080"
//     }
//   ]
// }
// 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.Printf("[ERROR] Not find any mapping in config\n")
		printHelpMessage()
		return
	}

	for _, tcpListen := range tcpListenConfig.TcpListens {
		go serverTcpListen(tcpListen)
	}
	exitChan := make(chan bool, 1)
	<-exitChan // never end, wait kill or Ctrl+C
}

func serverTcpListen(tcpListen TcpListen) {
	lis, err := net.Listen("tcp", tcpListen.Listen)
	if err != nil {
		fmt.Printf("[ERROR] Listen on port %s error: %s\n", tcpListen.Listen, err)
		return
	}
	defer lis.Close()
	fmt.Printf("[INFO] Listen TCP at: %v  --> %v\n", tcpListen.Listen, tcpListen.Backend)
	for {
		conn, err := lis.Accept()
		if err != nil {
			fmt.Printf("[ERROR] Listen error: %v\n", err)
			continue
		}
		fmt.Printf("[INFO] TCP connection from %v, local: %v\n", conn.RemoteAddr(), conn.LocalAddr())
		go handleTcpRequest(tcpListen, conn)
	}
}

func handleTcpRequest(tcpListen TcpListen, sconn net.Conn) {
	defer sconn.Close()
	ip := tcpListen.Backend
	dconn, err := net.Dial("tcp", ip)
	if err != nil {
		fmt.Printf("[ERROR] Connect to %v failed:%v\n", ip, err)
		return
	}
	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(dconn net.Conn, sconn net.Conn, ip string, exitChan chan bool) {
	if _, err := io.Copy(dconn, sconn); 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 [%s]:\n%s\n", ConfigFileName, string(out.Bytes()))
		}
	}
}

func readConfigFile() (*TcpListenConfig, error) {
	fmt.Printf("[INFO] Load config: %s\n", ConfigFileName)
	configFile, err := os.Open(ConfigFileName)
	if err != nil {
		return nil, err
	}
	defer configFile.Close()

	configFileBytes, err := ioutil.ReadAll(configFile)
	if err != nil {
		return nil, err
	}
	var tcpListenConfig TcpListenConfig
	if err = json.Unmarshal(configFileBytes, &tcpListenConfig); err != nil {
		return nil, err
	}
	return &tcpListenConfig, nil
}