From d65edcc47ada9ead0ac4bcb09d1c14a933d8ff96 Mon Sep 17 00:00:00 2001
From: Lakhan Samani <lakhan.m.samani@gmail.com>
Date: Thu, 22 Jul 2021 05:22:53 +0530
Subject: [PATCH] update getting hostname logic

---
 server/utils/cookie.go |  5 +++--
 server/utils/urls.go   | 24 +++++++++++++++++++++---
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/server/utils/cookie.go b/server/utils/cookie.go
index 281577e..bd8663c 100644
--- a/server/utils/cookie.go
+++ b/server/utils/cookie.go
@@ -14,10 +14,11 @@ func SetCookie(gc *gin.Context, token string) {
 	if !constants.IS_PROD {
 		secure = false
 	}
-
+	host := GetFrontendHost()
 	log.Println("-> is cookie secure", secure)
+	log.Println("-> host:", host)
 
-	gc.SetCookie(constants.COOKIE_NAME, token, 3600, "/", GetFrontendHost(), secure, httpOnly)
+	gc.SetCookie(constants.COOKIE_NAME, token, 3600, "/", host, secure, httpOnly)
 }
 
 func DeleteCookie(gc *gin.Context) {
diff --git a/server/utils/urls.go b/server/utils/urls.go
index 8a9cf0d..b7e8453 100644
--- a/server/utils/urls.go
+++ b/server/utils/urls.go
@@ -1,8 +1,8 @@
 package utils
 
 import (
-	"log"
 	"net/url"
+	"strings"
 
 	"github.com/yauthdev/yauth/server/constants"
 )
@@ -13,7 +13,25 @@ func GetFrontendHost() string {
 		return `localhost`
 	}
 
-	log.Println("hostname", "."+u.Hostname())
+	host := u.Hostname()
+	hostParts := strings.Split(host, ".")
+	hostPartsLen := len(hostParts)
 
-	return "." + u.Hostname()
+	if hostPartsLen == 1 {
+		return host
+	}
+
+	if hostPartsLen == 2 {
+		if hostParts[0] == "www" {
+			return hostParts[1]
+		} else {
+			return host
+		}
+	}
+
+	if hostPartsLen > 2 {
+		return strings.Join(hostParts[hostPartsLen-2:], ".")
+	}
+
+	return host
 }