zurg/pkg/http/ip.go

package http

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

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

type IPRepository struct {
	ipv4       []string
	ipv6       []string
	latencyMap map[string]float64
	log        *logutil.Logger
}

func NewIPRepository(log *logutil.Logger) *IPRepository {
	repo := &IPRepository{
		ipv4:       []string{},
		ipv6:       []string{},
		latencyMap: make(map[string]float64),
		log:        log,
	}

	repo.lookupDomains()

	return repo
}

func (r *IPRepository) NetworkTest(downloadClient *HTTPClient, forceRun bool) {
	latencyFile := "data/latency.json"
	if !forceRun {
		latencyData := r.readLatencyFile(latencyFile)
		if latencyData != nil {
			r.latencyMap = *latencyData
			return
		}
	}

	r.log.Info("Network test will start now. Note that it will only run once and record the latency of each domain for future use.")
	r.latencyTest(downloadClient)
	r.log.Infof("Network test completed. Saving the results to %s", latencyFile)
	r.writeLatencyFile(latencyFile)
}

func (r *IPRepository) lookupDomains() {
	limit := 99
	increment := 10
	start := 0
	for {
		lastDomainWorked := 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 {
				hasIPv6 := false
				for _, ip := range ips {
					if ip.To4() == nil {
						hasIPv6 = true
					}
				}
				// assume it always has ipv4
				r.ipv4 = append(r.ipv4, domain)
				if hasIPv6 {
					r.ipv6 = append(r.ipv6, domain)
				}
				if i == limit {
					lastDomainWorked = true
				}
			}

			domain2 := fmt.Sprintf("%d.download.real-debrid.cloud", i)
			ips2, err := net.LookupIP(domain2)
			if err == nil && len(ips2) > 0 {
				hasIPv6 := false
				for _, ip := range ips {
					if ip.To4() == nil {
						hasIPv6 = true
					}
				}
				r.ipv4 = append(r.ipv4, domain2)
				if hasIPv6 {
					r.ipv6 = append(r.ipv6, domain2)
				}
				if i == limit {
					lastDomainWorked = true
				}
			}
		}
		if lastDomainWorked {
			start = limit + 1
			limit += increment
		} else {
			break
		}
	}

	r.log.Infof("Found %d IPv4 domains and %d IPv6 domains", len(r.ipv4), len(r.ipv6))
}

func (r *IPRepository) latencyTest(downloadClient *HTTPClient) {
	const testFileSize = 1
	const iterations = 3

	for _, domain := range r.ipv4 {
		url := fmt.Sprintf("https://%s/speedtest/test.rar/%f", domain, rand.Float64())

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

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

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

			start := time.Now()
			resp, err := downloadClient.Do(req)
			if err != nil {
				r.log.Warnf("Failed to download from %s: %v", domain, err)
				hasError = true
				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)
				hasError = true
				break
			}

			duration := time.Since(start).Seconds()
			totalDuration += duration
		}
		if hasError {
			continue
		}

		r.latencyMap[domain] = totalDuration / 3

		r.log.Debugf("Latency from %s: %.5f seconds", domain, r.latencyMap[domain])
	}
}

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 latencyMap map[string]float64
		if err := json.Unmarshal(jsonData, &latencyMap); err != nil {
			return nil
		}
		return &latencyMap
	}
	return nil
}

func (r *IPRepository) writeLatencyFile(latencyFile string) {
	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(r.latencyMap)
	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
	}
}