From f264bd5604e593ecc9127e043be7a4a9884c1810 Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Sat, 13 Jan 2024 17:30:55 +0100
Subject: [PATCH] feat: Allow usage for top-level domains

---
 internal/dns/dns.go      | 46 +++++++++++++++++++++++++++++++++++++++-
 internal/dns/dns_test.go | 32 ++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 internal/dns/dns_test.go

diff --git a/internal/dns/dns.go b/internal/dns/dns.go
index 44e3285..5a8fc80 100644
--- a/internal/dns/dns.go
+++ b/internal/dns/dns.go
@@ -2,6 +2,7 @@ package dns
 
 import (
 	"net"
+	"errors"
 	"strings"
 	"time"
 
@@ -15,6 +16,10 @@ const (
 	// The key that the TXT record will have to start with, e.g.
 	// "repo=some-random-repo".
 	TxtRepoKey = "repo="
+
+	TxtCNAMERecord = "_rio-cname."
+
+	TxtCNAMEKey = "cname="
 )
 
 var (
@@ -61,14 +66,53 @@ func LookupCNAME(domain string) (string, error) {
 		return cname.(string), nil
 	}
 
+	cname, err := lookupCNAME(domain)
+	if err == nil {
+		cnameCache.Set(domain, cname, cache.DefaultExpiration)
+		return cname.(string), nil
+	}
+
+	altCname, err := lookupCNAMETxt(domain)
+	if err == nil {
+		cnameCache.Set(domain, altCname, cache.DefaultExpiration)
+		return altCname, nil
+	}
+
+	return "", err
+}
+
+// Lookup the CNAME by trying to find a CNAME RR. Contrary to net.LookupCNAME,
+// this method fails if we find no CNAME RR.
+func lookupCNAME(domain string) (string, error) {
 	query, err := net.LookupCNAME(domain)
 	if err == nil {
 		if query[len(query)-1] == '.' {
 			query = query[:len(query)-1]
 		}
-		cnameCache.Set(domain, query, cache.DefaultExpiration)
+
+		// Fail if we have no CNAME RR.
+		if query == domain {
+			return "", errors.New("CNAME is equal to domain")
+		}
+
 		return query, nil
 	}
 
 	return "", err
 }
+
+// Performs an alternative CNAME lookup by looking for a special TXT record.
+func lookupCNAMETxt(domain string) (string, error) {
+	txts, err := net.LookupTXT(TxtCNAMERecord+domain)
+	if err == nil {
+		for _, txt := range txts {
+			if !strings.HasPrefix(txt, TxtCNAMEKey) {
+				continue
+			}
+
+			return strings.TrimPrefix(txt, TxtCNAMEKey), nil
+		}
+	}
+	
+	return "", err
+}
\ No newline at end of file
diff --git a/internal/dns/dns_test.go b/internal/dns/dns_test.go
new file mode 100644
index 0000000..7b303b1
--- /dev/null
+++ b/internal/dns/dns_test.go
@@ -0,0 +1,32 @@
+package dns
+
+import "testing"
+
+func cleanCache() {
+	cnameCache.Flush()
+	txtRepoCache.Flush()
+}
+
+func TestAltCNAME(t *testing.T) {
+	defer cleanCache()
+
+	res, err := lookupCNAMETxt("moxxy.org")
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if res != "moxxy.pages.polynom.me" {
+		t.Fatalf("Unexpected alt CNAME: %s", res)
+	}
+}
+
+func TestLookupCNAMEWithAlt(t *testing.T) {
+	defer cleanCache()
+
+	res, err := LookupCNAME("moxxy.org")
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if res != "moxxy.pages.polynom.me" {
+		t.Fatalf("Unexpected alt CNAME: %s", res)
+	}
+}
\ No newline at end of file