Implement domain handling logic

Still lots of performance optimization required!

domains.go

@@ -1,15 +1,94 @@
 package main
 
-import "github.com/valyala/fasthttp"
+import (
+	"github.com/OrlovEvgeny/go-mcache"
+	"github.com/valyala/fasthttp"
+	"net"
+	"strings"
+	"time"
+)
 
-// getTargetFromDNS searches for CNAME entries on the request domain, optionally with a "www." prefix, and checks if
-// the domain is included in the repository's "domains.txt" file. If everything is fine, it returns the target data.
-// TODO: use TXT records with A/AAAA/ALIAS
-func getTargetFromDNS(ctx *fasthttp.RequestCtx) (targetOwner, targetRepo, targetBranch, targetPath string) {
-	// TODO: read CNAME record for host and "www.{host}" to get those values
-	// TODO: check domains.txt
+// DnsLookupCacheTimeout specifies the timeout for the DNS lookup cache.
+var DnsLookupCacheTimeout = 15*time.Minute
+// dnsLookupCache stores DNS lookups for custom domains
+var dnsLookupCache = mcache.New()
+
+// getTargetFromDNS searches for CNAME or TXT entries on the request domain ending with MainDomainSuffix, and checks if
+// the domain equals the repository's ".canonical-domain" file. If everything is fine, it returns the target data.
+func getTargetFromDNS(domain string) (targetOwner, targetRepo, targetBranch string) {
+	// Get CNAME or TXT
+	var cname string
+	var err error
+	if cachedName, ok := dnsLookupCache.Get(domain); ok {
+		cname = cachedName.(string)
+	} else {
+		cname, err = net.LookupCNAME(domain)
+		cname = strings.TrimSuffix(cname, ".")
+		if err != nil || !strings.HasSuffix(cname, string(MainDomainSuffix)) {
+			cname = ""
+			names, err := net.LookupTXT(domain)
+			if err == nil {
+				for _, name := range names {
+					name = strings.TrimSuffix(name, ".")
+					if strings.HasSuffix(name, string(MainDomainSuffix)) {
+						cname = name
+						break
+					}
+				}
+			}
+		}
+		_ = dnsLookupCache.Set(domain, cname, DnsLookupCacheTimeout)
+	}
+	if cname == "" {
+		return
+	}
+	cnameParts := strings.Split(strings.TrimSuffix(cname, string(MainDomainSuffix)), ".")
+	targetOwner = cnameParts[len(cnameParts)-1]
+	if len(cnameParts) > 1 {
+		targetRepo = cnameParts[len(cnameParts)-1]
+	}
+	if len(cnameParts) > 2 {
+		targetBranch = cnameParts[len(cnameParts)-2]
+	}
+	if targetRepo == "" {
+		targetRepo = "pages"
+	}
+	if targetBranch == "" && targetRepo != "pages" {
+		targetBranch = "pages"
+	}
+	// if targetBranch is still empty, the caller must find the default branch
 	return
 }
 
-// TODO: cache domains.txt for 15 minutes
-// TODO: canonical domains - redirect to first domain if domains.txt exists, also make sure owner.codeberg.page/pages/... redirects to /...
+
+// CanonicalDomainCacheTimeout specifies the timeout for the canonical domain cache.
+var CanonicalDomainCacheTimeout = 15*time.Minute
+// canonicalDomainCache stores canonical domains
+var canonicalDomainCache = mcache.New()
+
+// checkCanonicalDomain returns the canonical domain specified in the repo (using the file `.canonical-domain`).
+func checkCanonicalDomain(targetOwner, targetRepo, targetBranch string) (canonicalDomain string) {
+	// Check if the canonical domain matches
+	req := fasthttp.AcquireRequest()
+	req.SetRequestURI(string(GiteaRoot) + "/api/v1/repos/" + targetOwner + "/" + targetRepo + "/raw/" + targetBranch + "/.canonical-domain")
+	res := fasthttp.AcquireResponse()
+	if cachedValue, ok := canonicalDomainCache.Get(string(req.RequestURI())); ok {
+		canonicalDomain = cachedValue.(string)
+	} else {
+		err := upstreamClient.Do(req, res)
+		if err == nil && res.StatusCode() == fasthttp.StatusOK {
+			canonicalDomain = strings.TrimSpace(string(res.Body()))
+			if strings.Contains(canonicalDomain, "/") {
+				canonicalDomain = ""
+			}
+		}
+		if canonicalDomain == "" {
+			canonicalDomain = targetOwner + string(MainDomainSuffix)
+			if targetRepo != "" && targetRepo != "pages" {
+				canonicalDomain += "/" + targetRepo
+			}
+		}
+		_ = canonicalDomainCache.Set(string(req.RequestURI()), canonicalDomain, CanonicalDomainCacheTimeout)
+	}
+	return
+}

handler.go

@@ -98,6 +98,19 @@ func handler(ctx *fasthttp.RequestCtx) {
 
 	// tryUpstream forwards the target request to the Gitea API, and shows an error page on failure.
 	var tryUpstream = func() {
+		// check if a canonical domain exists on a request on MainDomain
+		if bytes.HasSuffix(ctx.Request.Host(), MainDomainSuffix) {
+			canonicalDomain := checkCanonicalDomain(targetOwner, targetRepo, targetBranch)
+			if !strings.HasSuffix(strings.SplitN(canonicalDomain, "/", 2)[0], string(MainDomainSuffix)) {
+				canonicalPath := string(ctx.RequestURI())
+				if targetRepo != "pages" {
+					canonicalPath = "/" + strings.SplitN(canonicalPath, "/", 3)[2]
+				}
+				ctx.Redirect("https://" + canonicalDomain + canonicalPath, fasthttp.StatusTemporaryRedirect)
+				return
+			}
+		}
+
 		// Try to request the file from the Gitea API
 		if !upstream(ctx, targetOwner, targetRepo, targetBranch, targetPath, targetOptions) {
 			returnErrorPage(ctx, ctx.Response.StatusCode())
@@ -151,7 +164,7 @@ func handler(ctx *fasthttp.RequestCtx) {
 		if len(pathElements) > 1 && strings.HasPrefix(pathElements[1], "@") {
 			if targetRepo == "pages" {
 				// example.codeberg.org/pages/@... redirects to example.codeberg.org/@...
-				ctx.Redirect("/" + strings.Join(pathElements[1:], "/"), fasthttp.StatusMovedPermanently)
+				ctx.Redirect("/" + strings.Join(pathElements[1:], "/"), fasthttp.StatusTemporaryRedirect)
 				return
 			}
 
@@ -196,12 +209,41 @@ func handler(ctx *fasthttp.RequestCtx) {
 		return
 	} else {
 		// Serve pages from external domains
-
-		targetOwner, targetRepo, targetBranch, targetPath = getTargetFromDNS(ctx)
+		targetOwner, targetRepo, targetBranch = getTargetFromDNS(string(ctx.Request.Host()))
 		if targetOwner == "" {
 			ctx.Redirect(BrokenDNSPage, fasthttp.StatusTemporaryRedirect)
 			return
 		}
+
+		pathElements := strings.Split(string(bytes.Trim(ctx.Request.URI().Path(), "/")), "/")
+		canonicalLink := ""
+		if strings.HasPrefix(pathElements[0], "@") {
+			targetBranch = pathElements[0][1:]
+			pathElements = pathElements[1:]
+			canonicalLink = "/%p"
+		}
+
+		// Try to use the given repo on the given branch or the default branch
+		if tryBranch(targetRepo, targetBranch, pathElements, canonicalLink) {
+			canonicalDomain := checkCanonicalDomain(targetOwner, targetRepo, targetBranch)
+			if canonicalDomain != string(ctx.Request.Host()) {
+				// only redirect if
+				targetOwner, _, _ = getTargetFromDNS(strings.SplitN(canonicalDomain, "/", 2)[0])
+				if targetOwner != "" {
+					ctx.Redirect("https://"+canonicalDomain+string(ctx.RequestURI()), fasthttp.StatusTemporaryRedirect)
+					return
+				} else {
+					ctx.Redirect(BrokenDNSPage, fasthttp.StatusTemporaryRedirect)
+					return
+				}
+			}
+
+			tryUpstream()
+			return
+		} else {
+			returnErrorPage(ctx, fasthttp.StatusFailedDependency)
+			return
+		}
 	}
 }