From 92b85c9483a877f8cb1befc0556be89ff4b14998 Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Thu, 4 Jul 2024 16:29:04 -0300
Subject: [PATCH] fix localhost trick from 8aaf5b.

---
 normalize.go      | 7 +++----
 normalize_test.go | 2 +-
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/normalize.go b/normalize.go
index 6e9333b..14a7585 100644
--- a/normalize.go
+++ b/normalize.go
@@ -15,13 +15,12 @@ func NormalizeURL(u string) string {
 	u = strings.TrimSpace(u)
 	u = strings.ToLower(u)
 
-	if strings.HasPrefix(u, "localhost") == true {
+	if strings.Split(u, ":")[0] == "localhost" {
 		u = "ws://" + u
-	}
-
-	if !strings.HasPrefix(u, "http") && !strings.HasPrefix(u, "ws") {
+	} else if !strings.HasPrefix(u, "http") && !strings.HasPrefix(u, "ws") {
 		u = "wss://" + u
 	}
+
 	p, err := url.Parse(u)
 	if err != nil {
 		return ""
diff --git a/normalize_test.go b/normalize_test.go
index 6b8d2ad..dc5a6b0 100644
--- a/normalize_test.go
+++ b/normalize_test.go
@@ -23,11 +23,11 @@ var urlTests = []urlTest{
 	{"x.com/?x=23", "wss://x.com?x=23"},
 	{"localhost:4036", "ws://localhost:4036"},
 	{"localhost:4036/relay", "ws://localhost:4036/relay"},
+	{"localhostmagnanimus.com", "wss://localhostmagnanimus.com"},
 	{NormalizeURL("localhost:4036/relay"), "ws://localhost:4036/relay"},
 }
 
 func TestNormalizeURL(t *testing.T) {
-
 	for _, test := range urlTests {
 		if output := NormalizeURL(test.url); output != test.expected {
 			t.Errorf("Output '%s' not equal to expected '%s'", output, test.expected)