EvoBGP/internal/pipeline/doh_resolve_test.go

package pipeline

import (
	"context"
	"encoding/base64"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"

	"evobgp/internal/store"

	"github.com/miekg/dns"
)

// TestResolveDomainIPs_ParallelAAndAAAA verifies that A and AAAA queries are issued
// concurrently: with a 250ms upstream latency the combined resolve must stay near
// one round-trip instead of two.
func TestResolveDomainIPs_ParallelAAndAAAA(t *testing.T) {
	const delay = 250 * time.Millisecond
	var inflight, maxInflight atomic.Int32
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cur := inflight.Add(1)
		for {
			old := maxInflight.Load()
			if cur <= old || maxInflight.CompareAndSwap(old, cur) {
				break
			}
		}
		defer inflight.Add(-1)
		time.Sleep(delay)

		if wire := r.URL.Query().Get("dns"); wire != "" {
			// RFC8484 dns-message: decode the query and answer on the wire.
			raw, err := base64.RawURLEncoding.DecodeString(wire)
			if err != nil {
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			msg := new(dns.Msg)
			if err := msg.Unpack(raw); err != nil {
				w.WriteHeader(http.StatusBadRequest)
				return
			}
			resp := new(dns.Msg)
			resp.SetReply(msg)
			switch msg.Question[0].Qtype {
			case dns.TypeA:
				resp.Answer = append(resp.Answer, &dns.A{
					Hdr: dns.RR_Header{Name: msg.Question[0].Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60},
					A:   []byte{203, 0, 113, 10},
				})
			case dns.TypeAAAA:
				resp.Answer = append(resp.Answer, &dns.AAAA{
					Hdr:  dns.RR_Header{Name: msg.Question[0].Name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 60},
					AAAA: []byte{0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x10},
				})
			}
			out, err := resp.Pack()
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				return
			}
			w.Header().Set("Content-Type", "application/dns-message")
			_, _ = w.Write(out)
			return
		}
		w.Header().Set("Content-Type", "application/dns-json")
		switch r.URL.Query().Get("type") {
		case "A":
			_, _ = w.Write([]byte(`{"Status":0,"Answer":[{"type":1,"data":"203.0.113.10"}]}`))
		default:
			_, _ = w.Write([]byte(`{"Status":0,"Answer":[{"type":28,"data":"2001:db8::10"}]}`))
		}
	}))
	defer srv.Close()

	prof := &store.DohProfile{URL: srv.URL, TimeoutMs: ptrInt(5000)}
	ctx := context.Background()
	start := time.Now()
	ips, err := resolveDomainIPsNoSystemFallback(ctx, srv.Client(), prof, "example.test")
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("resolve failed: %v", err)
	}
	if len(ips) != 2 {
		t.Fatalf("expected 2 addrs, got %+v", ips)
	}
	if maxInflight.Load() < 2 {
		t.Fatalf("expected concurrent A/AAAA queries, max inflight=%d", maxInflight.Load())
	}
	if elapsed >= 2*delay {
		t.Fatalf("resolve took %v; expected one round-trip (<2*%v)", elapsed, delay)
	}
}

func ptrInt(v int) *int { return &v }