From b2e862d65156b09802eac5624f00ad8e30f69aaf Mon Sep 17 00:00:00 2001
From: GuanM <1-GuanM@users.noreply.gitlab.54sxh.cn>
Date: Wed, 19 Jun 2024 14:00:24 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9C=80=E5=90=8E=E4=B8=80=E6=9D=BF=E8=87=AA?=
 =?UTF-8?q?=E5=B7=B1=E5=86=99=E7=9A=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 go.mod           |   5 +-
 go.sum           |   2 +
 main.go          |  24 +++----
 task/download.go | 178 ++++++++++++++++++++++++++++++++++++++++++++++-
 task/ip.go       |  11 +--
 task/tcpping.go  |  43 ++++--------
 6 files changed, 211 insertions(+), 52 deletions(-)

diff --git a/go.mod b/go.mod
index d6b646b..b8ad0b8 100644
--- a/go.mod
+++ b/go.mod
@@ -2,7 +2,10 @@ module DockerST
 
 go 1.22
 
-require github.com/schollz/progressbar/v3 v3.14.4
+require (
+	github.com/VividCortex/ewma v1.2.0
+	github.com/schollz/progressbar/v3 v3.14.4
+)
 
 require (
 	github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
diff --git a/go.sum b/go.sum
index 9af5570..9a8a09f 100644
--- a/go.sum
+++ b/go.sum
@@ -1,3 +1,5 @@
+github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow=
+github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
diff --git a/main.go b/main.go
index 807ded1..47ebca6 100644
--- a/main.go
+++ b/main.go
@@ -9,7 +9,6 @@ import (
 var (
 	VersionPrint bool
 	Version      string
-	DDownload    bool
 )
 
 func init() {
@@ -21,9 +20,10 @@ func init() {
 	flag.IntVar(&task.MinMS, "mis", 0, "只输出高于指定平均延迟的 IP")
 	flag.IntVar(&task.MaxMS, "mxs", 1000, "只输出低于指定平均延迟的 IP")
 	flag.IntVar(&task.DownloadNum, "dn", 10, "下载数量")
-	flag.StringVar(&task.DownloadUrl, "url", "https://github.com", "默认文件下载地址")
+	flag.StringVar(&task.URL, "url", "https://cf.xiu2.xyz/url", "默认文件下载地址")
+	flag.Float64Var(&task.MinSpeed, "md", 0, "最低下载速度")
 	flag.BoolVar(&task.TestAll, "ta", false, "测试所有 IP")
-	flag.BoolVar(&DDownload, "du", false, "禁止下载")
+	flag.BoolVar(&task.Disable, "dd", true, "禁止下载")
 	flag.BoolVar(&VersionPrint, "v", false, "输出版本")
 	flag.BoolVar(&task.IsOff, "om", false, "不下载子网列表")
 	flag.Parse()
@@ -40,18 +40,12 @@ func main() {
 	task.InitRandSeed()
 	// 输出版本
 	fmt.Printf("# DockerST %s \n", Version)
-	pingData := task.CreateData().Run().ExcludeInvalid()
-	// 按照延迟排序
-	sortedPingData := task.SortNodesDesc(pingData)
-	// 仅输出前10个结果
-	for i, ipDelay := range sortedPingData {
-		if i >= 10 {
-			break
+	pingData := task.CreateData().Run().ExcludeInvalid().SortNodesDesc()
+	DownloadData := task.TestDownloadSpeed(pingData)
+	for a, v := range DownloadData {
+		if a == 10 {
+			return
 		}
-		fmt.Printf("IP: %s, 延迟: %s\n", ipDelay.IP.String(), ipDelay.Delay)
-	}
-
-	if DDownload {
-
+		fmt.Printf("IP: %s, 延迟: %v, 下载速度: %.2f MB/s\n", v.IP.String(), v.Delay, v.DownloadSpeed)
 	}
 }
diff --git a/task/download.go b/task/download.go
index ac44e06..b5d986a 100644
--- a/task/download.go
+++ b/task/download.go
@@ -1,6 +1,178 @@
 package task
 
-var (
-	DownloadUrl = ""
-	DownloadNum int
+import (
+	"context"
+	"fmt"
+	"github.com/VividCortex/ewma"
+	"github.com/schollz/progressbar/v3"
+	"io"
+	"net"
+	"net/http"
+	"sort"
+	"time"
 )
+
+const (
+	bufferSize                     = 1024
+	defaultURL                     = "https://cf.xiu2.xyz/url"
+	defaultTimeout                 = 10 * time.Second
+	defaultDisableDownload         = false
+	defaultTestNum                 = 10
+	defaultMinSpeed        float64 = 0.0
+)
+
+var (
+	URL         = defaultURL
+	Timeout     = defaultTimeout
+	Disable     = defaultDisableDownload
+	DownloadNum = defaultTestNum
+	MinSpeed    = defaultMinSpeed
+)
+
+type BySpeed []IPDelay
+
+func TestDownloadSpeed(p []IPDelay) []IPDelay {
+	if Disable {
+		return p
+	}
+	if len(p) <= 0 { // IP数组长度(IP数量) 大于 0 时才会继续下载测速
+		fmt.Println("\n[信息] 延迟测速结果 IP 数量为 0，跳过下载测速。")
+		return p
+	}
+	speedSet := make([]IPDelay, 0)
+	testNum := DownloadNum
+	if len(p) < DownloadNum || MinSpeed > 0 { // 如果IP数组长度(IP数量) 小于下载测速数量（-dn），则次数修正为IP数
+		testNum = len(p)
+	}
+	if testNum < DownloadNum {
+		DownloadNum = testNum
+	}
+	fmt.Printf("开始下载测速（下限：%.2f MB/s, 数量：%d, 队列：%d）\n", MinSpeed, DownloadNum, testNum)
+	// 创建进度条
+	bar := progressbar.NewOptions(len(p),
+		progressbar.OptionSetWidth(50),
+		progressbar.OptionSetDescription("下载测速"),
+		progressbar.OptionShowCount(),
+		progressbar.OptionShowIts(),
+	)
+	for i := 0; i < testNum; i++ {
+		speed := downloadHandler(p[i].IP)
+		p[i].DownloadSpeed = speed
+		// 在每个 IP 下载测速后，以 [下载速度下限] 条件过滤结果
+		if speed >= MinSpeed*1024*1024 {
+			p[i].DownloadSpeed = speed
+			err := bar.Add(1)
+			if err != nil {
+				return nil
+			}
+		} else {
+			p = p[:i]
+			break
+		}
+	}
+	err := bar.Finish()
+	if err != nil {
+		return nil
+	}
+	if len(speedSet) == 0 { // 没有符合速度限制的数据，返回所有测试数据
+		speedSet = p
+	}
+	// 按速度排序
+	sort.Sort(BySpeed(speedSet))
+	return speedSet
+}
+
+func getDialContext(ip *net.IPAddr) func(ctx context.Context, network, address string) (net.Conn, error) {
+	var fakeSourceAddr string
+	if IsIpv4(ip.String()) {
+		fakeSourceAddr = fmt.Sprintf("%s:%d", ip.String(), TcpPort)
+	} else {
+		fakeSourceAddr = fmt.Sprintf("[%s]:%d", ip.String(), TcpPort)
+	}
+	return func(ctx context.Context, network, address string) (net.Conn, error) {
+		return (&net.Dialer{}).DialContext(ctx, network, fakeSourceAddr)
+	}
+}
+
+// return download Speed
+func downloadHandler(ip *net.IPAddr) float64 {
+	client := &http.Client{
+		Transport: &http.Transport{DialContext: getDialContext(ip)},
+		Timeout:   Timeout,
+		CheckRedirect: func(req *http.Request, via []*http.Request) error {
+			if len(via) > 10 { // 限制最多重定向 10 次
+				return http.ErrUseLastResponse
+			}
+			if req.Header.Get("Referer") == defaultURL { // 当使用默认下载测速地址时，重定向不携带 Referer
+				req.Header.Del("Referer")
+			}
+			return nil
+		},
+	}
+	req, err := http.NewRequest("GET", URL, nil)
+	if err != nil {
+		return 0.0
+	}
+
+	req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/98.0.4758.80 Safari/537.36")
+
+	response, err := client.Do(req)
+	if err != nil {
+		return 0.0
+	}
+	defer func(Body io.ReadCloser) {
+		err = Body.Close()
+		if err != nil {
+		}
+	}(response.Body)
+	if response.StatusCode != 200 {
+		return 0.0
+	}
+	timeStart := time.Now()           // 开始时间（当前）
+	timeEnd := timeStart.Add(Timeout) // 加上下载测速时间得到的结束时间
+
+	contentLength := response.ContentLength // 文件大小
+	buffer := make([]byte, bufferSize)
+
+	var (
+		contentRead     int64 = 0
+		timeSlice             = Timeout / 100
+		timeCounter           = 1
+		lastContentRead int64 = 0
+	)
+
+	var nextTime = timeStart.Add(timeSlice * time.Duration(timeCounter))
+	e := ewma.NewMovingAverage()
+
+	// 循环计算，如果文件下载完了（两者相等），则退出循环（终止测速）
+	for contentLength != contentRead {
+		currentTime := time.Now()
+		if currentTime.After(nextTime) {
+			timeCounter++
+			nextTime = timeStart.Add(timeSlice * time.Duration(timeCounter))
+			e.Add(float64(contentRead - lastContentRead))
+			lastContentRead = contentRead
+		}
+		// 如果超出下载测速时间，则退出循环（终止测速）
+		if currentTime.After(timeEnd) {
+			break
+		}
+		bufferRead, err := response.Body.Read(buffer)
+		if err != nil {
+			if err != io.EOF { // 如果文件下载过程中遇到报错（如 Timeout），且并不是因为文件下载完了，则退出循环（终止测速）
+				break
+			} else if contentLength == -1 { // 文件下载完成 且 文件大小未知，则退出循环（终止测速），例如：https://speed.cloudflare.com/__down?bytes=200000000 这样的，如果在 10 秒内就下载完成了，会导致测速结果明显偏低甚至显示为 0.00（下载速度太快时）
+				break
+			}
+			// 获取上个时间片
+			lastTimeSlice := timeStart.Add(timeSlice * time.Duration(timeCounter-1))
+			// 下载数据量 / (用当前时间 - 上个时间片/ 时间片)
+			e.Add(float64(contentRead-lastContentRead) / (float64(currentTime.Sub(lastTimeSlice)) / float64(timeSlice)))
+		}
+		contentRead += int64(bufferRead)
+	}
+	return e.Value() / (Timeout.Seconds() / 120)
+}
+func (a BySpeed) Len() int           { return len(a) }
+func (a BySpeed) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
+func (a BySpeed) Less(i, j int) bool { return a[i].DownloadSpeed > a[j].DownloadSpeed } // 这里使用 > 使其降序排序
diff --git a/task/ip.go b/task/ip.go
index fdce327..13dd339 100644
--- a/task/ip.go
+++ b/task/ip.go
@@ -117,7 +117,7 @@ func GetIPv4List() []string {
 }
 
 // ExcludeInvalid 排除不合格节点
-func (p *IPRangeList) ExcludeInvalid() []IPDelay {
+func (p *IPRangeList) ExcludeInvalid() *IPRangeList {
 	// 初始化一个空IPDelay切片
 	var delays []IPDelay
 	// 遍历IPRangeList的Delays切片
@@ -133,13 +133,14 @@ func (p *IPRangeList) ExcludeInvalid() []IPDelay {
 		// 将IPDelay的IP和Delay添加到delays切片中
 		delays = append(delays, IPDelay{IP: p.Ips[ip], Delay: delay.Delay, DownloadSpeed: 0})
 	}
-	return delays
+	p.Delays = delays
+	return p
 }
 
 // SortNodesDesc 按延迟降序排列
-func SortNodesDesc(p []IPDelay) []IPDelay {
-	sorted := make([]IPDelay, len(p))
-	_ = copy(sorted, p)
+func (p *IPRangeList) SortNodesDesc() []IPDelay {
+	sorted := make([]IPDelay, len(p.Delays))
+	_ = copy(sorted, p.Delays)
 	// 使用sort.Slice 对切片进行排序
 	sort.Slice(sorted, func(i, j int) bool {
 		return sorted[i].Delay < sorted[j].Delay
diff --git a/task/tcpping.go b/task/tcpping.go
index 87df460..1eecf79 100644
--- a/task/tcpping.go
+++ b/task/tcpping.go
@@ -59,36 +59,23 @@ func (p *IPRangeList) Run() *IPRangeList {
 
 // TCPing 通过TCP连接测试IP是否可用
 func TCPing(ip *net.IPAddr) (bool, time.Duration) {
-	var totalDuration time.Duration
-	var successCount int
-	for i := 0; i < 4; i++ { // 重试4次
-		func() {
-			startTime := time.Now()
-			var fullAddress string
-			if IsIpv4(ip.String()) {
-				fullAddress = fmt.Sprintf("%s:%d", ip.String(), TcpPort)
-			} else {
-				fullAddress = fmt.Sprintf("[%s]:%d", ip.String(), TcpPort)
-			}
-			conn, err := net.DialTimeout("tcp", fullAddress, TcpConnectTimeOut)
-			if err != nil {
-				return
-			}
-			defer func() {
-				err = conn.Close()
-				if err != nil {
-				}
-			}()
-			duration := time.Since(startTime)
-			totalDuration += duration
-			successCount++
-		}()
+	startTime := time.Now()
+	var fullAddress string
+	if IsIpv4(ip.String()) {
+		fullAddress = fmt.Sprintf("%s:%d", ip.String(), TcpPort)
+	} else {
+		fullAddress = fmt.Sprintf("[%s]:%d", ip.String(), TcpPort)
 	}
-
-	if successCount == 0 { // 如果所有尝试都失败，返回 false 和 0
+	conn, err := net.DialTimeout("tcp", fullAddress, TcpConnectTimeOut)
+	if err != nil {
 		return false, 0
 	}
+	defer func(conn net.Conn) {
+		err := conn.Close()
+		if err != nil {
 
-	averageDuration := totalDuration / time.Duration(successCount) // 计算平均持续时间
-	return true, averageDuration
+		}
+	}(conn)
+	duration := time.Since(startTime)
+	return true, duration
 }