From c20e1caf241545a6bde18d0cb5fa655c9e093e44 Mon Sep 17 00:00:00 2001
From: Aurora Gaffney <aurora@blinklabs.io>
Date: Fri, 9 Aug 2024 20:40:19 -0500
Subject: [PATCH] feat: support for CNAME record fallback (#187)

This adds support for falling back to a matching CNAME record when
there's no available A/AAAA record matching a query.

Fixes #186
---
 internal/dns/dns.go | 60 +++++++++++++++++++++++++--------------------
 1 file changed, 34 insertions(+), 26 deletions(-)

diff --git a/internal/dns/dns.go b/internal/dns/dns.go
index efb7dd2..63ee575 100644
--- a/internal/dns/dns.go
+++ b/internal/dns/dns.go
@@ -85,37 +85,45 @@ func handleQuery(w dns.ResponseWriter, r *dns.Msg) {
 	metricQueryTotal.Inc()
 
 	// Check for known record from local storage
-	records, err := state.GetState().LookupRecords(
-		[]string{dns.Type(r.Question[0].Qtype).String()},
-		strings.TrimSuffix(r.Question[0].Name, "."),
-	)
-	if err != nil {
-		slog.Error(
-			fmt.Sprintf("failed to lookup records in state: %s", err),
+	lookupRecordTypes := []uint16{r.Question[0].Qtype}
+	switch r.Question[0].Qtype {
+	case dns.TypeA, dns.TypeAAAA:
+		// If the query is for A/AAAA, also try looking up matching CNAME records
+		lookupRecordTypes = append(lookupRecordTypes, dns.TypeCNAME)
+	}
+	for _, lookupRecordType := range lookupRecordTypes {
+		records, err := state.GetState().LookupRecords(
+			[]string{dns.Type(lookupRecordType).String()},
+			strings.TrimSuffix(r.Question[0].Name, "."),
 		)
-		return
-	}
-	if records != nil {
-		// Assemble response
-		m.SetReply(r)
-		for _, tmpRecord := range records {
-			tmpRR, err := stateRecordToDnsRR(tmpRecord)
-			if err != nil {
+		if err != nil {
+			slog.Error(
+				fmt.Sprintf("failed to lookup records in state: %s", err),
+			)
+			return
+		}
+		if records != nil {
+			// Assemble response
+			m.SetReply(r)
+			for _, tmpRecord := range records {
+				tmpRR, err := stateRecordToDnsRR(tmpRecord)
+				if err != nil {
+					slog.Error(
+						fmt.Sprintf("failed to convert state record to dns.RR: %s", err),
+					)
+					return
+				}
+				m.Answer = append(m.Answer, tmpRR)
+			}
+			// Send response
+			if err := w.WriteMsg(m); err != nil {
 				slog.Error(
-					fmt.Sprintf("failed to convert state record to dns.RR: %s", err),
+					fmt.Sprintf("failed to write response: %s", err),
 				)
-				return
 			}
-			m.Answer = append(m.Answer, tmpRR)
-		}
-		// Send response
-		if err := w.WriteMsg(m); err != nil {
-			slog.Error(
-				fmt.Sprintf("failed to write response: %s", err),
-			)
+			// We found our answer, to return from handler
+			return
 		}
-		// We found our answer, to return from handler
-		return
 	}
 
 	// Check for any NS records for parent domains from local storage