feat: implement CLI flags and configuration handling (HAN-382) (#8)

internal/dns/query.go

@@ -10,18 +10,20 @@ import (
 )
 
 type QueryConfig struct {
-	UDPSize int
-	Timeout time.Duration
-	Retries int
-	UseTCP  bool
+	UDPSize  int
+	Timeout  time.Duration
+	Retries  int
+	UseTCP   bool
+	AllowTCP bool
 }
 
 func DefaultQueryConfig() *QueryConfig {
 	return &QueryConfig{
-		UDPSize: DefaultEDNS0UDPSize(),
-		Timeout: 5 * time.Second,
-		Retries: 3,
-		UseTCP:  false,
+		UDPSize:  DefaultEDNS0UDPSize(),
+		Timeout:  5 * time.Second,
+		Retries:  3,
+		UseTCP:   false,
+		AllowTCP: true,
 	}
 }
 
@@ -110,7 +112,7 @@ func QueryWithExchange(ctx context.Context, server net.IP, name string, qtype ui
 			continue
 		}
 
-		if resp.Truncated {
+		if resp.Truncated && cfg.AllowTCP {
 			resp, err = exchangeFn(ctx, serverStr, msg, true)
 			if err != nil {
 				lastErr = err
@@ -178,7 +180,7 @@ func IterativeQueryWithExchange(ctx context.Context, server net.IP, name string,
 			continue
 		}
 
-		if resp.Truncated {
+		if resp.Truncated && cfg.AllowTCP {
 			resp, err = exchangeFn(ctx, serverStr, msg, true)
 			if err != nil {
 				lastErr = err

internal/dns/query_test.go

@@ -110,10 +110,11 @@ func TestQueryTCPFallbackOnTruncation(t *testing.T) {
 	}
 
 	cfg := &QueryConfig{
-		UDPSize: 2048,
-		Timeout: 5,
-		Retries: 1,
-		UseTCP:  false,
+		UDPSize:  2048,
+		Timeout:  5,
+		Retries:  1,
+		UseTCP:   false,
+		AllowTCP: true,
 	}
 
 	server := net.ParseIP("8.8.8.8")
@@ -272,10 +273,11 @@ func TestQueryTCPFallbackFailsThenRetries(t *testing.T) {
 	}
 
 	cfg := &QueryConfig{
-		UDPSize: 2048,
-		Timeout: 5,
-		Retries: 2,
-		UseTCP:  false,
+		UDPSize:  2048,
+		Timeout:  5,
+		Retries:  2,
+		UseTCP:   false,
+		AllowTCP: true,
 	}
 
 	server := net.ParseIP("8.8.8.8")
@@ -287,3 +289,43 @@ func TestQueryTCPFallbackFailsThenRetries(t *testing.T) {
 		t.Errorf("expected 4 calls (2 retries x UDP+TCP), got %d", callCount)
 	}
 }
+
+func TestQueryNoTCPFallbackWhenDisabled(t *testing.T) {
+	truncatedResp := new(dns.Msg)
+	truncatedResp.Truncated = true
+	truncatedResp.SetReply(new(dns.Msg))
+
+	var mu sync.Mutex
+	calls := []bool{}
+	exchangeFn := func(ctx context.Context, server string, msg *dns.Msg, useTCP bool) (*dns.Msg, error) {
+		mu.Lock()
+		defer mu.Unlock()
+		calls = append(calls, useTCP)
+		return truncatedResp.Copy(), nil
+	}
+
+	cfg := &QueryConfig{
+		UDPSize:  2048,
+		Timeout:  5,
+		Retries:  1,
+		UseTCP:   false,
+		AllowTCP: false, // TCP fallback must be suppressed
+	}
+
+	server := net.ParseIP("8.8.8.8")
+	resp, err := QueryWithExchange(context.Background(), server, "example.com", TypeA, cfg, exchangeFn)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	// Only one UDP call; no TCP fallback.
+	if len(calls) != 1 {
+		t.Fatalf("expected 1 exchange call (no TCP fallback), got %d", len(calls))
+	}
+	if calls[0] != false {
+		t.Error("expected UDP-only call")
+	}
+	if !resp.Truncated {
+		t.Error("expected truncated response to be returned as-is")
+	}
+}
+