package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"strings"

	"ben.soroos.net/tsnet-proxy/netforward"
	"tailscale.com/tsnet"
)

var hostname = os.Args[1]


type Dialer struct{
	proto string;
	addr string;
}

func (dialer Dialer) Dial() (net.Conn, error) {
	return net.Dial(dialer.proto, dialer.addr)
}

func Forward(server *tsnet.Server, proto, port, dst string, finish chan error) {
	var ln net.Listener
	var err error
	if proto == "tls" {
		proto = "tcp"
		ln, err = server.ListenTLS(proto, fmt.Sprint(":", port))
	} else {
		ln, err = server.Listen(proto, fmt.Sprint(":", port))
	}
	if err != nil {
		finish <- err
	}
	defer ln.Close();
	for true {
		err = netforward.Forward(Dialer { proto: proto, addr: dst }, ln);
		log.Print(err);
	}
	finish <- err;
}

func main() {
	if !(hostname != "") {
		fmt.Println("Usage: tsnet-proxy HOSTNAME FORWARDS...")
		os.Exit(1)
	}

	dir, valid := os.LookupEnv("STATE_DIRECTORY")
	if !valid {
		log.Fatal("STATE_DIRECTORY must be set")
	}
	s := &tsnet.Server{
		Hostname: hostname,
		Dir: dir,
	}
	defer s.Close()

	finish := make(chan error)
	for _, arg := range os.Args[2:] {
		args := strings.SplitN(arg, ":", 3)
		proto, port, dst := args[0], args[1], args[2]
		go Forward(s, proto, port, dst, finish);
	}
	<- finish;
}
