github.com/nxtrace/NTrace-core@v1.3.1-0.20240513132635-39169291e8c9/util/latency.go (about)

     1  package util
     2  
     3  import (
     4  	"crypto/tls"
     5  	"fmt"
     6  	"io"
     7  	"log"
     8  	"net"
     9  	"net/http"
    10  	"strings"
    11  	"time"
    12  
    13  	"github.com/fatih/color"
    14  )
    15  
    16  type ResponseInfo struct {
    17  	IP      string
    18  	Latency string
    19  	Content string
    20  }
    21  
    22  var (
    23  	results = make(chan ResponseInfo)
    24  	timeout = 5 * time.Second
    25  )
    26  var FastIpCache = ""
    27  
    28  func GetFastIP(domain string, port string, enableOutput bool) string {
    29  	proxyUrl := GetProxy()
    30  	if proxyUrl != nil {
    31  		return "origin-fallback.nxtrace.org"
    32  	}
    33  	if FastIpCache != "" {
    34  		return FastIpCache
    35  	}
    36  
    37  	var ips []net.IP
    38  	var err error
    39  	if domain == "origin-fallback.nxtrace.org" {
    40  		ips, err = net.LookupIP("api.nxtrace.org")
    41  	} else {
    42  		ips, err = net.LookupIP(domain)
    43  	}
    44  
    45  	if err != nil {
    46  		log.Fatal("DNS resolution failed, please check your system DNS Settings")
    47  	}
    48  
    49  	if len(ips) == 0 {
    50  		// 添加默认IP 45.88.195.154
    51  		ips = append(ips, net.ParseIP("45.88.195.154"))
    52  	}
    53  
    54  	for _, ip := range ips {
    55  		go checkLatency(domain, ip.String(), port)
    56  	}
    57  
    58  	var result ResponseInfo
    59  
    60  	select {
    61  	case result = <-results:
    62  	//等待5s没有结果 视为连不上API了
    63  	case <-time.After(timeout):
    64  		log.Println("IP connection has been timeout, please check your network")
    65  
    66  	}
    67  
    68  	//if len(ips) > 0 {
    69  	if enableOutput {
    70  		_, _ = fmt.Fprintf(color.Output, "%s preferred API IP - %s - %s - %s",
    71  			color.New(color.FgWhite, color.Bold).Sprintf("[NextTrace API]"),
    72  			color.New(color.FgGreen, color.Bold).Sprintf("%s", result.IP),
    73  			color.New(color.FgCyan, color.Bold).Sprintf("%sms", result.Latency),
    74  			color.New(color.FgGreen, color.Bold).Sprintf("%s", result.Content),
    75  		)
    76  	}
    77  	//}
    78  
    79  	//有些时候真的啥都不通，还是挑一个顶上吧
    80  	if result.IP == "" {
    81  		result.IP = "45.88.195.154"
    82  	}
    83  
    84  	FastIpCache = result.IP
    85  	return result.IP
    86  }
    87  
    88  func checkLatency(domain string, ip string, port string) {
    89  	start := time.Now()
    90  	if !strings.Contains(ip, ".") {
    91  		ip = "[" + ip + "]"
    92  	}
    93  
    94  	// 自定义DialContext以使用指定的IP连接
    95  	transport := &http.Transport{
    96  		//DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
    97  		//	return net.DialTimeout(network, addr, 1*time.Second)
    98  		//},
    99  		TLSClientConfig: &tls.Config{
   100  			ServerName: domain,
   101  		},
   102  		TLSHandshakeTimeout: timeout,
   103  	}
   104  
   105  	client := &http.Client{
   106  		Transport: transport,
   107  		Timeout:   timeout,
   108  	}
   109  
   110  	//此处虽然是 https://domain/ 但是实际上会使用指定的IP连接
   111  	req, err := http.NewRequest("GET", "https://"+ip+":"+port+"/", nil)
   112  	if err != nil {
   113  		// !!! 此处不要给results返回任何值
   114  		//results <- ResponseInfo{IP: ip, Latency: "error", Content: ""}
   115  		return
   116  	}
   117  	req.Host = domain
   118  	resp, err := client.Do(req)
   119  	if err != nil {
   120  		//results <- ResponseInfo{IP: ip, Latency: "error", Content: ""}
   121  		return
   122  	}
   123  	defer resp.Body.Close()
   124  
   125  	bodyBytes, err := io.ReadAll(resp.Body)
   126  	if err != nil {
   127  		//results <- ResponseInfo{IP: ip, Latency: "error", Content: ""}
   128  		return
   129  	}
   130  	bodyString := string(bodyBytes)
   131  
   132  	latency := fmt.Sprintf("%.2f", float64(time.Since(start))/float64(time.Millisecond))
   133  	results <- ResponseInfo{IP: ip, Latency: latency, Content: bodyString}
   134  }