proxy_util/handler.go

package main

import (
	"context"
	"fmt"
	"io"
	"math/rand"
	"net"
	"net/http"
	"time"

	log "github.com/fangdingjun/go-log/v5"
)

func getLocalIP() string {
	n := rand.Int() % len(iplist)
	return iplist[n]
}

func proxyHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		handleConnect(w, r)
		return
	}
	if r.Method == http.MethodGet || r.Method == http.MethodPost {
		handleHTTP(w, r)
		return
	}
	http.Error(w, "unsupported method", http.StatusBadRequest)
}

func localDialer(ctx context.Context, network, addr string) (net.Conn, error) {
	ip := getLocalIP()
	laddr := &net.TCPAddr{
		IP:   net.ParseIP(ip),
		Port: 0,
	}

	dialer := &net.Dialer{
		LocalAddr: laddr,
		Timeout:   10 * time.Second,
	}
	return dialer.DialContext(ctx, network, addr)
}

func handleHTTP(w http.ResponseWriter, r *http.Request) {
	r.Header.Del("proxy-connection")
	r.Header.Del("proxy-authorization")
	tr := &http.Transport{
		DialContext: localDialer,
	}

	log.Debugf("%s %s", r.Method, r.RequestURI)

	resp, err := tr.RoundTrip(r)
	if err != nil {
		log.Errorln(err)
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer resp.Body.Close()

	hdr := w.Header()
	resp.Header.Del("connection")

	for k, v := range resp.Header {
		for _, v1 := range v {
			hdr.Add(k, v1)
		}
	}

	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

func mydail(uri string) (net.Conn, error) {
	addr := &net.TCPAddr{}
	dialer := &net.Dialer{
		LocalAddr: addr,
		Timeout:   3 * time.Second,
	}

	for i := 0; i < 3; i++ {
		ip := getLocalIP()
		addr.IP = net.ParseIP(ip)

		log.Debugf("connect to %s, use local ip %s", uri, ip)

		conn, err := dialer.Dial("tcp", uri)
		if err == nil {
			return conn, err
		}
		log.Errorf("connect to %s, %s, retrying", uri, err)
	}
	return nil, fmt.Errorf("connect to %s, max retries reached, give up", uri)
}

func handleConnect(w http.ResponseWriter, r *http.Request) {
	conn, err := mydail(r.RequestURI)
	if err != nil {
		log.Errorln(err)
		http.Error(w, "connection failed", http.StatusBadGateway)
		return
	}

	defer conn.Close()

	hj := w.(http.Hijacker)
	conn1, _, _ := hj.Hijack()
	defer conn1.Close()

	fmt.Fprintf(conn1, "HTTP/1.1 200 connection ok\r\n\r\n")

	ch := make(chan struct{}, 2)
	go func() {
		io.Copy(conn, conn1)
		ch <- struct{}{}
	}()

	go func() {
		io.Copy(conn1, conn)
		ch <- struct{}{}
	}()
	<-ch
}

type handler struct {
}

func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	proxyHandler(w, r)
}