package main

import (
	"context"
	"flag"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"strings"
	"time"
)

var verbose bool
var dnsServer string
var localIp string

type ReverseProxy struct{}

func (h *ReverseProxy) ProxyHTTPS(rw http.ResponseWriter, req *http.Request) {
	hij, ok := rw.(http.Hijacker)
	if !ok {
		log.Printf("http server does not support hijacker\n")
		return
	}

	clientConn, _, err := hij.Hijack()
	if err != nil {
		log.Printf("http: proxy error: %v\n", err)
		return
	}

	d := net.Dialer{}
	if len(localIp) > 0 {
		d.LocalAddr = &net.TCPAddr{
			IP:   net.ParseIP(localIp),
			Port: 0,
		}
	}

	if len(dnsServer) > 0 {
		r := &net.Resolver{
			PreferGo: true,
			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
				d := net.Dialer{
					Timeout: time.Millisecond * time.Duration(10000),
				}
				return d.DialContext(ctx, "udp", dnsServer)
			},
		}
		d.Resolver = r
	}
	proxyConn, err := d.Dial("tcp", req.URL.Host)
	if err != nil {
		log.Printf("http: proxy error: %v\n", err)
		return
	}

	deadline := time.Now()
	deadline = deadline.Add(time.Second * 30)

	err = clientConn.SetDeadline(deadline)
	if err != nil {
		log.Printf("http: proxy error: %v\n", err)
		return
	}
	err = proxyConn.SetDeadline(deadline)
	if err != nil {
		log.Printf("http: proxy error: %v\n", err)
		return
	}

	_, err = clientConn.Write([]byte("HTTP/1.0 200 OK\r\n\r\n"))
	if err != nil {
		log.Printf("http: proxy error: %v\n", err)
		return
	}

	go func() {
		io.Copy(clientConn, proxyConn)
		clientConn.Close()
		proxyConn.Close()
	}()

	io.Copy(proxyConn, clientConn)
	proxyConn.Close()
	clientConn.Close()
}

func (h *ReverseProxy) ProxyHTTP(w http.ResponseWriter, r *http.Request) {
	director := func(req *http.Request) {
	}
	proxy := &httputil.ReverseProxy{Director: director}

	if len(dnsServer) > 0 || len(localIp) > 0 {
		var r *net.Resolver
		var addr *net.TCPAddr

		if len(dnsServer) > 0 {
			r = &net.Resolver{
				PreferGo: true,
				Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
					d := net.Dialer{
						Timeout: time.Millisecond * time.Duration(10000),
					}
					return d.DialContext(ctx, "udp", dnsServer)
				},
			}
		}

		if len(localIp) > 0 {
			addr = &net.TCPAddr{
				IP:   net.ParseIP(localIp),
				Port: 0,
			}
		}

		t := &http.Transport{
			Proxy: http.ProxyFromEnvironment,
			DialContext: (&net.Dialer{
				Timeout:   30 * time.Second,
				KeepAlive: 30 * time.Second,
				DualStack: true,
				Resolver:  r,
				LocalAddr: addr,
			}).DialContext,
			ForceAttemptHTTP2:     true,
			MaxIdleConns:          100,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   10 * time.Second,
			ExpectContinueTimeout: 1 * time.Second,
		}
		proxy.Transport = t
	}

	proxy.ServeHTTP(w, r)
}

func (h *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if verbose {
		log.Println("", r.URL)
	}

	if r.Method == "CONNECT" {
		h.ProxyHTTPS(w, r)
	} else {
		h.ProxyHTTP(w, r)
	}
}

func main() {
	var bindAddr string
	flag.BoolVar(&verbose, "v", false, "print url")
	flag.StringVar(&bindAddr, "b", "127.0.0.1:1080", "bind address")
	flag.StringVar(&dnsServer, "d", "", "dns server to use(1.1.1.1:53)")
	flag.StringVar(&localIp, "l", "", "local ip address for http request")
	flag.Parse()

	if len(dnsServer) > 0 {
		if strings.Index(dnsServer, ":") < 0 {
			dnsServer += ":53"
		}
	}
	var h ReverseProxy
	log.Fatal(http.ListenAndServe(bindAddr, &h))
}