navi/internal/protocols/awg/proxy.go

package awg

import (
	"bufio"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// DialFunc dials TCP through the AWG userspace stack.
type DialFunc func(ctx context.Context, network, address string) (net.Conn, error)

type proxyServer struct {
	mu       sync.Mutex
	socksLn  net.Listener
	httpLn   net.Listener
	dial     DialFunc
	wg       sync.WaitGroup
	closed   bool
	socksAddr string
	httpAddr  string
}

func startProxies(socksHostPort, httpHostPort string, dial DialFunc) (*proxyServer, error) {
	s := &proxyServer{dial: dial}
	if socksHostPort != "" {
		ln, err := net.Listen("tcp", socksHostPort)
		if err != nil {
			return nil, fmt.Errorf("awg socks listen %s: %w", socksHostPort, err)
		}
		s.socksLn = ln
		s.socksAddr = ln.Addr().String()
		s.wg.Add(1)
		go s.serveSOCKS()
	}
	if httpHostPort != "" {
		ln, err := net.Listen("tcp", httpHostPort)
		if err != nil {
			_ = s.Close()
			return nil, fmt.Errorf("awg http listen %s: %w", httpHostPort, err)
		}
		s.httpLn = ln
		s.httpAddr = ln.Addr().String()
		s.wg.Add(1)
		go s.serveHTTP()
	}
	return s, nil
}

func (s *proxyServer) Close() error {
	s.mu.Lock()
	if s.closed {
		s.mu.Unlock()
		return nil
	}
	s.closed = true
	if s.socksLn != nil {
		_ = s.socksLn.Close()
	}
	if s.httpLn != nil {
		_ = s.httpLn.Close()
	}
	s.mu.Unlock()
	s.wg.Wait()
	return nil
}

func (s *proxyServer) serveSOCKS() {
	defer s.wg.Done()
	for {
		c, err := s.socksLn.Accept()
		if err != nil {
			return
		}
		s.wg.Add(1)
		go func(conn net.Conn) {
			defer s.wg.Done()
			_ = handleSOCKS(conn, s.dial)
		}(c)
	}
}

func (s *proxyServer) serveHTTP() {
	defer s.wg.Done()
	for {
		c, err := s.httpLn.Accept()
		if err != nil {
			return
		}
		s.wg.Add(1)
		go func(conn net.Conn) {
			defer s.wg.Done()
			_ = handleHTTPProxy(conn, s.dial)
		}(c)
	}
}

func handleSOCKS(client net.Conn, dial DialFunc) error {
	defer client.Close()
	_ = client.SetDeadline(time.Now().Add(30 * time.Second))
	buf := make([]byte, 258)
	if _, err := io.ReadFull(client, buf[:2]); err != nil {
		return err
	}
	if buf[0] != 0x05 {
		return fmt.Errorf("socks: not v5")
	}
	nmethods := int(buf[1])
	if _, err := io.ReadFull(client, buf[:nmethods]); err != nil {
		return err
	}
	if _, err := client.Write([]byte{0x05, 0x00}); err != nil {
		return err
	}
	if _, err := io.ReadFull(client, buf[:4]); err != nil {
		return err
	}
	if buf[0] != 0x05 || buf[1] != 0x01 {
		_, _ = client.Write([]byte{0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return fmt.Errorf("socks: only CONNECT supported")
	}
	var host string
	switch buf[3] {
	case 0x01: // IPv4
		if _, err := io.ReadFull(client, buf[:4]); err != nil {
			return err
		}
		host = net.IP(buf[:4]).String()
	case 0x03: // domain
		if _, err := io.ReadFull(client, buf[:1]); err != nil {
			return err
		}
		l := int(buf[0])
		if _, err := io.ReadFull(client, buf[:l]); err != nil {
			return err
		}
		host = string(buf[:l])
	case 0x04: // IPv6
		if _, err := io.ReadFull(client, buf[:16]); err != nil {
			return err
		}
		host = net.IP(buf[:16]).String()
	default:
		_, _ = client.Write([]byte{0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return fmt.Errorf("socks: bad atyp")
	}
	if _, err := io.ReadFull(client, buf[:2]); err != nil {
		return err
	}
	port := binary.BigEndian.Uint16(buf[:2])
	_ = client.SetDeadline(time.Time{})

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	remote, err := dial(ctx, "tcp", net.JoinHostPort(host, fmt.Sprintf("%d", port)))
	if err != nil {
		_, _ = client.Write([]byte{0x05, 0x05, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
		return err
	}
	defer remote.Close()
	if _, err := client.Write([]byte{0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0}); err != nil {
		return err
	}
	return relay(client, remote)
}

func handleHTTPProxy(client net.Conn, dial DialFunc) error {
	defer client.Close()
	_ = client.SetDeadline(time.Now().Add(30 * time.Second))
	br := bufio.NewReader(client)
	req, err := http.ReadRequest(br)
	if err != nil {
		return err
	}
	_ = client.SetDeadline(time.Time{})

	if strings.EqualFold(req.Method, http.MethodConnect) {
		host := req.Host
		if !strings.Contains(host, ":") {
			host += ":443"
		}
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		remote, err := dial(ctx, "tcp", host)
		if err != nil {
			_, _ = client.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n"))
			return err
		}
		defer remote.Close()
		if _, err := client.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")); err != nil {
			return err
		}
		return relay(client, remote)
	}

	// Absolute-form HTTP proxy request.
	u := req.URL
	host := u.Host
	if host == "" {
		host = req.Host
	}
	if host == "" {
		_, _ = client.Write([]byte("HTTP/1.1 400 Bad Request\r\n\r\n"))
		return fmt.Errorf("http proxy: no host")
	}
	if !strings.Contains(host, ":") {
		host += ":80"
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	remote, err := dial(ctx, "tcp", host)
	if err != nil {
		_, _ = client.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n"))
		return err
	}
	defer remote.Close()
	req.RequestURI = ""
	if err := req.Write(remote); err != nil {
		return err
	}
	return relay(client, remote)
}

func relay(a, b net.Conn) error {
	errc := make(chan error, 2)
	go func() {
		_, err := io.Copy(a, b)
		errc <- err
	}()
	go func() {
		_, err := io.Copy(b, a)
		errc <- err
	}()
	err := <-errc
	_ = a.Close()
	_ = b.Close()
	<-errc
	return err
}