package main

import (
	"go.megvii-inc.com/transparent-proxy/server"
	"fmt"
	"net"
	"io"
)

func main()  {
	proxy_listener, err := server.NewIPv4TransparentListener(":9091")
	if err != nil {
		panic(err)
	}
	for {
		conn, err := proxy_listener.Accept()
		if err != nil {
			fmt.Println("Accept Error: ", err.Error())
			continue
		}
		fmt.Printf("Accepted connection: %s => %s\n", conn.TCPConn().RemoteAddr().String(), conn.RealAddr().String())
		go handle(conn)
	}
}

func pipeThenClose(reader io.Reader, writer io.WriteCloser) {
	defer writer.Close()
	buff := make([]byte, 4096)
	for {
		n, err := reader.Read(buff)
		if n > 0 {
			_, err := writer.Write(buff[:n])
			if err != nil {
				fmt.Printf("Write Error: %#v\n", err.Error())
				return
			}
		}
		if err == io.EOF {
			fmt.Println("Reader reached EOF, closing.")
			return
		}
		if err != nil {
			// Enable this line ONLY FOR DEBUG PURPOSE.
			// fmt.Printf("Read Error: %#v\n", err.Error())
			return
		}
	}
}

func handle(conn server.TransparentConnection) {
	rconn, err := net.DialTCP("tcp4", nil, conn.RealAddr())
	if err != nil {
		fmt.Println("Connect to ", conn.RealAddr().String(), " error: ", err.Error())
		conn.TCPConn().Close()
		return
	}
	go pipeThenClose(conn.TCPConn(), rconn)
	go pipeThenClose(rconn, conn.TCPConn())
}