package main

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

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"
//     }
//   ]
// }
// 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)
		return
	}

	if len(tcpListenConfig.TcpListens) == 0 {
		fmt.Println("[ERROR] Not find any mapping in config")
		return
	}

	for _, tcpListen := range tcpListenConfig.TcpListens {
		go serverTcpListen(tcpListen)
	}
	ExitChan := make(chan bool, 1)
	<-ExitChan
}

func serverTcpListen(tcpListen TcpListen) {
	lis, err := net.Listen("tcp", tcpListen.Listen)
	if err != nil {
		fmt.Println(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)
	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
	dconn.Close()
}

func readConfigFile() (*TcpListenConfig, error) {
	configFile, err := os.Open("tcp_listen_config.json")
	if err != nil {
		return nil, err
	}
	defer configFile.Close()

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

	return &tcpListenConfig, nil
}