From 501fc5c05907ed876d66901983026d20b87a47c3 Mon Sep 17 00:00:00 2001
From: Ben Toogood <ben@micro.mu>
Date: Tue, 7 Apr 2020 10:28:39 +0100
Subject: [PATCH] Refactor to use publicsuffix

---
 api/server/auth/auth.go      | 34 ++++++++++++++++++++++------------
 api/server/auth/auth_test.go |  5 ++++-
 go.sum                       |  1 +
 3 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/api/server/auth/auth.go b/api/server/auth/auth.go
index ab93443f..eb064421 100644
--- a/api/server/auth/auth.go
+++ b/api/server/auth/auth.go
@@ -12,6 +12,7 @@ import (
 	"github.com/micro/go-micro/v2/api/resolver/path"
 	"github.com/micro/go-micro/v2/auth"
 	"github.com/micro/go-micro/v2/logger"
+	"golang.org/x/net/publicsuffix"
 )
 
 // CombinedAuthHandler wraps a server and authenticates requests
@@ -39,7 +40,7 @@ type authHandler struct {
 
 func (h authHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	// Determine the namespace and set it in the header
-	namespace := h.namespaceFromRequest(req)
+	namespace := h.NamespaceFromRequest(req)
 	req.Header.Set(auth.NamespaceKey, namespace)
 
 	// Extract the token from the request
@@ -131,7 +132,7 @@ func (h authHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	http.Redirect(w, req, loginWithRedirect, http.StatusTemporaryRedirect)
 }
 
-func (h authHandler) namespaceFromRequest(req *http.Request) string {
+func (h authHandler) NamespaceFromRequest(req *http.Request) string {
 	// check to see what the provided namespace is, we only do
 	// domain mapping if the namespace is set to 'domain'
 	if h.namespace != "domain" {
@@ -161,18 +162,27 @@ func (h authHandler) namespaceFromRequest(req *http.Request) string {
 		return auth.DefaultNamespace
 	}
 
-	// TODO: this logic needs to be replaced with usage of publicsuffix
-	// if host is not a subdomain, deturn default namespace
-	comps := strings.Split(host, ".")
-	if len(comps) < 3 {
+	// extract the top level domain plus one (e.g. 'myapp.com')
+	domain, err := publicsuffix.EffectiveTLDPlusOne(host)
+	if err != nil {
+		logger.Debugf("Unable to extract domain from %v", host)
 		return auth.DefaultNamespace
 	}
 
-	// return the reversed subdomain as the namespace
-	nComps := comps[0 : len(comps)-2]
-	for i := len(nComps)/2 - 1; i >= 0; i-- {
-		opp := len(nComps) - 1 - i
-		nComps[i], nComps[opp] = nComps[opp], nComps[i]
+	// check to see if the domain is the host, in this
+	// case we return the default namespace
+	if domain == host {
+		return auth.DefaultNamespace
 	}
-	return strings.Join(nComps, ".")
+
+	// remove the domain from the host, leaving the subdomain
+	subdomain := strings.TrimSuffix(host, "."+domain)
+
+	// return the reversed subdomain as the namespace
+	comps := strings.Split(subdomain, ".")
+	for i := len(comps)/2 - 1; i >= 0; i-- {
+		opp := len(comps) - 1 - i
+		comps[i], comps[opp] = comps[opp], comps[i]
+	}
+	return strings.Join(comps, ".")
 }
diff --git a/api/server/auth/auth_test.go b/api/server/auth/auth_test.go
index cf454e28..04923e35 100644
--- a/api/server/auth/auth_test.go
+++ b/api/server/auth/auth_test.go
@@ -13,6 +13,7 @@ func TestNamespaceFromRequest(t *testing.T) {
 		Namespace string
 	}{
 		{Host: "micro.mu", Namespace: auth.DefaultNamespace},
+		{Host: "micro.com.au", Namespace: auth.DefaultNamespace},
 		{Host: "web.micro.mu", Namespace: auth.DefaultNamespace},
 		{Host: "api.micro.mu", Namespace: auth.DefaultNamespace},
 		{Host: "myapp.com", Namespace: auth.DefaultNamespace},
@@ -23,9 +24,11 @@ func TestNamespaceFromRequest(t *testing.T) {
 		{Host: "81.151.101.146", Namespace: auth.DefaultNamespace},
 	}
 
+	h := &authHandler{namespace: "domain"}
+
 	for _, tc := range tt {
 		t.Run(tc.Host, func(t *testing.T) {
-			ns := namespaceFromRequest(&http.Request{Host: tc.Host})
+			ns := h.NamespaceFromRequest(&http.Request{Host: tc.Host})
 			if ns != tc.Namespace {
 				t.Errorf("Expected namespace %v for host %v, actually got %v", tc.Namespace, tc.Host, ns)
 			}
diff --git a/go.sum b/go.sum
index 7d8a1ee1..b7d9c90b 100644
--- a/go.sum
+++ b/go.sum
@@ -461,6 +461,7 @@ golang.org/x/net v0.0.0-20190930134127-c5a3c61f89f3/go.mod h1:z5CRVTTTmAJ677TzLL
 golang.org/x/net v0.0.0-20191027093000-83d349e8ac1a/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200222125558-5a598a2470a0 h1:MsuvTghUPjX762sGLnGsxC3HM0B5r83wEtYcYR8/vRs=
 golang.org/x/net v0.0.0-20200222125558-5a598a2470a0/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e h1:3G+cUijn7XD+S4eJFddp53Pv7+slrESplyjG25HgL+k=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
 golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=