add support for multipart domains and TLDs (#283)

puma · Oct 10, 2021 · 03601ea · 03601ea
1 parent 5e49c47
commit 03601ea
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 0 deletions.
diff --git a/cmd/puma-dev/main_darwin.go b/cmd/puma-dev/main_darwin.go
@@ -179,6 +179,7 @@ func main() {
 	http.Pool = &pool
 	http.Debug = *fDebug
 	http.Events = &events
+	http.Domains = domains
 	if len(*fNoServePublicPaths) > 0 {
 		http.IgnoredStaticPaths = strings.Split(*fNoServePublicPaths, ":")
 		fmt.Printf("* Ignoring files under: public{%s}\n", strings.Join(http.IgnoredStaticPaths, ", "))

diff --git a/cmd/puma-dev/main_linux.go b/cmd/puma-dev/main_linux.go
@@ -100,6 +100,7 @@ func main() {
 	http.Pool = &pool
 	http.Debug = *fDebug
 	http.Events = &events
+	http.Domains = domains
 	if len(*fNoServePublicPaths) > 0 {
 		http.IgnoredStaticPaths = strings.Split(*fNoServePublicPaths, ":")
 		fmt.Printf("* Ignoring files under: public{%s}\n", strings.Join(http.IgnoredStaticPaths, ", "))

diff --git a/dev/http.go b/dev/http.go
@@ -9,6 +9,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"sort"
 	"strings"
 	"time"
 
@@ -24,6 +25,7 @@ type HTTPServer struct {
 	Debug              bool
 	Events             *Events
 	IgnoredStaticPaths []string
+	Domains            []string
 
 	mux       *pat.PatternServeMux
 	transport *httpu.Transport
@@ -106,6 +108,14 @@ func (h *HTTPServer) findApp(name string) (*App, error) {
 	return app, nil
 }
 
+type ByDecreasingTLDComplexity []string
+
+func (a ByDecreasingTLDComplexity) Len() int      { return len(a) }
+func (a ByDecreasingTLDComplexity) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
+func (a ByDecreasingTLDComplexity) Less(i, j int) bool {
+	return strings.Count(a[i], ".") > strings.Count(a[j], ".")
+}
+
 func (h *HTTPServer) removeTLD(host string) string {
 	colon := strings.LastIndexByte(host, ':')
 	if colon != -1 {
@@ -125,6 +135,15 @@ func (h *HTTPServer) removeTLD(host string) string {
 		return name
 	}
 
+	if !sort.IsSorted(ByDecreasingTLDComplexity(h.Domains)) {
+		sort.Sort(ByDecreasingTLDComplexity(h.Domains))
+	}
+	for _, tld := range h.Domains {
+		if strings.HasSuffix(host, "."+tld) {
+			return strings.TrimSuffix(host, "."+tld)
+		}
+	}
+
 	dot := strings.LastIndexByte(host, '.')
 
 	if dot == -1 {

diff --git a/dev/http_test.go b/dev/http_test.go
@@ -49,3 +49,17 @@ func TestHttp_removeTLD_nipIoDots(t *testing.T) {
 
 	assert.Equal(t, "effective-invention", str)
 }
+
+func TestHttp_removeTLD_multipartTLD(t *testing.T) {
+	testHttp.Domains = []string{"co.test"}
+	str := testHttp.removeTLD("confusing-riddle.co.test")
+
+	assert.Equal(t, "confusing-riddle", str)
+}
+
+func TestHttp_removeTLD_multipartTLDSimilarToShorterOne(t *testing.T) {
+	testHttp.Domains = []string{"test", "co.test"}
+	str := testHttp.removeTLD("confusing-riddle.co.test")
+
+	assert.Equal(t, "confusing-riddle", str)
+}