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) {
	ipv4latencyFile := "data/latency4.json"
	ipv6latencyFile := "data/latency6.json"
	if !forceRun {
		ipv4Loaded := false
		ipv6Loaded := false
		latencyData := r.readLatencyFile(ipv4latencyFile)
		if latencyData != nil {
			r.ipv4latencyMap = *latencyData
			ipv4Loaded = true
		}
		latencyData = r.readLatencyFile(ipv6latencyFile)
		if latencyData != nil {
			r.ipv6latencyMap = *latencyData
			ipv6Loaded = true
		}
		if ipv4Loaded && ipv6Loaded {
			return
		} else {
			r.log.Warn("Network test files not found")
		}
	}

	r.log.Info("Network test will start now. IGNORE THE WARNINGS!")
	r.runLatencyTest()
	r.log.Info("Network test completed!")
	r.log.Infof("To rerun the network test, run 'zurg network-test', or delete the files %s and %s and run zurg again", ipv4latencyFile, ipv6latencyFile)
	r.writeLatencyFile(ipv4latencyFile, r.ipv4latencyMap)
	r.writeLatencyFile(ipv6latencyFile, r.ipv6latencyMap)
}

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

	// Convert the latency map to a slice of key-value pairs
	type kv struct {
		Key   string
		Value float64
	}

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

	// Sort the slice by latency values
	sort.Slice(kvList, func(i, j int) bool {
		return kvList[i].Value < kvList[j].Value
	})

	// Calculate the number of hosts to return (top 50%)
	n := len(kvList) / 5
	if len(kvList)%5 != 0 {
		n++
	}

	// Collect the keys of the top 50% hosts
	var optimalHosts []string
	for i := 0; i < n; i++ {
		optimalHosts = append(optimalHosts, kvList[i].Key)
	}

	return optimalHosts
}

func (r *IPRepository) runLatencyTest() {
	limit := 99
	start := 0
	for {
		lastDomainsWorked := false
		for i := start; i <= limit; i++ {
			domain := fmt.Sprintf("%d.download.real-debrid.com", i)
			ips, err := net.LookupIP(domain)
			if err != nil || len(ips) == 0 {
				continue
			}

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

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

			domain = fmt.Sprintf("%d.download.real-debrid.cloud", i)
			ips, err = net.LookupIP(domain)
			if err != nil || len(ips) == 0 {
				continue
			}

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

			latency, err = r.testDomainLatency(r.ipv6client, domain)
			if err == nil {
				r.ipv6latencyMap[domain] = latency
				r.log.Debugf("Latency from ipv6 %s: %.5f seconds", domain, 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 testFileSize = 1 // byte
	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.MethodGet, testURL, nil)
		if err != nil {
			r.log.Warnf("Failed to create request for %s: %v", domain, err)
			retErr = err
			break
		}

		headers := make(http.Header)
		headers.Set("Range", fmt.Sprintf("bytes=0-%d", testFileSize-1))
		req.Header = headers

		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
		}

		limitedReader := io.LimitReader(resp.Body, testFileSize)

		_, err = io.Copy(io.Discard, limitedReader)
		resp.Body.Close()

		if err != nil && err != io.EOF {
			r.log.Warnf("Failed to read from %s: %v", domain, err)
			retErr = err
			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
	}
}