From 25eb0de1e7dec06e4c72b418feb8382aead3cdd3 Mon Sep 17 00:00:00 2001
From: "Alexander \"PapaTutuWawa" <papatutuwawa@polynom.me>
Date: Thu, 11 Jan 2024 20:37:56 +0100
Subject: [PATCH] fix: Handle usernames with a dot in them

Fixes #1.
---
 .gitignore           |  2 +-
 cmd/rio/main.go      | 22 +++++++++++++++++-----
 cmd/rio/main_test.go | 23 +++++++++++++++++++++++
 3 files changed, 41 insertions(+), 6 deletions(-)
 create mode 100644 cmd/rio/main_test.go

diff --git a/.gitignore b/.gitignore
index 6693397..c997da2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,5 @@
 # Artificats
-rio
+/rio
 
 # Testing stuff
 *.json
diff --git a/cmd/rio/main.go b/cmd/rio/main.go
index aa6e170..7bdcd0b 100644
--- a/cmd/rio/main.go
+++ b/cmd/rio/main.go
@@ -24,14 +24,26 @@ import (
 	"github.com/urfave/cli/v2"
 )
 
-func handleSubdomain(domain, cname, path, giteaUrl, defaultCsp string, giteaClient *repo.GiteaClient, w http.ResponseWriter) {
+// 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 = strings.Split(cname, ".")[0]
+		username = extractUsername(pagesDomain, cname)
 	} else {
 		// If we are directly accessed, then domain contains our <user>.<pages domain> value.
-		username = strings.Split(domain, ".")[0]
+		username = extractUsername(pagesDomain, domain)
 	}
 
 	// Strip the leading /
@@ -79,7 +91,7 @@ func Handler(pagesDomain, giteaUrl, defaultCsp string, giteaClient *repo.GiteaCl
 		// Is a direct subdomain requested?
 		if strings.HasSuffix(req.Host, pagesDomain) {
 			log.Debug("Domain can be directly handled")
-			handleSubdomain(req.Host, "", req.URL.Path, giteaUrl, defaultCsp, giteaClient, w)
+			handleSubdomain(pagesDomain, req.Host, "", req.URL.Path, giteaUrl, defaultCsp, giteaClient, w)
 			return
 		}
 
@@ -96,7 +108,7 @@ func Handler(pagesDomain, giteaUrl, defaultCsp string, giteaClient *repo.GiteaCl
 		//       pages domain makes no sense.
 		if strings.HasSuffix(cname, "."+pagesDomain) {
 			log.Debugf("%s is alias of %s and can be handled after a CNAME query", req.Host, cname)
-			handleSubdomain(req.Host, cname, req.URL.Path, giteaUrl, defaultCsp, giteaClient, w)
+			handleSubdomain(pagesDomain, req.Host, cname, req.URL.Path, giteaUrl, defaultCsp, giteaClient, w)
 			return
 		}
 
diff --git a/cmd/rio/main_test.go b/cmd/rio/main_test.go
new file mode 100644
index 0000000..19ab538
--- /dev/null
+++ b/cmd/rio/main_test.go
@@ -0,0 +1,23 @@
+package main
+
+import "testing"
+
+func TestExtractUsernameSimple(t *testing.T) {
+	username := extractUsername(
+		"pages.local",
+		"papatutuwawa.pages.local",
+	)
+	if username != "papatutuwawa" {
+		t.Fatalf("Unexpected username: '%s'", username)
+	}
+}
+
+func TestExtractUsernameDot(t *testing.T) {
+	username := extractUsername(
+		"pages.local",
+		"polynom.me.pages.local",
+	)
+	if username != "polynom.me" {
+		t.Fatalf("Unexpected username: '%s'", username)
+	}
+}