Files
zurg/pkg/http/ip.go
2024-08-25 13:29:50 +02:00

287 lines
7.2 KiB
Go

package http
import (
"context"
"encoding/json"
"fmt"
"io"
"math/rand"
"net"
"net/http"
"net/url"
"os"
"sort"
"strings"
"sync"
"time"
"github.com/debridmediamanager/zurg/pkg/logutil"
"github.com/panjf2000/ants/v2"
)
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 by latency from lowest to highest
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() {
var wg sync.WaitGroup
ipv6Enabled := false
if r.canConnectToIPv6() {
r.log.Info("Your network supports IPv6")
ipv6Enabled = true
} else {
r.log.Info("Your network does not support IPv6")
}
p, _ := ants.NewPoolWithFunc(8, func(h interface{}) {
defer wg.Done()
host := h.(string)
ips, err := net.LookupIP(host)
if err == nil && len(ips) > 0 {
latency, err := r.testDomainLatency(r.ipv4client, host)
if err == nil {
r.ipv4latencyMap[host] = latency
r.log.Debugf("Latency from ipv4 %s: %.5f seconds", host, latency)
}
if ipv6Enabled {
latency, err = r.testDomainLatency(r.ipv6client, host)
if err == nil {
r.ipv6latencyMap[host] = latency
r.log.Debugf("Latency from ipv6 %s: %.5f seconds", host, latency)
}
}
if strings.HasSuffix(host, ".download.real-debrid.com") {
ipv4Host := strings.Replace(host, ".download.real-debrid.com", "-4.download.real-debrid.com", 1)
latency, err := r.testDomainLatency(r.ipv4client, ipv4Host)
if err == nil {
r.ipv4latencyMap[ipv4Host] = latency
r.log.Debugf("Latency from ipv4 %s: %.5f seconds", ipv4Host, latency)
}
if ipv6Enabled {
ipv6Host := strings.Replace(host, ".download.real-debrid.com", "-6.download.real-debrid.com", 1)
latency, err = r.testDomainLatency(r.ipv6client, ipv6Host)
if err == nil {
r.ipv6latencyMap[ipv6Host] = latency
r.log.Debugf("Latency from ipv6 %s: %.5f seconds", ipv6Host, latency)
}
}
}
}
})
defer p.Release()
for i := 1; i <= 99; i++ {
wg.Add(2)
_ = p.Invoke(fmt.Sprintf("%d.download.real-debrid.com", i))
_ = p.Invoke(fmt.Sprintf("%d.download.real-debrid.cloud", i))
}
// for i := 'a'; i <= 'z'; i++ {
// for j := 'a'; j <= 'z'; j++ {
// for k := 'a'; k <= 'z'; k++ {
// wg.Add(2)
// _ = p.Invoke(fmt.Sprintf("%c%c%c.download.real-debrid.com", i, j, k))
// _ = p.Invoke(fmt.Sprintf("%c%c%c1.download.real-debrid.com", i, j, k))
// }
// }
// }
wg.Wait()
r.log.Infof("Found %d ipv4 hosts", len(r.ipv4latencyMap))
if ipv6Enabled {
r.log.Infof("Found %d ipv6 hosts", len(r.ipv6latencyMap))
}
}
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
}
}
func (r *IPRepository) canConnectToIPv6() bool {
address := "[2001:4860:4860::8888]:53" // Google Public DNS IPv6 address on port 53
timeout := 5 * time.Second // Timeout duration
conn, err := net.DialTimeout("tcp6", address, timeout)
if err != nil {
fmt.Printf("Failed to connect to IPv6 address %s: %v\n", address, err)
return false
}
defer conn.Close()
fmt.Printf("Successfully connected to IPv6 address %s\n", address)
return true
}