elmprodvpn/selective-vpn-api/app/dns_smartdns_handlers_prewarm.go

package app

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"os"
	dnscfgpkg "selective-vpn-api/app/dnscfg"
	"time"
)

// ---------------------------------------------------------------------
// EN: `handleSmartdnsPrewarm` forces DNS lookups for wildcard domains via SmartDNS.
// EN: This warms agvpn_dyn4 in realtime through SmartDNS nftset runtime integration.
// RU: `handleSmartdnsPrewarm` принудительно резолвит wildcard-домены через SmartDNS.
// RU: Это прогревает agvpn_dyn4 в realtime через runtime-интеграцию SmartDNS nftset.
// ---------------------------------------------------------------------
func handleSmartdnsPrewarm(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}
	var body struct {
		Limit          int  `json:"limit"`
		Workers        int  `json:"workers"`
		TimeoutMS      int  `json:"timeout_ms"`
		AggressiveSubs bool `json:"aggressive_subs"`
	}
	if r.Body != nil {
		defer r.Body.Close()
		_ = json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&body)
	}
	writeJSON(w, http.StatusOK, runSmartdnsPrewarm(body.Limit, body.Workers, body.TimeoutMS, body.AggressiveSubs))
}

func runSmartdnsPrewarm(limit, workers, timeoutMS int, aggressiveSubs bool) cmdResult {
	mode := loadDNSMode()
	runtimeEnabled := smartDNSRuntimeEnabled()
	source := "resolver"
	if runtimeEnabled {
		source = "smartdns_runtime"
	}
	smartdnsAddr := normalizeSmartDNSAddr(mode.SmartDNSAddr)
	if smartdnsAddr == "" {
		smartdnsAddr = resolveDefaultSmartDNSAddr()
	}

	aggressive := aggressiveSubs || prewarmAggressiveFromEnv()
	subs := []string{}
	subsPerBaseLimit := 0
	if aggressive {
		subs = loadList(domainDir + "/subs.txt")
		subsPerBaseLimit = envInt("RESOLVE_SUBS_PER_BASE_LIMIT", 0)
		if subsPerBaseLimit < 0 {
			subsPerBaseLimit = 0
		}
	}

	res := dnscfgpkg.RunPrewarm(
		dnscfgpkg.PrewarmInput{
			Mode:             string(mode.Mode),
			Source:           source,
			RuntimeEnabled:   runtimeEnabled,
			SmartDNSAddr:     smartdnsAddr,
			Wildcards:        loadSmartDNSWildcardDomains(nil),
			AggressiveSubs:   aggressive,
			Subs:             subs,
			SubsPerBaseLimit: subsPerBaseLimit,
			Limit:            limit,
			Workers:          workers,
			TimeoutMS:        timeoutMS,
			EnvWorkers:       envInt("SMARTDNS_PREWARM_WORKERS", 24),
			EnvTimeoutMS:     envInt("SMARTDNS_PREWARM_TIMEOUT_MS", 1800),
			MaxHostsLog:      200,
			WildcardMapPath:  lastIPsMapDyn,
		},
		dnscfgpkg.PrewarmDeps{
			IsGoogleLike: isGoogleLike,
			EnsureRuntimeSet: func() {
				_, _, _, _ = runCommandTimeout(5*time.Second, "nft", "add", "table", "inet", "agvpn")
				_, _, _, _ = runCommandTimeout(5*time.Second, "nft", "add", "set", "inet", "agvpn", "agvpn_dyn4", "{", "type", "ipv4_addr", ";", "flags", "interval", ";", "}")
			},
			DigA: func(host string, dnsList []string, timeout time.Duration) ([]string, dnscfgpkg.PrewarmDNSMetrics) {
				ips, stats := digA(host, dnsList, timeout, nil)
				return ips, prewarmMetricsFromDNSMetrics(stats)
			},
			ReadDynSet: func() ([]string, error) {
				return readNftSetElements("agvpn_dyn4")
			},
			ApplyDynSet: func(ips []string) error {
				ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
				defer cancel()
				return nftUpdateSetIPsSmart(ctx, "agvpn_dyn4", ips, nil)
			},
			Logf: func(message string) {
				appendTraceLineTo(smartdnsLogPath, "smartdns", message)
			},
		},
	)

	return cmdResult{
		OK:       res.OK,
		Message:  res.Message,
		ExitCode: res.ExitCode,
	}
}

func prewarmAggressiveFromEnv() bool {
	return dnscfgpkg.SmartDNSForced(os.Getenv("SMARTDNS_PREWARM_AGGRESSIVE"))
}

func prewarmMetricsFromDNSMetrics(in dnsMetrics) dnscfgpkg.PrewarmDNSMetrics {
	out := dnscfgpkg.PrewarmDNSMetrics{
		Attempts:  in.Attempts,
		OK:        in.OK,
		NXDomain:  in.NXDomain,
		Timeout:   in.Timeout,
		Temporary: in.Temporary,
		Other:     in.Other,
		Skipped:   in.Skipped,
	}
	if len(in.PerUpstream) > 0 {
		out.PerUpstream = make(map[string]dnscfgpkg.PrewarmDNSUpstreamMetrics, len(in.PerUpstream))
		for upstream, stats := range in.PerUpstream {
			if stats == nil {
				continue
			}
			out.PerUpstream[upstream] = dnscfgpkg.PrewarmDNSUpstreamMetrics{
				Attempts:  stats.Attempts,
				OK:        stats.OK,
				NXDomain:  stats.NXDomain,
				Timeout:   stats.Timeout,
				Temporary: stats.Temporary,
				Other:     stats.Other,
				Skipped:   stats.Skipped,
			}
		}
	}
	return out
}