/**
* simple tcp proxy with timeout failover
* thanks to:
* https://github.com/xtaci/kcptun
* https://github.com/jpillora/go-tcp-proxy
*/
package main

import (
	"flag"
	"net"
	"io"
	"log"
	"time"
)

type addresses []string

func (i *addresses) String() string {
	return ""
}

func (i *addresses) Set(value string) error {
    *i = append(*i, value)
    return nil
}

var localAddr *string = flag.String("l", ":8989", "local address")
var remoteConnectTimeout *int = flag.Int("t", 2, "remoe connect timeout")
var remoteAddresses addresses


func handleIoCopy(p1, p2 io.ReadWriteCloser){
	log.Println("stream opened")
	defer log.Println("stream closed")
	defer p1.Close()
	defer p2.Close()

	// start tunnel
	p1die := make(chan struct{})
	go func() { io.Copy(p1, p2); close(p1die) }()

	p2die := make(chan struct{})
	go func() { io.Copy(p2, p1); close(p2die) }()

	// wait for tunnel termination
	select {
	case <-p1die:
	case <-p2die:
	}
}


func findAvaliableServer() (*net.TCPConn) {
	for _,v := range remoteAddresses {
		rConn, err := net.DialTimeout("tcp", v,time.Duration(*remoteConnectTimeout)*time.Second)
		if err != nil {
			continue
		} else {
			log.Println("connect to server ", v)
			return rConn.(*net.TCPConn)
		}
	}
	return nil
}

func handle(conn *net.TCPConn){
	connRemote := findAvaliableServer()
	if connRemote == nil {
		log.Println("all server connect failed")
	} else {
		go handleIoCopy(conn, connRemote)
	}
}

func main() {
	flag.Var(&remoteAddresses, "r", "remote servers, and specify multiple times")
	flag.Parse()

	log.Println(remoteAddresses)

	addr, err := net.ResolveTCPAddr("tcp", *localAddr)
	if err != nil {
		panic(err)
	}
	log.Println(addr)
	listener, err := net.ListenTCP("tcp", addr)
	if err != nil {
		panic(err)
	}
	
	for {
		conn, err := listener.AcceptTCP()
		if err != nil {
			log.Println("accept error", err)
		}else{
			log.Println("accept connect")
		}
		go handle(conn)	
	}
}