package main

import (
	"context"
	"flag"
	"io"
	"log"
	"net"
	"sync"
	"time"

	"git.maze.io/go/dpi/protocol"
)

func main() {
	listenFlag := flag.String("listen", "localhost:4080", "proxy listen address")
	targetFlag := flag.String("target", "localhost:22", "proxy target address")
	flag.Parse()

	l, err := net.Listen("tcp", *listenFlag)
	if err != nil {
		log.Fatalf("listen error: %v", err)
	}

	log.Printf("listening on %s", l.Addr())
	for {
		c, err := l.Accept()
		if err != nil {
			log.Fatalln("accept error:", err)
		}
		go proxy(c, *targetFlag)
	}
}

func proxy(client net.Conn, target string) {
	log.Printf("new connection from %s", client.RemoteAddr())

	// Hangup client if we return
	defer func() {
		log.Printf("closing connection to %s: %v", client.RemoteAddr(), client.Close())
	}()

	log.Printf("dialing %s", target)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	var dialer net.Dialer

	server, err := dialer.DialContext(ctx, "tcp", target)
	if err != nil {
		cancel()
		log.Printf("error connecting to %s: %v", target, err)
		return
	}
	cancel()

	// Hangup server if we return
	defer func() {
		log.Printf("closing connection to %s: %v", server.RemoteAddr(), server.Close())
	}()

	// Setup interceptor and wrap the client and server connections.
	interceptor := protocol.NewInterceptor()
	client = interceptor.Client(client)
	server = interceptor.Server(server)

	// Request a return channel and start the detection before doing anything
	// else with the client and server connections.
	intercepted := interceptor.Detect(10 * time.Second)

	log.Printf("client %s connected to %s; proxying", client.RemoteAddr(), server.RemoteAddr())

	// Create a wait group and copy between client and server bidirectionally,
	// either side needs to generate data for the detection to work.
	var group sync.WaitGroup
	group.Go(func() { _, _ = io.Copy(client, server) })
	group.Go(func() { _, _ = io.Copy(server, client) })

	// Wait until the interceptor produces data.
	result := <-intercepted
	if result.Error != nil {
		log.Printf("protocol detection failed: %v", result.Error)
	} else {
		log.Printf("detected protocol %s version %s initiated by %s",
			result.Protocol.Name, result.Protocol.Version, result.Direction)
	}

	// Wait for the multiplexing to finish.
	group.Wait()
}