dial.go 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. // Copyright 2019 The Go Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package proxy
  5. import (
  6. "context"
  7. "net"
  8. )
  9. // A ContextDialer dials using a context.
  10. type ContextDialer interface {
  11. DialContext(ctx context.Context, network, address string) (net.Conn, error)
  12. }
  13. // Dial works like DialContext on net.Dialer but using a dialer returned by FromEnvironment.
  14. //
  15. // The passed ctx is only used for returning the Conn, not the lifetime of the Conn.
  16. //
  17. // Custom dialers (registered via RegisterDialerType) that do not implement ContextDialer
  18. // can leak a goroutine for as long as it takes the underlying Dialer implementation to timeout.
  19. //
  20. // A Conn returned from a successful Dial after the context has been cancelled will be immediately closed.
  21. func Dial(ctx context.Context, network, address string) (net.Conn, error) {
  22. d := FromEnvironment()
  23. if xd, ok := d.(ContextDialer); ok {
  24. return xd.DialContext(ctx, network, address)
  25. }
  26. return dialContext(ctx, d, network, address)
  27. }
  28. // WARNING: this can leak a goroutine for as long as the underlying Dialer implementation takes to timeout
  29. // A Conn returned from a successful Dial after the context has been cancelled will be immediately closed.
  30. func dialContext(ctx context.Context, d Dialer, network, address string) (net.Conn, error) {
  31. var (
  32. conn net.Conn
  33. done = make(chan struct{}, 1)
  34. err error
  35. )
  36. go func() {
  37. conn, err = d.Dial(network, address)
  38. close(done)
  39. if conn != nil && ctx.Err() != nil {
  40. conn.Close()
  41. }
  42. }()
  43. select {
  44. case <-ctx.Done():
  45. err = ctx.Err()
  46. case <-done:
  47. }
  48. return conn, err
  49. }