@@ -10,12 +10,14 @@ import (
1010 "net/url"
1111 "os"
1212 "strings"
13+ "sync"
1314 "time"
1415
1516 "bufio"
1617
1718 "github.com/coder/websocket"
1819 "golang.org/x/net/context"
20+ "golang.org/x/net/http/httpproxy"
1921 "golang.org/x/net/proxy"
2022)
2123
@@ -25,6 +27,63 @@ type Dialer interface {
2527 String () string
2628}
2729
30+ // ProxyFunc returns a proxy URL for the given address, or nil for direct connection.
31+ // This follows the same pattern as http.Transport.Proxy.
32+ type ProxyFunc func (addr string ) (* url.URL , error )
33+
34+ var (
35+ envProxyOnce sync.Once
36+ envProxyFunc func (* url.URL ) (* url.URL , error )
37+ )
38+
39+ // TCPProxyFromEnvironment returns the proxy URL for the given address based on
40+ // HTTP_PROXY, HTTPS_PROXY, and NO_PROXY environment variables.
41+ // If the address should be connected to directly, nil is returned.
42+ func TCPProxyFromEnvironment (addr string ) (* url.URL , error ) {
43+ envProxyOnce .Do (func () {
44+ // Use Go's standard httpproxy package
45+ cfg := httpproxy .FromEnvironment ()
46+ envProxyFunc = cfg .ProxyFunc ()
47+ })
48+
49+ // Create a fake URL with the address to check proxy rules
50+ // We use http scheme as default for TCP connections
51+ u := & url.URL {
52+ Scheme : "http" ,
53+ Host : addr ,
54+ }
55+
56+ return envProxyFunc (u )
57+ }
58+
59+ // normalizeProxyURL adds the appropriate scheme to a proxy URL if missing
60+ func normalizeProxyURL (proxyStr string , isSocks bool ) string {
61+ if proxyStr == "" || strings .Contains (proxyStr , "://" ) {
62+ return proxyStr
63+ }
64+
65+ if isSocks {
66+ return "socks5://" + proxyStr
67+ }
68+ return "http://" + proxyStr
69+ }
70+
71+ // parseProxyURL tries to parse proxy URL from environment variables
72+ func parseProxyURL (vars []string ) * url.URL {
73+ for _ , v := range vars {
74+ if val := os .Getenv (v ); val != "" {
75+ // Add default scheme if missing
76+ isSocks := strings .Contains (strings .ToUpper (v ), "SOCKS" )
77+ val = normalizeProxyURL (val , isSocks )
78+
79+ if u , err := url .Parse (val ); err == nil {
80+ return u
81+ }
82+ }
83+ }
84+ return nil
85+ }
86+
2887// DirectDialer connects directly to the target
2988type DirectDialer struct {
3089 addr string
@@ -150,13 +209,7 @@ func NewSOCKS5Dialer(proxyURL *url.URL, addr string, timeout time.Duration) *SOC
150209
151210func (s * SOCKS5Dialer ) Connect () (net.Conn , error ) {
152211 dialer := & net.Dialer {Timeout : s .timeout }
153- auth := & proxy.Auth {}
154- if s .proxyURL .User != nil {
155- auth .User = s .proxyURL .User .Username ()
156- if pass , ok := s .proxyURL .User .Password (); ok {
157- auth .Password = pass
158- }
159- }
212+ auth := extractSOCKS5Auth (s .proxyURL )
160213
161214 socks5 , err := proxy .SOCKS5 ("tcp" , s .proxyURL .Host , auth , dialer )
162215 if err != nil {
@@ -203,19 +256,14 @@ func (s *WebSocketDialer) Connect() (net.Conn, error) {
203256 }
204257
205258 // Configure proxies
206- if s .proxy .Scheme == "http" || s .proxy .Scheme == "https" {
259+ switch s .proxy .Scheme {
260+ case "http" , "https" :
207261 transport .Proxy = func (req * http.Request ) (* url.URL , error ) {
208262 return s .proxy , nil
209263 }
210264
211- } else if s .proxy .Scheme == "socks5" {
212- auth := & proxy.Auth {}
213- if s .proxy .User != nil {
214- auth .User = s .proxy .User .Username ()
215- if pass , ok := s .proxy .User .Password (); ok {
216- auth .Password = pass
217- }
218- }
265+ case "socks5" :
266+ auth := extractSOCKS5Auth (s .proxy )
219267 dialer , err := proxy .SOCKS5 ("tcp" , s .proxy .Host , auth , & net.Dialer {
220268 Timeout : s .timeout ,
221269 })
@@ -227,7 +275,7 @@ func (s *WebSocketDialer) Connect() (net.Conn, error) {
227275 return dialer .(proxy.ContextDialer ).DialContext (ctx , network , addr )
228276 }
229277
230- } else {
278+ default :
231279 return nil , fmt .Errorf ("unsupported proxy scheme: %s" , s .proxy .Scheme )
232280 }
233281
@@ -291,6 +339,21 @@ func generateProxyAuth(user *url.Userinfo) string {
291339 return "Basic " + base64 .StdEncoding .EncodeToString ([]byte (auth ))
292340}
293341
342+ // extractSOCKS5Auth extracts authentication from URL for SOCKS5 proxy
343+ func extractSOCKS5Auth (u * url.URL ) * proxy.Auth {
344+ if u == nil || u .User == nil {
345+ return & proxy.Auth {}
346+ }
347+
348+ auth := & proxy.Auth {
349+ User : u .User .Username (),
350+ }
351+ if pass , ok := u .User .Password (); ok {
352+ auth .Password = pass
353+ }
354+ return auth
355+ }
356+
294357// DialConfig is a configuration for creating dialers
295358type DialConfig struct {
296359 // Configuration options
@@ -306,8 +369,8 @@ func NewDialConfig() *DialConfig {
306369 enableTLS : os .Getenv ("ENABLE_TLS" ) == "true" ,
307370 enableWS : os .Getenv ("ENABLE_WS" ) == "true" ,
308371 enableWSS : os .Getenv ("ENABLE_WSS" ) == "true" ,
309- httpProxy : parseHTTPProxy (),
310- socksProxy : parseSOCKSProxy (),
372+ httpProxy : ParseHTTPProxy (),
373+ socksProxy : ParseSOCKSProxy (),
311374 }
312375}
313376
@@ -348,13 +411,9 @@ func (dc *DialConfig) WithHTTPProxyString(proxyURLStr string) *DialConfig {
348411 return dc
349412 }
350413
351- if ! strings .HasPrefix (proxyURLStr , "http://" ) && ! strings .HasPrefix (proxyURLStr , "https://" ) {
352- proxyURLStr = "http://" + proxyURLStr
353- }
354-
355- proxyURL , err := url .Parse (proxyURLStr )
356- if err == nil {
357- dc .httpProxy = proxyURL
414+ proxyURLStr = normalizeProxyURL (proxyURLStr , false )
415+ if u , err := url .Parse (proxyURLStr ); err == nil {
416+ dc .httpProxy = u
358417 }
359418 return dc
360419}
@@ -366,13 +425,9 @@ func (dc *DialConfig) WithSOCKSProxyString(proxyURLStr string) *DialConfig {
366425 return dc
367426 }
368427
369- if ! strings .HasPrefix (proxyURLStr , "socks5://" ) {
370- proxyURLStr = "socks5://" + proxyURLStr
371- }
372-
373- proxyURL , err := url .Parse (proxyURLStr )
374- if err == nil {
375- dc .socksProxy = proxyURL
428+ proxyURLStr = normalizeProxyURL (proxyURLStr , true )
429+ if u , err := url .Parse (proxyURLStr ); err == nil {
430+ dc .socksProxy = u
376431 }
377432 return dc
378433}
@@ -381,25 +436,27 @@ func (dc *DialConfig) WithSOCKSProxyString(proxyURLStr string) *DialConfig {
381436func (dc * DialConfig ) Dialers (addr string , timeout time.Duration ) []Dialer {
382437 var dialers []Dialer
383438
384- // If HTTP proxy is configured, add proxy strategies first
439+ // Helper to create WebSocket URLs
440+ makeWSURL := func (secure bool ) * url.URL {
441+ scheme := "ws"
442+ if secure {
443+ scheme = "wss"
444+ }
445+ return & url.URL {Scheme : scheme , Host : addr , Path : "/" }
446+ }
447+
448+ // If HTTP proxy is configured, add proxy strategies
385449 if dc .httpProxy != nil {
386450 dialers = append (dialers ,
387451 NewHTTPProxyDialer (dc .httpProxy , addr , timeout ))
388452
389- if dc .enableTLS {
390- dialers = append (dialers ,
391- NewHTTPProxyDialer (dc .httpProxy , addr , timeout ))
392- }
393-
394453 if dc .enableWS {
395- wsURL := & url.URL {Scheme : "ws" , Host : addr , Path : "/" }
396454 dialers = append (dialers ,
397- NewWebSocketDialer (wsURL , dc .httpProxy , false , timeout ))
455+ NewWebSocketDialer (makeWSURL ( false ) , dc .httpProxy , false , timeout ))
398456
399457 if dc .enableWSS {
400- wssURL := & url.URL {Scheme : "wss" , Host : addr , Path : "/" }
401458 dialers = append (dialers ,
402- NewWebSocketDialer (wssURL , dc .httpProxy , false , timeout ))
459+ NewWebSocketDialer (makeWSURL ( true ) , dc .httpProxy , false , timeout ))
403460 }
404461 }
405462 } else if dc .socksProxy != nil {
@@ -419,53 +476,76 @@ func (dc *DialConfig) Dialers(addr string, timeout time.Duration) []Dialer {
419476
420477 // Add WebSocket direct if enabled
421478 if dc .enableWS {
422- wsURL := & url.URL {Scheme : "ws" , Host : addr , Path : "/" }
423479 dialers = append (dialers ,
424- NewWebSocketDialer (wsURL , nil , false , timeout ))
480+ NewWebSocketDialer (makeWSURL ( false ) , nil , false , timeout ))
425481
426482 if dc .enableWSS {
427- wssURL := & url.URL {Scheme : "wss" , Host : addr , Path : "/" }
428483 dialers = append (dialers ,
429- NewWebSocketDialer (wssURL , nil , false , timeout ))
484+ NewWebSocketDialer (makeWSURL ( true ) , nil , false , timeout ))
430485 }
431486 }
432487
433488 return dialers
434489}
435490
436- func parseHTTPProxy () * url.URL {
437- proxyURLStr := os .Getenv ("HTTP_PROXY" )
438- if proxyURLStr == "" {
439- proxyURLStr = os .Getenv ("HTTPS_PROXY" )
440- }
441- if proxyURLStr == "" {
442- return nil
443- }
491+ func ParseHTTPProxy () * url.URL {
492+ // Try each proxy variable in order
493+ vars := []string {"http_proxy" , "HTTP_PROXY" , "https_proxy" , "HTTPS_PROXY" , "all_proxy" , "ALL_PROXY" }
494+ return parseProxyURL (vars )
495+ }
444496
445- if ! strings . HasPrefix ( proxyURLStr , "http://" ) && ! strings . HasPrefix ( proxyURLStr , "https://" ) {
446- proxyURLStr = "http://" + proxyURLStr
447- }
497+ func ParseSOCKSProxy () * url. URL {
498+ return parseProxyURL ([] string { "SOCKS5_PROXY" })
499+ }
448500
449- proxyURL , err := url .Parse (proxyURLStr )
450- if err != nil {
451- return nil
452- }
453- return proxyURL
501+ // TCPDialer is a smart dialer that handles proxy configuration automatically
502+ type TCPDialer struct {
503+ proxyFunc ProxyFunc
504+ timeout time.Duration
454505}
455506
456- func parseSOCKSProxy () * url.URL {
457- proxyURLStr := os .Getenv ("SOCKS5_PROXY" )
458- if proxyURLStr == "" {
459- return nil
507+ // NewTCPDialer creates a new TCP dialer with the given proxy function and timeout
508+ func NewTCPDialer (proxyFunc ProxyFunc , timeout time.Duration ) * TCPDialer {
509+ if proxyFunc == nil {
510+ // Default to environment-based proxy configuration.
511+ proxyFunc = TCPProxyFromEnvironment
460512 }
461-
462- if ! strings . HasPrefix ( proxyURLStr , "socks5://" ) {
463- proxyURLStr = "socks5://" + proxyURLStr
513+ return & TCPDialer {
514+ proxyFunc : proxyFunc ,
515+ timeout : timeout ,
464516 }
517+ }
465518
466- proxyURL , err := url .Parse (proxyURLStr )
519+ // DialContext connects to the given address, using a proxy if configured
520+ func (d * TCPDialer ) DialContext (ctx context.Context , network , addr string ) (net.Conn , error ) {
521+ // Determine if we should use a proxy
522+ proxyURL , err := d .proxyFunc (addr )
467523 if err != nil {
468- return nil
524+ return nil , fmt .Errorf ("proxy resolution failed: %w" , err )
525+ }
526+
527+ if proxyURL != nil {
528+ // Use proxy connection
529+ switch proxyURL .Scheme {
530+ case "http" , "https" :
531+ proxyDialer := NewHTTPProxyDialer (proxyURL , addr , d .timeout )
532+ return proxyDialer .Connect ()
533+ case "socks5" :
534+ socksDialer := NewSOCKS5Dialer (proxyURL , addr , d .timeout )
535+ return socksDialer .Connect ()
536+ default :
537+ return nil , fmt .Errorf ("unsupported proxy scheme: %s" , proxyURL .Scheme )
538+ }
469539 }
470- return proxyURL
540+
541+ // Direct connection
542+ dialer := & net.Dialer {
543+ Timeout : d .timeout ,
544+ }
545+ return dialer .DialContext (ctx , network , addr )
546+ }
547+
548+ // Dial connects to the given address (non-context version)
549+ func (d * TCPDialer ) Dial (network , addr string ) (net.Conn , error ) {
550+ return d .DialContext (context .Background (), network , addr )
471551}
0 commit comments