CloudflareSpeedTest/task/tcping.go

package task

import (
	"fmt"
	"net"
	"sort"
	"sync"
	"time"

	"CloudflareSpeedTest/utils"
)

const (
	tcpConnectTimeout = time.Second * 1
	maxRoutine        = 1000
	defaultRoutines   = 200
	defaultPort       = 443
	defaultPingTimes  = 4
)

var (
	Routines      = defaultRoutines
	TCPPort   int = defaultPort
	PingTimes int = defaultPingTimes
)

type Ping struct {
	wg      *sync.WaitGroup
	m       *sync.Mutex
	ips     []*net.IPAddr
	csv     utils.PingDelaySet
	control chan bool
	bar     *utils.Bar
}

func checkPingDefault() {
	if Routines <= 0 {
		Routines = defaultRoutines
	}
	if TCPPort <= 0 || TCPPort >= 65535 {
		TCPPort = defaultPort
	}
	if PingTimes <= 0 {
		PingTimes = defaultPingTimes
	}
}

func NewPing() *Ping {
	checkPingDefault()
	ips := loadIPRanges()
	return &Ping{
		wg:      &sync.WaitGroup{},
		m:       &sync.Mutex{},
		ips:     ips,
		csv:     make(utils.PingDelaySet, 0),
		control: make(chan bool, Routines),
		bar:     utils.NewBar(len(ips)),
	}
}

func (p *Ping) Run() utils.PingDelaySet {
	if len(p.ips) == 0 {
		return p.csv
	}
	ipVersion := "IPv4"
	if IPv6 { // IPv6 模式判断
		ipVersion = "IPv6"
	}
	fmt.Printf("开始延迟测速（模式：TCP %s，端口：%d，平均延迟上限：%vms，平均延迟下限：%vms)\n", ipVersion, TCPPort, utils.InputMaxDelay.Milliseconds(), utils.InputMinDelay.Milliseconds())
	for _, ip := range p.ips {
		p.wg.Add(1)
		p.control <- false
		go p.start(ip)
	}
	p.wg.Wait()
	p.bar.Done()
	sort.Sort(p.csv)
	return p.csv
}

func (p *Ping) start(ip *net.IPAddr) {
	defer p.wg.Done()
	p.tcpingHandler(ip)
	<-p.control
}

//bool connectionSucceed float32 time
func (p *Ping) tcping(ip *net.IPAddr) (bool, time.Duration) {
	startTime := time.Now()
	fullAddress := fmt.Sprintf("%s:%d", ip.String(), TCPPort)
	//fmt.Println(ip.String())
	if IPv6 { // IPv6 需要加上 []
		fullAddress = fmt.Sprintf("[%s]:%d", ip.String(), TCPPort)
	}
	conn, err := net.DialTimeout("tcp", fullAddress, tcpConnectTimeout)
	if err != nil {
		return false, 0
	}
	defer conn.Close()
	duration := time.Since(startTime)
	return true, duration
}

//pingReceived pingTotalTime
func (p *Ping) checkConnection(ip *net.IPAddr) (recv int, totalDelay time.Duration) {
	for i := 0; i < PingTimes; i++ {
		if ok, delay := p.tcping(ip); ok {
			recv++
			totalDelay += delay
		}
	}
	return
}

func (p *Ping) appendIPData(data *utils.PingData) {
	p.m.Lock()
	defer p.m.Unlock()
	p.csv = append(p.csv, utils.CloudflareIPData{
		PingData: data,
	})
}

//return Success packetRecv averagePingTime specificIPAddr
func (p *Ping) tcpingHandler(ip *net.IPAddr) {
	ipCanConnect := false
	pingRecv := 0
	var delay time.Duration
	for !ipCanConnect {
		recv, totalDlay := p.checkConnection(ip)
		if recv > 0 {
			ipCanConnect = true
			pingRecv = recv
			delay = totalDlay
		} else {
			ip.IP[15]++
			if ip.IP[15] == 0 {
				break
			}
			break
		}
	}
	p.bar.Grow(1)
	if !ipCanConnect {
		return
	}
	// for i := 0; i < PingTimes; i++ {
	// 	pingSuccess, pingTimeCurrent := p.tcping(ip)
	// 	progressHandler(utils.NormalPing)
	// 	if pingSuccess {
	// 		pingRecv++
	// 		pingTime += pingTimeCurrent
	// 	}
	// }
	data := &utils.PingData{
		IP:       ip,
		Sended:   PingTimes,
		Received: pingRecv,
		Delay:    delay / time.Duration(pingRecv),
	}
	p.appendIPData(data)
	return
}