starnet/proxy_local_helpers_test.go

package starnet

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
)

type connectProxyServer struct {
	*httptest.Server
	mu      sync.Mutex
	targets []string
}

func newIPv4Server(t testing.TB, handler http.Handler) *httptest.Server {
	t.Helper()

	listener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen tcp4: %v", err)
	}

	server := httptest.NewUnstartedServer(handler)
	server.Listener = listener
	server.Start()
	return server
}

func newIPv4TLSServer(t testing.TB, handler http.Handler) *httptest.Server {
	t.Helper()

	listener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen tcp4: %v", err)
	}

	server := httptest.NewUnstartedServer(handler)
	server.Listener = listener
	server.StartTLS()
	return server
}

func newTrustedIPv4TLSServer(t testing.TB, dnsName string, handler http.Handler) (*httptest.Server, *x509.CertPool) {
	t.Helper()

	testT, ok := t.(*testing.T)
	if !ok {
		t.Fatal("newTrustedIPv4TLSServer requires *testing.T")
	}

	certPEM, keyPEM := genSelfSignedCertPEM(testT, dnsName)
	cert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		t.Fatalf("X509KeyPair: %v", err)
	}

	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(certPEM) {
		t.Fatal("AppendCertsFromPEM returned false")
	}

	server := httptest.NewUnstartedServer(handler)
	listener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen tcp4: %v", err)
	}
	server.Listener = listener
	server.TLS = &tls.Config{
		Certificates: []tls.Certificate{cert},
	}
	server.StartTLS()
	return server, pool
}

func httpsURLForHost(t testing.TB, server *httptest.Server, host string) string {
	t.Helper()

	_, port, err := net.SplitHostPort(server.Listener.Addr().String())
	if err != nil {
		t.Fatalf("split host port: %v", err)
	}
	return fmt.Sprintf("https://%s:%s", host, port)
}

func newIPv4ConnectProxyServer(t testing.TB, dialTarget func(target string) (net.Conn, error)) *connectProxyServer {
	t.Helper()

	proxy := &connectProxyServer{}
	if dialTarget == nil {
		dialTarget = func(target string) (net.Conn, error) {
			return net.Dial("tcp", target)
		}
	}

	proxy.Server = newIPv4Server(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodConnect {
			http.Error(w, "connect required", http.StatusMethodNotAllowed)
			return
		}

		proxy.mu.Lock()
		proxy.targets = append(proxy.targets, r.Host)
		proxy.mu.Unlock()

		targetConn, err := dialTarget(r.Host)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadGateway)
			return
		}

		hijacker, ok := w.(http.Hijacker)
		if !ok {
			targetConn.Close()
			t.Fatal("proxy response writer is not a hijacker")
		}

		clientConn, rw, err := hijacker.Hijack()
		if err != nil {
			targetConn.Close()
			t.Fatalf("hijack proxy conn: %v", err)
		}
		if _, err := rw.WriteString("HTTP/1.1 200 Connection Established\r\n\r\n"); err != nil {
			clientConn.Close()
			targetConn.Close()
			t.Fatalf("write connect response: %v", err)
		}
		if err := rw.Flush(); err != nil {
			clientConn.Close()
			targetConn.Close()
			t.Fatalf("flush connect response: %v", err)
		}

		relayProxyConns(clientConn, targetConn)
	}))
	return proxy
}

func (p *connectProxyServer) Targets() []string {
	p.mu.Lock()
	defer p.mu.Unlock()
	return append([]string(nil), p.targets...)
}

type socks5ProxyServer struct {
	ln      net.Listener
	addr    string
	dial    func(target string) (net.Conn, error)
	stopCh  chan struct{}
	wg      sync.WaitGroup
	mu      sync.Mutex
	targets []string
}

func newSOCKS5ProxyServer(t testing.TB, dialTarget func(target string) (net.Conn, error)) *socks5ProxyServer {
	t.Helper()

	if dialTarget == nil {
		dialTarget = func(target string) (net.Conn, error) {
			return net.Dial("tcp", target)
		}
	}

	ln, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen tcp4 socks5: %v", err)
	}

	proxy := &socks5ProxyServer{
		ln:     ln,
		addr:   ln.Addr().String(),
		dial:   dialTarget,
		stopCh: make(chan struct{}),
	}

	proxy.wg.Add(1)
	go func() {
		defer proxy.wg.Done()
		for {
			conn, err := ln.Accept()
			if err != nil {
				select {
				case <-proxy.stopCh:
					return
				default:
					return
				}
			}

			proxy.wg.Add(1)
			go func(c net.Conn) {
				defer proxy.wg.Done()
				proxy.handleConn(t, c)
			}(conn)
		}
	}()

	return proxy
}

func (p *socks5ProxyServer) URL() string {
	return "socks5://" + p.addr
}

func (p *socks5ProxyServer) Targets() []string {
	p.mu.Lock()
	defer p.mu.Unlock()
	return append([]string(nil), p.targets...)
}

func (p *socks5ProxyServer) Close() {
	close(p.stopCh)
	_ = p.ln.Close()
	p.wg.Wait()
}

func (p *socks5ProxyServer) handleConn(t testing.TB, conn net.Conn) {
	t.Helper()
	closeConn := true
	defer func() {
		if closeConn {
			_ = conn.Close()
		}
	}()

	header := make([]byte, 2)
	if _, err := io.ReadFull(conn, header); err != nil {
		return
	}
	if header[0] != 0x05 {
		return
	}

	methods := make([]byte, int(header[1]))
	if _, err := io.ReadFull(conn, methods); err != nil {
		return
	}
	if _, err := conn.Write([]byte{0x05, 0x00}); err != nil {
		return
	}

	reqHeader := make([]byte, 4)
	if _, err := io.ReadFull(conn, reqHeader); err != nil {
		return
	}
	if reqHeader[0] != 0x05 || reqHeader[1] != 0x01 {
		_, _ = conn.Write([]byte{0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return
	}

	host, err := readSOCKS5Addr(conn, reqHeader[3])
	if err != nil {
		_, _ = conn.Write([]byte{0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return
	}
	portBytes := make([]byte, 2)
	if _, err := io.ReadFull(conn, portBytes); err != nil {
		return
	}
	target := net.JoinHostPort(host, fmt.Sprintf("%d", binary.BigEndian.Uint16(portBytes)))

	p.mu.Lock()
	p.targets = append(p.targets, target)
	p.mu.Unlock()

	targetConn, err := p.dial(target)
	if err != nil {
		_, _ = conn.Write([]byte{0x05, 0x05, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return
	}

	if _, err := conn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0}); err != nil {
		targetConn.Close()
		return
	}

	closeConn = false
	relayProxyConns(conn, targetConn)
}

func readSOCKS5Addr(r io.Reader, atyp byte) (string, error) {
	switch atyp {
	case 0x01:
		buf := make([]byte, 4)
		if _, err := io.ReadFull(r, buf); err != nil {
			return "", err
		}
		return net.IP(buf).String(), nil
	case 0x03:
		var size [1]byte
		if _, err := io.ReadFull(r, size[:]); err != nil {
			return "", err
		}
		buf := make([]byte, int(size[0]))
		if _, err := io.ReadFull(r, buf); err != nil {
			return "", err
		}
		return string(buf), nil
	case 0x04:
		buf := make([]byte, 16)
		if _, err := io.ReadFull(r, buf); err != nil {
			return "", err
		}
		return net.IP(buf).String(), nil
	default:
		return "", fmt.Errorf("unsupported atyp: %d", atyp)
	}
}

func relayProxyConns(left, right net.Conn) {
	var once sync.Once
	closeBoth := func() {
		_ = left.Close()
		_ = right.Close()
	}

	go func() {
		_, _ = io.Copy(left, right)
		once.Do(closeBoth)
	}()
	go func() {
		_, _ = io.Copy(right, left)
		once.Do(closeBoth)
	}()
}