tss/proxy/proxy.go

package proxy

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"net/url"

	"golang.org/x/net/proxy"
)

func GetTransport(proxy string) (http.RoundTripper, error) {
	if proxy == "" {
		return http.DefaultTransport, nil
	}
	proxyUrl, err := url.Parse(proxy)
	if err != nil {
		return nil, fmt.Errorf("failed to parse proxyUrl, url:%s, err: %w", proxy, err)
	}
	if proxyUrl.Host == "unix" && proxyUrl.Scheme == "socks5" &&
		len(proxyUrl.Path) > 1 /* Path cannot be empty or just / */ {
		return unixSocks5Proxy(proxyUrl)
	}
	return proxyTcp(proxyUrl)
}

type forwardDialer struct {
	dialContext func(ctx context.Context, network, address string) (net.Conn, error)
}

func (d *forwardDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	return d.dialContext(ctx, network, address)
}
func (d *forwardDialer) Dial(network, address string) (net.Conn, error) {
	panic("Dial should not be called")
	//default Dial is nil
	//return defaultTransport().Dial(network, address)
}

func unixSocks5Proxy(proxyUrl *url.URL) (http.RoundTripper, error) {
	trans := defaultTransport().Clone()
	if trans.DialContext == nil {
		panic("DefaultTransport has nil DialContext")
	}
	var auth *proxy.Auth
	username := proxyUrl.User.Username()
	password, _ := proxyUrl.User.Password()
	// Both username and password should have atleast one char
	if username != "" && password != "" {
		auth = &proxy.Auth{
			User:     username,
			Password: password,
		}
	}
	dialer, err := proxy.SOCKS5("unix", proxyUrl.Path, auth, &forwardDialer{trans.DialContext})
	if err != nil {
		return nil, fmt.Errorf("failed to make socks proxy, url: %s, err: %w", proxyUrl, err)
	}
	ctxDialer, ok := dialer.(proxy.ContextDialer)
	if !ok {
		panic("proxy.SOCKS5 did not return a ContextDialer")
	}
	trans.DialContext = ctxDialer.DialContext
	trans.Proxy = nil
	return trans, nil
}

func proxyTcp(proxyUrl *url.URL) (http.RoundTripper, error) {
	trans := defaultTransport().Clone()
	trans.Proxy = http.ProxyURL(proxyUrl)
	return trans, nil
}

func defaultTransport() *http.Transport {
	transPtr, ok := http.DefaultTransport.(*http.Transport)
	if !ok {
		panic("http.DefaultTransport is not a *http.Transport")
	}
	return transPtr
}