zurg/pkg/realdebrid/network.go

package realdebrid

import (
	"fmt"
	"os/exec"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/debridmediamanager/zurg/pkg/hosts"
)

type HostInfo struct {
	Address string
	Hops    int
	Latency time.Duration
}

func measureLatency(host, testType string) (int, time.Duration, error) {
	traceroutePath := "traceroute"
	pingPath := "ping"
	if testType == "ipv6" {
		traceroutePath = "traceroute6"
		pingPath = "ping6"
	}

	// Try executing traceroute
	cmd := exec.Command(traceroutePath, "-n", "-q", "1", "-w", "1", host)
	out, err := cmd.CombinedOutput()

	if err == nil {
		// Traceroute executed successfully
		output := string(out)
		lines := strings.Split(output, "\n")
		hopCount := len(lines) - 1
		var latency time.Duration

		if hopCount > 0 {
			lastLine := lines[hopCount-1]
			fmt.Print(".")
			parts := strings.Fields(lastLine)
			if len(parts) >= 3 {
				latencyValue, parseErr := strconv.ParseFloat(parts[2], 64)
				if parseErr == nil {
					latency = time.Duration(latencyValue * float64(time.Millisecond))
				}
			}
		}
		return hopCount, latency, nil
	}

	// Traceroute not successful, measure latency using ping
	pingCmd := exec.Command(pingPath, "-c", "1", host)
	pingOut, pingErr := pingCmd.CombinedOutput()

	if pingErr != nil {
		return 0, 0, pingErr
	}

	fmt.Print(".")
	pingOutput := string(pingOut)
	pingLines := strings.Split(pingOutput, "\n")
	for _, line := range pingLines {
		if strings.Contains(line, "time=") {
			parts := strings.Split(line, " ")
			for _, part := range parts {
				if strings.Contains(part, "time=") {
					timeStr := strings.Split(part, "=")[1]
					timeStr = strings.TrimSuffix(timeStr, " ms")
					latencyValue, parseErr := strconv.ParseFloat(timeStr, 64)
					if parseErr == nil {
						latency := time.Duration(latencyValue * float64(time.Millisecond))
						return -1, latency, nil
					}
				}
			}
		}
	}

	return 0, 0, fmt.Errorf("failed to measure latency")
}

func RunTest(testType string) {
	fmt.Print("Running network test...")

	var ipv4Hosts, ipv6Hosts []string
	var err error

	if testType == "ipv4" || testType == "both" {
		ipv4Hosts, err = hosts.FetchHosts(hosts.IPV4)
		if err != nil {
			fmt.Println("Error fetching IPv4 hosts:", err)
			return
		}
	}

	if testType == "ipv6" || testType == "both" {
		ipv6Hosts, err = hosts.FetchHosts(hosts.IPV6)
		if err != nil {
			fmt.Println("Error fetching IPv6 hosts:", err)
			return
		}
	}

	var totalHosts int
	var hostInfos []HostInfo // Declare the slice to hold IPInfo objects
	infoChan := make(chan HostInfo)

	if testType == "ipv4" || testType == "both" {
		totalHosts += len(ipv4Hosts)
		go runLatencyTests(ipv4Hosts, "ipv4", infoChan)
	}
	if testType == "ipv6" {
		totalHosts += len(ipv6Hosts)
		go runLatencyTests(ipv6Hosts, "ipv6", infoChan)
	}
	if testType == "both" {
		totalHosts += len(ipv6Hosts)
		go runLatencyTests(ipv6Hosts, "ipv4", infoChan)
	}

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for i := 0; i < totalHosts; i++ {
			hostInfo := <-infoChan
			hostInfos = append(hostInfos, hostInfo)
		}
		close(infoChan)
	}()

	wg.Wait()

	sort.Slice(hostInfos, func(i, j int) bool {
		return hostInfos[i].Latency < hostInfos[j].Latency
	})

	const minResults = 10
	const maxResults = 20
	var okHosts []HostInfo

	if len(hostInfos) > 0 {
		// Start by adding the best IPs based on hops and latency up to the minResults
		for i := 0; i < min(len(hostInfos), minResults); i++ {
			okHosts = append(okHosts, hostInfos[i])
		}

		// Find the highest latency in the current okHosts list
		highestLatency := okHosts[len(okHosts)-1].Latency

		// Add any additional hosts that have latency within a reasonable range of the highest latency
		for _, info := range hostInfos[minResults:] {
			if len(okHosts) >= maxResults {
				break // Stop adding hosts if maxResults is reached
			}
			if info.Latency <= highestLatency+(highestLatency/3) {
				okHosts = append(okHosts, info)
			}
		}
	}

	fmt.Printf("complete!\n\nHere are the results, you can copy-paste the following to your config.yml:\n\n")
	fmt.Println("preferred_hosts:")
	for _, info := range okHosts {
		fmt.Printf("  - %s # hops: %d latency: %v\n", info.Address, info.Hops, info.Latency)
	}
}

func runLatencyTests(hosts []string, testType string, infoChan chan<- HostInfo) {
	var wg sync.WaitGroup
	semaphore := make(chan struct{}, 10)

	for _, host := range hosts {
		wg.Add(1)
		semaphore <- struct{}{}
		go func(host string) {
			defer wg.Done()
			hops, latency, err := measureLatency(host, testType)
			if err != nil {
				fmt.Printf("Error measuring latency for %s: %s\n", host, err)
			} else {
				infoChan <- HostInfo{Address: host, Hops: hops, Latency: latency}
			}
			<-semaphore
		}(host)
	}

	wg.Wait()
}