zurg/pkg/http/ip.go

package http

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"sort"
	"time"

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

type IPRepository struct {
	ipv4client     *HTTPClient
	ipv6client     *HTTPClient
	ipv4latencyMap map[string]float64
	ipv6latencyMap map[string]float64
	testURL        string
	log            *logutil.Logger
}

func NewIPRepository(ipv4client *HTTPClient, ipv6client *HTTPClient, testURL string, log *logutil.Logger) *IPRepository {
	repo := &IPRepository{
		ipv4client:     ipv4client,
		ipv6client:     ipv6client,
		ipv4latencyMap: make(map[string]float64),
		ipv6latencyMap: make(map[string]float64),
		testURL:        testURL,
		log:            log,
	}

	return repo
}

func (r *IPRepository) NetworkTest(forceRun bool, persist bool) {
	ipv4HostsFile := "data/ipv4-hosts.json"
	ipv6HostsFile := "data/ipv6-hosts.json"
	if !forceRun {
		ipv4Loaded := false
		ipv6Loaded := false
		latencyData := r.readLatencyFile(ipv4HostsFile)
		if latencyData != nil {
			r.ipv4latencyMap = *latencyData
			ipv4Loaded = true
		}
		latencyData = r.readLatencyFile(ipv6HostsFile)
		if latencyData != nil {
			r.ipv6latencyMap = *latencyData
			ipv6Loaded = true
		}
		if ipv4Loaded && ipv6Loaded {
			return
		} else {
			r.log.Warn("Network test files not found, running network test")
		}
	}

	r.log.Info("zurg will check for all reachable download servers. You can set 'cache_network_test_results: true' in your config to skip this test in the future.")
	r.log.Warn("IGNORE THE WARNINGS!")
	r.runLatencyTest()
	r.log.Info("Network test completed!")

	if persist {
		r.log.Infof("To rerun the network test, run 'zurg network-test', or delete the files %s and %s and run zurg again", ipv4HostsFile, ipv6HostsFile)
		r.writeLatencyFile(ipv4HostsFile, r.ipv4latencyMap)
		r.writeLatencyFile(ipv6HostsFile, r.ipv6latencyMap)
	}
}

func (r *IPRepository) GetHosts(numberOfHosts int, ipv6 bool) []string {
	latencyMap := r.ipv4latencyMap
	if ipv6 {
		latencyMap = r.ipv6latencyMap
	}

	type kv struct {
		Key   string
		Value float64
	}

	var kvList []kv
	for k, v := range latencyMap {
		kvList = append(kvList, kv{k, v})
	}

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

	var hosts []string
	if numberOfHosts == 0 {
		numberOfHosts = len(kvList)
	}
	for i := 0; i < numberOfHosts && i < len(kvList); i++ {
		hosts = append(hosts, kvList[i].Key)
	}

	return hosts
}

func (r *IPRepository) runLatencyTest() {
	limit := 99
	start := 0
	for {
		lastDomainsWorked := false
		for i := start; i <= limit; i++ {
			// EU servers
			ipv4Domain := fmt.Sprintf("%d-4.download.real-debrid.com", i)
			ipv6Domain := fmt.Sprintf("%d-6.download.real-debrid.com", i)

			ips, err := net.LookupIP(ipv4Domain)
			if err != nil || len(ips) == 0 {
				continue
			}

			latency, err := r.testDomainLatency(r.ipv4client, ipv4Domain)
			if err == nil {
				r.ipv4latencyMap[ipv4Domain] = latency
				r.log.Debugf("Latency from ipv4 %s: %.5f seconds", ipv4Domain, latency)
				if i >= limit-2 {
					lastDomainsWorked = true
				}
			}

			latency, err = r.testDomainLatency(r.ipv6client, ipv6Domain)
			if err == nil {
				r.ipv6latencyMap[ipv6Domain] = latency
				r.log.Debugf("Latency from ipv6 %s: %.5f seconds", ipv6Domain, latency)
				if i >= limit-2 {
					lastDomainsWorked = true
				}
			}

			// EU servers (old)
			ipv4Domain = fmt.Sprintf("%d.download.real-debrid.com", i)
			ipv6Domain = fmt.Sprintf("%d.download.real-debrid.com", i)

			ips, err = net.LookupIP(ipv4Domain)
			if err != nil || len(ips) == 0 {
				continue
			}

			latency, err = r.testDomainLatency(r.ipv4client, ipv4Domain)
			if err == nil {
				r.ipv4latencyMap[ipv4Domain] = latency
				r.log.Debugf("Latency from ipv4 %s: %.5f seconds", ipv4Domain, latency)
				if i >= limit-2 {
					lastDomainsWorked = true
				}
			}

			latency, err = r.testDomainLatency(r.ipv6client, ipv6Domain)
			if err == nil {
				r.ipv6latencyMap[ipv6Domain] = latency
				r.log.Debugf("Latency from ipv6 %s: %.5f seconds", ipv6Domain, latency)
				if i >= limit-2 {
					lastDomainsWorked = true
				}
			}

			// Cloudflare servers
			ipv4Domain = fmt.Sprintf("%d.download.real-debrid.cloud", i)
			ipv6Domain = fmt.Sprintf("%d.download.real-debrid.cloud", i)

			ips, err = net.LookupIP(ipv4Domain)
			if err != nil || len(ips) == 0 {
				continue
			}

			latency, err = r.testDomainLatency(r.ipv4client, ipv4Domain)
			if err == nil {
				r.ipv4latencyMap[ipv4Domain] = latency
				r.log.Debugf("Latency from ipv4 %s: %.5f seconds", ipv4Domain, latency)
				if i >= limit-2 {
					lastDomainsWorked = true
				}
			}

			latency, err = r.testDomainLatency(r.ipv6client, ipv6Domain)
			if err == nil {
				r.ipv6latencyMap[ipv6Domain] = latency
				r.log.Debugf("Latency from ipv6 %s: %.5f seconds", ipv6Domain, latency)
				if i >= limit-2 {
					lastDomainsWorked = true
				}
			}
		}
		if lastDomainsWorked {
			start = limit + 1
			limit += 10
		} else {
			break
		}
	}
}

func (r *IPRepository) testDomainLatency(client *HTTPClient, domain string) (float64, error) {
	const iterations = 3
	testURL := fmt.Sprintf("https://%s/speedtest/test.rar/%f", domain, rand.Float64())
	if r.testURL != "" {
		parsedURL, err := url.Parse(r.testURL)
		if err == nil {
			testURL = fmt.Sprintf("https://%s%s", domain, parsedURL.Path)
		}
	}

	var totalDuration float64
	var retErr error
	for i := 0; i < iterations; i++ {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		req, err := http.NewRequestWithContext(ctx, http.MethodHead, testURL, nil)
		if err != nil {
			r.log.Warnf("Failed to create request for %s: %v", domain, err)
			retErr = err
			break
		}

		start := time.Now()
		resp, err := client.Do(req)
		if err != nil {
			r.log.Warnf("Failed to download from %s: %v", domain, err)
			retErr = err
			break
		}
		if resp.StatusCode != http.StatusOK {
			r.log.Warnf("Failed to download from %s: %s", domain, resp.Status)
			retErr = fmt.Errorf("status code: %s", resp.Status)
			break
		}
		duration := time.Since(start).Seconds()
		totalDuration += duration
	}
	if retErr != nil {
		return 0, retErr
	}

	avgDuration := totalDuration / 3
	return avgDuration, nil
}

func (r *IPRepository) readLatencyFile(latencyFile string) *map[string]float64 {
	if _, err := os.Stat(latencyFile); err == nil {
		file, err := os.Open(latencyFile)
		if err != nil {
			return nil
		}
		defer file.Close()
		jsonData, err := io.ReadAll(file)
		if err != nil {
			return nil
		}
		var ipv4latencyMap map[string]float64
		if err := json.Unmarshal(jsonData, &ipv4latencyMap); err != nil {
			return nil
		}
		return &ipv4latencyMap
	}
	return nil
}

func (r *IPRepository) writeLatencyFile(latencyFile string, data interface{}) {
	file, err := os.Create(latencyFile)
	if err != nil {
		r.log.Warnf("Cannot create latency file %s: %v", latencyFile, err)
		return
	}
	defer file.Close()

	jsonData, err := json.Marshal(data)
	if err != nil {
		r.log.Warnf("Cannot marshal latency map: %v", err)
		return
	}

	if _, err := file.Write(jsonData); err != nil {
		r.log.Warnf("Cannot write to latency file %s: %v", latencyFile, err)
		return
	}
}