From 07df97e8b8359f8ae818cd7f8933f174b7f5fceb Mon Sep 17 00:00:00 2001
From: Tai Groot <tai@taigrr.com>
Date: Mon, 4 May 2026 06:02:33 +0000
Subject: [PATCH] fix(browser): honor forwarded websocket scheme

---
 browser/browser.go      | 23 +++++++++++++++++------
 browser/browser_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/browser/browser.go b/browser/browser.go
index c6a02c4..d0e3040 100644
--- a/browser/browser.go
+++ b/browser/browser.go
@@ -11,14 +11,25 @@ import (
 var webpage string
 
 func LogSocketViewHandler(w http.ResponseWriter, r *http.Request) {
-	wsResource := r.Host + r.URL.Path
-	if r.TLS != nil {
-		wsResource = "wss://" + wsResource
-	} else {
-		wsResource = "ws://" + wsResource
-	}
+	wsResource := websocketScheme(r) + r.Host + r.URL.Path
 	wsResource = strings.TrimSuffix(wsResource, "/") + "/ws"
 	homeTemplate.Execute(w, wsResource)
 }
 
+func websocketScheme(r *http.Request) string {
+	if forwardedProto := r.Header.Get("X-Forwarded-Proto"); forwardedProto != "" {
+		protocol := strings.ToLower(strings.TrimSpace(strings.Split(forwardedProto, ",")[0]))
+		if protocol == "https" {
+			return "wss://"
+		}
+		if protocol == "http" {
+			return "ws://"
+		}
+	}
+	if r.TLS != nil {
+		return "wss://"
+	}
+	return "ws://"
+}
+
 var homeTemplate = template.Must(template.New("").Parse(webpage))
diff --git a/browser/browser_test.go b/browser/browser_test.go
index 40b1122..01f9063 100644
--- a/browser/browser_test.go
+++ b/browser/browser_test.go
@@ -1,6 +1,7 @@
 package browser
 
 import (
+	"crypto/tls"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -51,3 +52,28 @@ func TestLogSocketViewHandler_TrailingSlashTrimmed(t *testing.T) {
 		t.Error("expected escaped ws://example.com/ws in body")
 	}
 }
+
+func TestLogSocketViewHandler_ForwardedHTTPSUsesWSS(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "http://example.com/logs/", nil)
+	req.Header.Set("X-Forwarded-Proto", "https")
+	w := httptest.NewRecorder()
+	LogSocketViewHandler(w, req)
+
+	body := w.Body.String()
+	if !strings.Contains(body, `wss:\/\/example.com\/logs\/ws`) {
+		t.Error("expected escaped wss://example.com/logs/ws in body")
+	}
+}
+
+func TestLogSocketViewHandler_ForwardedHTTPOverridesTLS(t *testing.T) {
+	req := httptest.NewRequest(http.MethodGet, "https://example.com/logs/", nil)
+	req.TLS = &tls.ConnectionState{}
+	req.Header.Set("X-Forwarded-Proto", "http")
+	w := httptest.NewRecorder()
+	LogSocketViewHandler(w, req)
+
+	body := w.Body.String()
+	if !strings.Contains(body, `ws:\/\/example.com\/logs\/ws`) {
+		t.Error("expected escaped ws://example.com/logs/ws in body")
+	}
+}