github.com/laof/lite-speed-test@v0.0.0-20230930011949-1f39b7037845/transport/dialer/dialer.go (about)

     1  package dialer
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"net"
     7  
     8  	"github.com/laof/lite-speed-test/transport/resolver"
     9  )
    10  
    11  func DialContext(ctx context.Context, network, address string) (net.Conn, error) {
    12  	switch network {
    13  	case "tcp4", "tcp6", "udp4", "udp6":
    14  		host, port, err := net.SplitHostPort(address)
    15  		if err != nil {
    16  			return nil, err
    17  		}
    18  
    19  		dialer, err := Dialer()
    20  		if err != nil {
    21  			return nil, err
    22  		}
    23  
    24  		var ip net.IP
    25  		switch network {
    26  		case "tcp4", "udp4":
    27  			ip, err = resolver.ResolveIPv4(host)
    28  		default:
    29  			ip, err = resolver.ResolveIPv6(host)
    30  		}
    31  
    32  		if err != nil {
    33  			return nil, err
    34  		}
    35  
    36  		return dialer.DialContext(ctx, network, net.JoinHostPort(ip.String(), port))
    37  	case "tcp", "udp":
    38  		return dualStackDialContext(ctx, network, address)
    39  	default:
    40  		return nil, errors.New("network invalid")
    41  	}
    42  }
    43  
    44  func ListenPacket(network, address string) (net.PacketConn, error) {
    45  	// cfg := &net.ListenConfig{}
    46  	// return cfg.ListenPacket(context.Background(), network, address)
    47  	return effectiveListener.ListenPacket(context.Background(), network, address)
    48  }
    49  
    50  func dualStackDialContext(ctx context.Context, network, address string) (net.Conn, error) {
    51  	host, port, err := net.SplitHostPort(address)
    52  	if err != nil {
    53  		return nil, err
    54  	}
    55  
    56  	returned := make(chan struct{})
    57  	defer close(returned)
    58  
    59  	type dialResult struct {
    60  		net.Conn
    61  		error
    62  		resolved bool
    63  		ipv6     bool
    64  		done     bool
    65  	}
    66  	results := make(chan dialResult)
    67  	var primary, fallback dialResult
    68  
    69  	startRacer := func(ctx context.Context, network, host string, ipv6 bool) {
    70  		result := dialResult{ipv6: ipv6, done: true}
    71  		defer func() {
    72  			select {
    73  			case results <- result:
    74  			case <-returned:
    75  				if result.Conn != nil {
    76  					result.Conn.Close()
    77  				}
    78  			}
    79  		}()
    80  
    81  		dialer, err := Dialer()
    82  		if err != nil {
    83  			result.error = err
    84  			return
    85  		}
    86  
    87  		var ip net.IP
    88  		if ipv6 {
    89  			ip, result.error = resolver.ResolveIPv6(host)
    90  		} else {
    91  			ip, result.error = resolver.ResolveIPv4(host)
    92  		}
    93  		if result.error != nil {
    94  			return
    95  		}
    96  		result.resolved = true
    97  
    98  		result.Conn, result.error = dialer.DialContext(ctx, network, net.JoinHostPort(ip.String(), port))
    99  	}
   100  
   101  	go startRacer(ctx, network+"4", host, false)
   102  	go startRacer(ctx, network+"6", host, true)
   103  
   104  	for res := range results {
   105  		if res.error == nil {
   106  			return res.Conn, nil
   107  		}
   108  
   109  		if !res.ipv6 {
   110  			primary = res
   111  		} else {
   112  			fallback = res
   113  		}
   114  
   115  		if primary.done && fallback.done {
   116  			if primary.resolved {
   117  				return nil, primary.error
   118  			} else if fallback.resolved {
   119  				return nil, fallback.error
   120  			} else {
   121  				return nil, primary.error
   122  			}
   123  		}
   124  	}
   125  
   126  	return nil, errors.New("never touched")
   127  }