EvoBGP/internal/pipeline/doh_resolve.go

package pipeline

import (
	"context"
	"fmt"
	"net/http"
	"net/netip"
	"strings"
	"sync"
	"time"

	"evobgp/internal/store"

	"github.com/miekg/dns"
)

func loadModuleDohProfiles(st store.Backend, tenantID string, mod *store.Module) ([]*store.DohProfile, string, error) {
	if mod == nil {
		return nil, store.DohPolicyPrimaryOnly, nil
	}
	policy := store.NormalizeDohResolverPolicy(mod.DohResolverPolicy)
	var profiles []*store.DohProfile
	for _, id := range mod.EffectiveDohProfileIDs() {
		prof, err := st.GetDohProfile(tenantID, id)
		if err != nil {
			return nil, "", fmt.Errorf("get doh profile %s: %w", id, err)
		}
		profiles = append(profiles, prof)
	}
	return profiles, policy, nil
}

func resolveDomainIPsWithPolicy(ctx context.Context, hc *http.Client, profiles []*store.DohProfile, policy, fqdn string) ([]netip.Addr, error) {
	policy = store.NormalizeDohResolverPolicy(policy)
	if len(profiles) == 0 {
		return resolveDomainIPs(ctx, hc, nil, fqdn)
	}
	if len(profiles) == 1 {
		return resolveDomainIPs(ctx, hc, profiles[0], fqdn)
	}

	switch policy {
	case store.DohPolicyUnion:
		return resolveDomainIPsUnion(ctx, hc, profiles, fqdn)
	case store.DohPolicyFailover:
		return resolveDomainIPsFailover(ctx, hc, profiles, fqdn)
	default:
		return resolveDomainIPs(ctx, hc, profiles[0], fqdn)
	}
}

func resolveDomainIPsUnion(ctx context.Context, hc *http.Client, profiles []*store.DohProfile, fqdn string) ([]netip.Addr, error) {
	var merged []netip.Addr
	var errs []error
	for _, prof := range profiles {
		if prof == nil {
			continue
		}
		ips, err := resolveDomainIPsNoSystemFallback(ctx, hc, prof, fqdn)
		if err != nil {
			errs = append(errs, fmt.Errorf("%s: %w", strings.TrimSpace(prof.URL), err))
			continue
		}
		merged = append(merged, ips...)
	}
	merged = uniqAddrs(merged)
	if len(merged) > 0 {
		return merged, nil
	}
	if len(errs) > 0 {
		return nil, fmt.Errorf("doh union failed: %v", errs)
	}
	return nil, nil
}

func resolveDomainIPsFailover(ctx context.Context, hc *http.Client, profiles []*store.DohProfile, fqdn string) ([]netip.Addr, error) {
	var lastErr error
	for _, prof := range profiles {
		if prof == nil {
			continue
		}
		ips, err := resolveDomainIPsNoSystemFallback(ctx, hc, prof, fqdn)
		if err != nil {
			lastErr = err
			continue
		}
		if len(ips) > 0 {
			return ips, nil
		}
	}
	if lastErr != nil {
		return nil, lastErr
	}
	return resolveDomainIPs(ctx, hc, nil, fqdn)
}

// resolveDomainIPsNoSystemFallback queries one DoH profile without falling back to OS resolver.
func resolveDomainIPsNoSystemFallback(ctx context.Context, hc *http.Client, profile *store.DohProfile, fqdn string) ([]netip.Addr, error) {
	host := strings.TrimSpace(strings.TrimSuffix(fqdn, "."))
	if host == "" {
		return nil, nil
	}
	if profile == nil || strings.TrimSpace(profile.URL) == "" {
		return nil, fmt.Errorf("empty doh profile")
	}

	timeout := dohProfileTimeout(profile)
	dctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	baseURL := strings.TrimSpace(profile.URL)
	v4, v6, err4, err6 := resolveDOHMessagePair(dctx, hc, baseURL, host)
	if err4 != nil {
		v4, err4 = resolveDomainWithDOHJSON(dctx, hc, baseURL, host, "A")
	}
	if err6 != nil {
		v6, err6 = resolveDomainWithDOHJSON(dctx, hc, baseURL, host, "AAAA")
	}
	if err4 != nil && err6 != nil {
		return nil, fmt.Errorf("doh failed for A and AAAA: %v; %v", err4, err6)
	}
	return uniqAddrs(append(v4, v6...)), nil
}

// resolveDOHMessagePair issues RFC8484 dns-message A and AAAA queries concurrently
// and waits for both (fallbacks are handled by the caller).
func resolveDOHMessagePair(ctx context.Context, hc *http.Client, baseURL, host string) (v4, v6 []netip.Addr, err4, err6 error) {
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		v4, err4 = resolveDomainWithDOHMessage(ctx, hc, baseURL, host, dns.TypeA)
	}()
	go func() {
		defer wg.Done()
		v6, err6 = resolveDomainWithDOHMessage(ctx, hc, baseURL, host, dns.TypeAAAA)
	}()
	wg.Wait()
	return v4, v6, err4, err6
}

func dohProfileTimeout(profile *store.DohProfile) time.Duration {
	timeout := 10 * time.Second
	if profile != nil && profile.TimeoutMs != nil && *profile.TimeoutMs > 0 {
		timeout = time.Duration(*profile.TimeoutMs) * time.Millisecond
	}
	return timeout
}