From 9b971daf282dc792323e4818b76552a00a874070 Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Thu, 11 Jan 2024 22:10:32 +0100
Subject: [PATCH] fix: Fix username extraction in the TLS handler

---
 cmd/rio/main.go                                  | 16 ++--------------
 internal/dns/username.go                         | 15 +++++++++++++++
 .../dns/username_test.go                         |  6 +++---
 internal/server/tls.go                           | 10 ++--------
 4 files changed, 22 insertions(+), 25 deletions(-)
 create mode 100644 internal/dns/username.go
 rename cmd/rio/main_test.go => internal/dns/username_test.go (83%)

diff --git a/cmd/rio/main.go b/cmd/rio/main.go
index 7bdcd0b..8d0caf6 100644
--- a/cmd/rio/main.go
+++ b/cmd/rio/main.go
@@ -24,26 +24,14 @@ import (
 	"github.com/urfave/cli/v2"
 )
 
-// Extract the username from the domain name @domain that we're processing
-// at the moment.
-func extractUsername(pagesDomain, domain string) string {
-	suffixlessDomain := strings.TrimSuffix(domain, "."+pagesDomain)
-	usernameParts := strings.Split(suffixlessDomain, ".")
-	if len(usernameParts) == 1 {
-		return usernameParts[0]
-	}
-
-	return strings.Join(usernameParts, ".")
-}
-
 func handleSubdomain(pagesDomain, domain, cname, path, giteaUrl, defaultCsp string, giteaClient *repo.GiteaClient, w http.ResponseWriter) {
 	username := ""
 	if cname != "" {
 		// If we are accessed via a CNAME, then CNAME contains our <user>.<pages domain> value.
-		username = extractUsername(pagesDomain, cname)
+		username = dns.ExtractUsername(pagesDomain, cname)
 	} else {
 		// If we are directly accessed, then domain contains our <user>.<pages domain> value.
-		username = extractUsername(pagesDomain, domain)
+		username = dns.ExtractUsername(pagesDomain, domain)
 	}
 
 	// Strip the leading /
diff --git a/internal/dns/username.go b/internal/dns/username.go
new file mode 100644
index 0000000..3d995e4
--- /dev/null
+++ b/internal/dns/username.go
@@ -0,0 +1,15 @@
+package dns
+
+import "strings"
+
+// Extract the username from the domain name @domain that we're processing
+// at the moment.
+func ExtractUsername(pagesDomain, domain string) string {
+	suffixlessDomain := strings.TrimSuffix(domain, "."+pagesDomain)
+	usernameParts := strings.Split(suffixlessDomain, ".")
+	if len(usernameParts) == 1 {
+		return usernameParts[0]
+	}
+
+	return strings.Join(usernameParts, ".")
+}
diff --git a/cmd/rio/main_test.go b/internal/dns/username_test.go
similarity index 83%
rename from cmd/rio/main_test.go
rename to internal/dns/username_test.go
index 19ab538..e605ea4 100644
--- a/cmd/rio/main_test.go
+++ b/internal/dns/username_test.go
@@ -1,9 +1,9 @@
-package main
+package dns
 
 import "testing"
 
 func TestExtractUsernameSimple(t *testing.T) {
-	username := extractUsername(
+	username := ExtractUsername(
 		"pages.local",
 		"papatutuwawa.pages.local",
 	)
@@ -13,7 +13,7 @@ func TestExtractUsernameSimple(t *testing.T) {
 }
 
 func TestExtractUsernameDot(t *testing.T) {
-	username := extractUsername(
+	username := ExtractUsername(
 		"pages.local",
 		"polynom.me.pages.local",
 	)
diff --git a/internal/server/tls.go b/internal/server/tls.go
index 7474563..ae500f7 100644
--- a/internal/server/tls.go
+++ b/internal/server/tls.go
@@ -44,7 +44,6 @@ func unlockDomain(domain string) {
 }
 
 func buildDomainList(domain, pagesDomain string) []string {
-	// TODO: For wildcards, we MUST use DNS01
 	if domain == pagesDomain || strings.HasSuffix(domain, pagesDomain) {
 		return []string{
 			pagesDomain,
@@ -56,7 +55,6 @@ func buildDomainList(domain, pagesDomain string) []string {
 }
 
 func getDomainKey(domain, pagesDomain string) string {
-	// TODO: For wildcards, we MUST use DNS01
 	if domain == pagesDomain || strings.HasSuffix(domain, pagesDomain) {
 		return "*." + pagesDomain
 	}
@@ -64,10 +62,6 @@ func getDomainKey(domain, pagesDomain string) string {
 	return domain
 }
 
-func usernameFromDomain(domain string) string {
-	return strings.Split(domain, ".")[0]
-}
-
 func getUsername(sni, pagesDomain string) (string, error) {
 	if !strings.HasSuffix(sni, pagesDomain) {
 		log.Debugf("'%s' is not a subdomain of '%s'", sni, pagesDomain)
@@ -81,10 +75,10 @@ func getUsername(sni, pagesDomain string) (string, error) {
 			return "", errors.New("CNAME does not resolve to subdomain of pages domain")
 		}
 
-		return usernameFromDomain(query), nil
+		return dns.ExtractUsername(pagesDomain, query), nil
 	}
 
-	return usernameFromDomain(sni), nil
+	return dns.ExtractUsername(pagesDomain, sni), nil
 }
 
 func MakeTlsConfig(pagesDomain, cachePath string, cache *certificates.CertificatesCache, acmeClient *lego.Client, giteaClient *repo.GiteaClient) *tls.Config {