proxy-pool/internal/gateway/transport/transport.go

package transport

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httptrace"
	"net/url"
	"sync"
	"sync/atomic"
	"time"

	proxyDomain "proxy-pool/internal/domain/proxy"
)

const (
	defaultDialTimeout            = 10 * time.Second
	defaultHandshakeTimeout       = 15 * time.Second
	defaultResponseHeaderTimeout  = 30 * time.Second
	defaultIdleConnTimeout        = 90 * time.Second
	defaultMaxErrorResponseBytes  = int64(64 << 10)
	defaultMaxResponseHeaderBytes = int64(64 << 10)
)

var (
	ErrUnsupportedProxyScheme = errors.New("unsupported upstream proxy scheme")
	ErrProxyResponseTooLarge  = errors.New("upstream proxy response headers exceed the configured limit")
)

type Config struct {
	DialTimeout            time.Duration
	HandshakeTimeout       time.Duration
	ResponseHeaderTimeout  time.Duration
	IdleConnTimeout        time.Duration
	MaxIdleConns           int
	MaxIdleConnsPerHost    int
	MaxErrorResponseBytes  int64
	MaxResponseHeaderBytes int64
	TunnelBufferBytes      int
	TunnelIdleTimeout      time.Duration
	TLSClientConfig        *tls.Config
}

type Credentials struct {
	Username string
	Password string
}

func (Credentials) Format(state fmt.State, _ rune) {
	_, _ = state.Write([]byte("transport.Credentials{Username:<redacted>, Password:<redacted>}"))
}

type CredentialResolver interface {
	Resolve(context.Context, proxyDomain.Proxy) (Credentials, error)
}

type CredentialResolverFunc func(context.Context, proxyDomain.Proxy) (Credentials, error)

func (resolve CredentialResolverFunc) Resolve(ctx context.Context, selected proxyDomain.Proxy) (Credentials, error) {
	return resolve(ctx, selected)
}

type ProxyResponseError struct {
	StatusCode int
	Status     string
	Header     http.Header
	Body       []byte
}

func (err *ProxyResponseError) Error() string {
	return fmt.Sprintf("upstream proxy CONNECT response: %s", err.Status)
}

func (err *ProxyResponseError) Retryable() bool {
	return err.StatusCode == http.StatusRequestTimeout ||
		err.StatusCode == http.StatusTooEarly ||
		err.StatusCode == http.StatusTooManyRequests ||
		(err.StatusCode >= http.StatusInternalServerError && err.StatusCode <= 599)
}

type Transport struct {
	config   Config
	resolver CredentialResolver
	client   *http.Transport
	buffers  sync.Pool
}

func New(config Config, resolver CredentialResolver) *Transport {
	applyDefaults(&config)
	transport := &Transport{config: config, resolver: resolver}
	tlsConfig := config.TLSClientConfig
	if tlsConfig != nil {
		tlsConfig = tlsConfig.Clone()
	}
	transport.client = &http.Transport{
		Proxy: func(request *http.Request) (*url.URL, error) {
			proxyURL, ok := request.Context().Value(proxyURLContextKey{}).(*url.URL)
			if !ok || proxyURL == nil {
				return nil, errors.New("upstream proxy URL is missing from request context")
			}
			return proxyURL, nil
		},
		DialContext:            (&net.Dialer{Timeout: config.DialTimeout, KeepAlive: 30 * time.Second}).DialContext,
		ForceAttemptHTTP2:      true,
		MaxIdleConns:           config.MaxIdleConns,
		MaxIdleConnsPerHost:    config.MaxIdleConnsPerHost,
		IdleConnTimeout:        config.IdleConnTimeout,
		TLSHandshakeTimeout:    config.HandshakeTimeout,
		ResponseHeaderTimeout:  config.ResponseHeaderTimeout,
		TLSClientConfig:        tlsConfig,
		MaxResponseHeaderBytes: config.MaxResponseHeaderBytes,
	}
	transport.buffers.New = func() any {
		return make([]byte, config.TunnelBufferBytes)
	}
	return transport
}

func (transport *Transport) RoundTrip(
	ctx context.Context,
	selected proxyDomain.Proxy,
	request *http.Request,
	commit ...func() error,
) (*http.Response, error) {
	if request == nil {
		return nil, errors.New("round trip through proxy: nil request")
	}
	credentials, err := transport.resolveCredentials(ctx, selected)
	if err != nil {
		return nil, fmt.Errorf("resolve upstream proxy credentials: %w", err)
	}
	proxyURL, err := upstreamURL(selected, credentials)
	if err != nil {
		return nil, err
	}
	var commitOnce sync.Once
	var commitMu sync.Mutex
	var commitErr error
	if len(commit) > 0 && commit[0] != nil {
		trace := &httptrace.ClientTrace{GotConn: func(httptrace.GotConnInfo) {
			commitOnce.Do(func() {
				commitMu.Lock()
				commitErr = commit[0]()
				commitMu.Unlock()
			})
		}}
		ctx = httptrace.WithClientTrace(ctx, trace)
	}
	ctx = context.WithValue(ctx, proxyURLContextKey{}, proxyURL)
	clone := request.Clone(ctx)
	clone.RequestURI = ""
	clone.Header = request.Header.Clone()
	clone.Header.Del("Proxy-Authorization")
	response, err := transport.client.RoundTrip(clone)
	commitMu.Lock()
	deferredCommitErr := commitErr
	commitMu.Unlock()
	if deferredCommitErr != nil {
		if response != nil {
			_ = response.Body.Close()
		}
		return nil, fmt.Errorf("commit upstream proxy allocation: %w", deferredCommitErr)
	}
	return response, err
}

func (transport *Transport) OpenTunnel(
	ctx context.Context,
	selected proxyDomain.Proxy,
	target string,
) (net.Conn, error) {
	credentials, err := transport.resolveCredentials(ctx, selected)
	if err != nil {
		return nil, fmt.Errorf("resolve upstream proxy credentials: %w", err)
	}
	handshakeContext, cancelHandshake := context.WithTimeout(ctx, transport.config.HandshakeTimeout)
	defer cancelHandshake()
	connection, err := transport.dialProxy(handshakeContext, selected)
	if err != nil {
		return nil, err
	}
	succeeded := false
	defer func() {
		if !succeeded {
			_ = connection.Close()
		}
	}()

	deadline := time.Now().Add(transport.config.HandshakeTimeout)
	if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) {
		deadline = contextDeadline
	}
	if err := connection.SetDeadline(deadline); err != nil {
		return nil, fmt.Errorf("set upstream proxy handshake deadline: %w", err)
	}
	stopCancellation := context.AfterFunc(ctx, func() {
		_ = connection.SetDeadline(time.Now())
	})
	defer stopCancellation()

	request := &http.Request{
		Method: http.MethodConnect,
		URL:    &url.URL{Opaque: target},
		Host:   target,
		Header: make(http.Header),
	}
	if credentials.Username != "" || credentials.Password != "" {
		request.Header.Set("Proxy-Authorization", basicAuth(credentials))
	}
	if err := request.Write(connection); err != nil {
		return nil, contextError(ctx, fmt.Errorf("write upstream CONNECT: %w", err))
	}

	limited := &io.LimitedReader{R: connection, N: transport.config.MaxResponseHeaderBytes + 1}
	reader := bufio.NewReader(limited)
	response, err := http.ReadResponse(reader, request)
	if err != nil {
		if limited.N == 0 {
			return nil, ErrProxyResponseTooLarge
		}
		return nil, contextError(ctx, fmt.Errorf("read upstream CONNECT response: %w", err))
	}
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		body, readErr := io.ReadAll(io.LimitReader(response.Body, transport.config.MaxErrorResponseBytes))
		_ = response.Body.Close()
		if readErr != nil {
			return nil, fmt.Errorf("read upstream CONNECT error response: %w", readErr)
		}
		return nil, &ProxyResponseError{
			StatusCode: response.StatusCode,
			Status:     response.Status,
			Header:     response.Header.Clone(),
			Body:       body,
		}
	}
	_ = response.Body.Close()
	if err := connection.SetDeadline(time.Time{}); err != nil {
		return nil, fmt.Errorf("clear upstream proxy handshake deadline: %w", err)
	}
	buffered := make([]byte, reader.Buffered())
	if _, err := io.ReadFull(reader, buffered); err != nil {
		return nil, fmt.Errorf("preserve buffered CONNECT bytes: %w", err)
	}
	succeeded = true
	return &bufferedConn{Conn: connection, reader: io.MultiReader(bytes.NewReader(buffered), connection)}, nil
}

func (transport *Transport) CloseIdleConnections() {
	if transport == nil {
		return
	}
	transport.client.CloseIdleConnections()
}

func (transport *Transport) Relay(ctx context.Context, left, right net.Conn) error {
	if left == nil || right == nil {
		return errors.New("relay tunnel: nil connection")
	}
	stopCancellation := context.AfterFunc(ctx, func() {
		deadline := time.Now()
		_ = left.SetDeadline(deadline)
		_ = right.SetDeadline(deadline)
	})
	defer stopCancellation()

	activity := newTunnelActivity()
	idleStopped := make(chan struct{})
	go transport.enforceTunnelIdle(ctx, idleStopped, left, right, activity)
	defer close(idleStopped)

	activeLeft := &activityConn{Conn: left, activity: activity}
	activeRight := &activityConn{Conn: right, activity: activity}
	errorsByDirection := make(chan error, 2)
	go func() { errorsByDirection <- transport.copyTunnel(activeRight, activeLeft) }()
	go func() { errorsByDirection <- transport.copyTunnel(activeLeft, activeRight) }()

	first := <-errorsByDirection
	if first != nil {
		_ = left.Close()
		_ = right.Close()
	}
	second := <-errorsByDirection
	if err := ctx.Err(); err != nil {
		return err
	}
	return errors.Join(first, second)
}

func (transport *Transport) enforceTunnelIdle(
	ctx context.Context,
	stopped <-chan struct{},
	left, right net.Conn,
	activity *tunnelActivity,
) {
	timer := time.NewTimer(transport.config.TunnelIdleTimeout)
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-stopped:
			return
		case <-timer.C:
			remaining := activity.remaining(transport.config.TunnelIdleTimeout)
			if remaining > 0 {
				timer.Reset(remaining)
				continue
			}
			deadline := time.Now()
			_ = left.SetDeadline(deadline)
			_ = right.SetDeadline(deadline)
			return
		}
	}
}

func (transport *Transport) copyTunnel(destination, source net.Conn) error {
	buffer := transport.buffers.Get().([]byte)
	defer transport.buffers.Put(buffer)
	_, err := io.CopyBuffer(destination, source, buffer)
	if halfCloser, ok := destination.(interface{ CloseWrite() error }); ok {
		_ = halfCloser.CloseWrite()
	}
	if halfCloser, ok := source.(interface{ CloseRead() error }); ok {
		_ = halfCloser.CloseRead()
	}
	if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) {
		return nil
	}
	return err
}

func (transport *Transport) dialProxy(ctx context.Context, selected proxyDomain.Proxy) (net.Conn, error) {
	if selected.Scheme != proxyDomain.SchemeHTTP && selected.Scheme != proxyDomain.SchemeHTTPS {
		return nil, fmt.Errorf("%w: %s", ErrUnsupportedProxyScheme, selected.Scheme)
	}
	dialer := &net.Dialer{Timeout: transport.config.DialTimeout, KeepAlive: 30 * time.Second}
	connection, err := dialer.DialContext(ctx, "tcp", selected.Address())
	if err != nil {
		return nil, contextError(ctx, fmt.Errorf("dial upstream proxy %s: %w", selected.ID, err))
	}
	if selected.Scheme == proxyDomain.SchemeHTTP {
		return connection, nil
	}

	tlsConfig := &tls.Config{MinVersion: tls.VersionTLS12, ServerName: selected.Host}
	if transport.config.TLSClientConfig != nil {
		tlsConfig = transport.config.TLSClientConfig.Clone()
		if tlsConfig.ServerName == "" {
			tlsConfig.ServerName = selected.Host
		}
	}
	tlsConnection := tls.Client(connection, tlsConfig)
	if err := tlsConnection.HandshakeContext(ctx); err != nil {
		_ = connection.Close()
		return nil, contextError(ctx, fmt.Errorf("TLS handshake with upstream proxy %s: %w", selected.ID, err))
	}
	return tlsConnection, nil
}

func (transport *Transport) resolveCredentials(
	ctx context.Context,
	selected proxyDomain.Proxy,
) (Credentials, error) {
	if transport.resolver == nil {
		return Credentials{Username: selected.Username}, nil
	}
	credentials, err := transport.resolver.Resolve(ctx, selected)
	if credentials.Username == "" {
		credentials.Username = selected.Username
	}
	return credentials, err
}

func upstreamURL(selected proxyDomain.Proxy, credentials Credentials) (*url.URL, error) {
	if selected.Scheme != proxyDomain.SchemeHTTP && selected.Scheme != proxyDomain.SchemeHTTPS {
		return nil, fmt.Errorf("%w: %s", ErrUnsupportedProxyScheme, selected.Scheme)
	}
	proxyURL := &url.URL{Scheme: string(selected.Scheme), Host: selected.Address()}
	if credentials.Username != "" || credentials.Password != "" {
		proxyURL.User = url.UserPassword(credentials.Username, credentials.Password)
	}
	return proxyURL, nil
}

func applyDefaults(config *Config) {
	if config.DialTimeout <= 0 {
		config.DialTimeout = defaultDialTimeout
	}
	if config.HandshakeTimeout <= 0 {
		config.HandshakeTimeout = defaultHandshakeTimeout
	}
	if config.ResponseHeaderTimeout <= 0 {
		config.ResponseHeaderTimeout = defaultResponseHeaderTimeout
	}
	if config.IdleConnTimeout <= 0 {
		config.IdleConnTimeout = defaultIdleConnTimeout
	}
	if config.MaxIdleConns <= 0 {
		config.MaxIdleConns = 1024
	}
	if config.MaxIdleConnsPerHost <= 0 {
		config.MaxIdleConnsPerHost = 64
	}
	if config.MaxErrorResponseBytes <= 0 {
		config.MaxErrorResponseBytes = defaultMaxErrorResponseBytes
	}
	if config.MaxResponseHeaderBytes <= 0 {
		config.MaxResponseHeaderBytes = defaultMaxResponseHeaderBytes
	}
	if config.TunnelBufferBytes <= 0 {
		config.TunnelBufferBytes = 32 << 10
	}
	if config.TunnelIdleTimeout <= 0 {
		config.TunnelIdleTimeout = 5 * time.Minute
	}
}

func basicAuth(credentials Credentials) string {
	request := &http.Request{Header: make(http.Header)}
	request.SetBasicAuth(credentials.Username, credentials.Password)
	return request.Header.Get("Authorization")
}

func contextError(ctx context.Context, fallback error) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	var networkError net.Error
	if deadline, ok := ctx.Deadline(); ok && !time.Now().Before(deadline) &&
		errors.As(fallback, &networkError) && networkError.Timeout() {
		return context.DeadlineExceeded
	}
	return fallback
}

type bufferedConn struct {
	net.Conn
	reader io.Reader
}

type proxyURLContextKey struct{}

type tunnelActivity struct {
	started     time.Time
	lastElapsed atomic.Int64
}

func newTunnelActivity() *tunnelActivity {
	activity := &tunnelActivity{started: time.Now()}
	activity.touch()
	return activity
}

func (activity *tunnelActivity) touch() {
	activity.lastElapsed.Store(int64(time.Since(activity.started)))
}

func (activity *tunnelActivity) remaining(timeout time.Duration) time.Duration {
	elapsedSinceActivity := time.Since(activity.started) - time.Duration(activity.lastElapsed.Load())
	return timeout - elapsedSinceActivity
}

type activityConn struct {
	net.Conn
	activity *tunnelActivity
}

func (connection *activityConn) Read(buffer []byte) (int, error) {
	read, err := connection.Conn.Read(buffer)
	if read > 0 {
		connection.activity.touch()
	}
	return read, err
}

func (connection *activityConn) Write(buffer []byte) (int, error) {
	written, err := connection.Conn.Write(buffer)
	if written > 0 {
		connection.activity.touch()
	}
	return written, err
}

func (connection *activityConn) CloseWrite() error {
	if halfCloser, ok := connection.Conn.(interface{ CloseWrite() error }); ok {
		return halfCloser.CloseWrite()
	}
	return connection.Close()
}

func (connection *activityConn) CloseRead() error {
	if halfCloser, ok := connection.Conn.(interface{ CloseRead() error }); ok {
		return halfCloser.CloseRead()
	}
	return nil
}

func (connection *bufferedConn) Read(buffer []byte) (int, error) {
	return connection.reader.Read(buffer)
}

func (connection *bufferedConn) CloseWrite() error {
	if halfCloser, ok := connection.Conn.(interface{ CloseWrite() error }); ok {
		return halfCloser.CloseWrite()
	}
	return connection.Close()
}

func (connection *bufferedConn) CloseRead() error {
	if halfCloser, ok := connection.Conn.(interface{ CloseRead() error }); ok {
		return halfCloser.CloseRead()
	}
	return nil
}