zurg/pkg/http/ip.go

package http

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

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

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

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

	return repo
}

func (r *IPRepository) NetworkTest(forceRun bool) {
	ipv4latencyFile := "data/latency4.json"
	ipv6latencyFile := "data/latency6.json"
	if !forceRun {
		latencyData := r.readLatencyFile(ipv4latencyFile)
		if latencyData != nil {
			r.ipv4latencyMap = *latencyData
		}
		latencyData = r.readLatencyFile(ipv6latencyFile)
		if latencyData != nil {
			r.ipv6latencyMap = *latencyData
		}
	}

	r.log.Info("Network test will start now. IGNORE THE WARNINGS!")
	r.runLatencyTest()
	r.log.Infof("Network test completed. Saving the results to %s and %s", ipv4latencyFile, ipv6latencyFile)
	r.log.Debugf("ipv4 %v", r.ipv4latencyMap)
	r.log.Debugf("ipv6 %v", r.ipv6latencyMap)
	r.writeLatencyFile(ipv4latencyFile, r.ipv4latencyMap)
	r.writeLatencyFile(ipv6latencyFile, r.ipv6latencyMap)
}

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
	url := fmt.Sprintf("https://%s/speedtest/test.rar/%f", domain, rand.Float64())

	var totalDuration float64
	var retErr error
	for i := 0; i < iterations; i++ {
		ctx, cancel := context.WithTimeout(context.Background(), 1*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)
			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
	}
}