From e9ac3cd1a9c5510796b09e00cc6a53a377f1c9f8 Mon Sep 17 00:00:00 2001
From: yusing <yusing.wys@gmail.com>
Date: Fri, 7 Nov 2025 15:48:38 +0800
Subject: [PATCH] refactor: fix incorrect logic introduced in previous commits
 and improve error handling

---
 internal/idlewatcher/events.go                |  1 +
 internal/idlewatcher/handle_http.go           |  3 +-
 internal/idlewatcher/html/loading.js          | 36 ++++++++++++++++++-
 internal/idlewatcher/html/loading_page.html   |  2 +-
 internal/idlewatcher/html/style.css           |  8 ++++-
 internal/idlewatcher/loading_page.go          |  6 ++++
 internal/idlewatcher/watcher.go               | 18 +++++++---
 internal/net/gphttp/loadbalancer/ip_hash.go   | 29 +++++++--------
 .../net/gphttp/loadbalancer/least_conn.go     | 13 +++----
 .../net/gphttp/loadbalancer/loadbalancer.go   |  6 ----
 .../net/gphttp/loadbalancer/round_robin.go    |  5 +--
 internal/net/gphttp/loadbalancer/sticky.go    |  3 +-
 internal/net/gphttp/middleware/middleware.go  |  7 ++++
 internal/watcher/health/monitor/raw.go        | 12 ++++---
 14 files changed, 99 insertions(+), 50 deletions(-)

diff --git a/internal/idlewatcher/events.go b/internal/idlewatcher/events.go
index bb408671..39d6e23b 100644
--- a/internal/idlewatcher/events.go
+++ b/internal/idlewatcher/events.go
@@ -55,6 +55,7 @@ func (w *Watcher) clearEventHistory() {
 }
 
 func (w *Watcher) sendEvent(eventType WakeEventType, message string, err error) {
+	// NOTE: events will be cleared on stop/pause
 	event := w.newWakeEvent(eventType, message, err)
 
 	w.l.Debug().Str("event", string(eventType)).Str("message", message).Err(err).Msg("sending event")
diff --git a/internal/idlewatcher/handle_http.go b/internal/idlewatcher/handle_http.go
index efdbe3e8..84f4ba55 100644
--- a/internal/idlewatcher/handle_http.go
+++ b/internal/idlewatcher/handle_http.go
@@ -15,8 +15,6 @@ import (
 	_ "unsafe"
 )
 
-// FIXME: html and js ccannot be separte
-
 type ForceCacheControl struct {
 	expires string
 	http.ResponseWriter
@@ -179,6 +177,7 @@ func (w *Watcher) wakeFromHTTP(rw http.ResponseWriter, r *http.Request) (shouldN
 
 	if !acceptHTML {
 		serveStaticContent(rw, http.StatusOK, "text/plain", []byte("Container woken"))
+		return false
 	}
 
 	// Send a loading response to the client
diff --git a/internal/idlewatcher/html/loading.js b/internal/idlewatcher/html/loading.js
index aaf903aa..8eb485d0 100644
--- a/internal/idlewatcher/html/loading.js
+++ b/internal/idlewatcher/html/loading.js
@@ -4,6 +4,11 @@ window.onload = async function () {
   const consoleEl = document.getElementById("console");
   const loadingDotsEl = document.getElementById("loading-dots");
 
+  if (!consoleEl || !loadingDotsEl) {
+    console.error("Required DOM elements not found");
+    return;
+  }
+
   function formatTimestamp(timestamp) {
     const date = new Date(timestamp);
     return date.toLocaleTimeString("en-US", {
@@ -34,11 +39,40 @@ window.onload = async function () {
     consoleEl.scrollTop = consoleEl.scrollHeight;
   }
 
+  if (typeof wakeEventsPath === "undefined" || !wakeEventsPath) {
+    addConsoleLine(
+      "error",
+      "Configuration error: wakeEventsPath not defined",
+      new Date().toISOString()
+    );
+    loadingDotsEl.style.display = "none";
+    return;
+  }
+
+  if (typeof EventSource === "undefined") {
+    addConsoleLine(
+      "error",
+      "Browser does not support Server-Sent Events",
+      new Date().toISOString()
+    );
+    loadingDotsEl.style.display = "none";
+    return;
+  }
+
   // Connect to SSE endpoint
   const eventSource = new EventSource(wakeEventsPath);
 
   eventSource.onmessage = function (event) {
-    const data = JSON.parse(event.data);
+    try {
+      const data = JSON.parse(event.data);
+    } catch (error) {
+      addConsoleLine(
+        "error",
+        "Invalid event data: " + event.data,
+        new Date().toISOString()
+      );
+      return;
+    }
 
     if (data.type === "ready") {
       ready = true;
diff --git a/internal/idlewatcher/html/loading_page.html b/internal/idlewatcher/html/loading_page.html
index f5f82004..1cb990ba 100644
--- a/internal/idlewatcher/html/loading_page.html
+++ b/internal/idlewatcher/html/loading_page.html
@@ -26,7 +26,7 @@
         <div class="dot"></div>
         <div class="dot"></div>
       </div>
-      <div id="console" class="console" />
+      <div id="console" class="console"></div>
     </div>
   </body>
 </html>
diff --git a/internal/idlewatcher/html/style.css b/internal/idlewatcher/html/style.css
index f261afd8..d7d79728 100644
--- a/internal/idlewatcher/html/style.css
+++ b/internal/idlewatcher/html/style.css
@@ -42,6 +42,13 @@ body {
   min-height: 400px;
 }
 
+@media (max-width: 768px) {
+  .container {
+    min-width: auto;
+    max-width: 90%;
+  }
+}
+
 /* Spinner Styles */
 .loading-dots {
   display: flex;
@@ -81,7 +88,6 @@ body {
   color: #f8f9fa;
   max-width: 100%;
   letter-spacing: 0.3px;
-  white-space: nowrap;
 }
 
 /* Console Styles */
diff --git a/internal/idlewatcher/loading_page.go b/internal/idlewatcher/loading_page.go
index 908e7684..8fdd71a7 100644
--- a/internal/idlewatcher/loading_page.go
+++ b/internal/idlewatcher/loading_page.go
@@ -38,6 +38,12 @@ func (w *Watcher) writeLoadingPage(rw http.ResponseWriter) error {
 	data.LoadingPageCSSPath = idlewatcher.LoadingPageCSSPath
 	data.LoadingPageJSPath = idlewatcher.LoadingPageJSPath
 	data.WakeEventsPath = idlewatcher.WakeEventsPath
+
+	rw.Header().Set("Content-Type", "text/html; charset=utf-8")
+	rw.Header().Set("Cache-Control", "no-cache")
+	rw.Header().Add("Cache-Control", "no-store")
+	rw.Header().Add("Cache-Control", "must-revalidate")
+	rw.Header().Add("Connection", "close")
 	err := loadingPageTmpl.Execute(rw, data)
 	return err
 }
diff --git a/internal/idlewatcher/watcher.go b/internal/idlewatcher/watcher.go
index 33fdfb8c..a589cd2e 100644
--- a/internal/idlewatcher/watcher.go
+++ b/internal/idlewatcher/watcher.go
@@ -262,10 +262,14 @@ func NewWatcher(parent task.Parent, r types.Route, cfg *types.IdlewatcherConfig)
 		p, err = provider.NewProxmoxProvider(cfg.Proxmox.Node, cfg.Proxmox.VMID)
 		kind = "proxmox"
 	}
+	targetURL := r.TargetURL()
+	if targetURL != nil {
+		return nil, errors.New("target URL is not set")
+	}
 	w.l = log.With().
 		Str("kind", kind).
 		Str("container", cfg.ContainerName()).
-		Str("url", r.TargetURL().String()).
+		Str("url", targetURL.String()).
 		Logger()
 
 	if cfg.IdleTimeout != neverTick {
@@ -312,7 +316,7 @@ func NewWatcher(parent task.Parent, r types.Route, cfg *types.IdlewatcherConfig)
 		watcherMap[key] = w
 
 		go func() {
-			cause := w.watchUntilDestroy(p)
+			cause := w.watchUntilDestroy()
 
 			watcherMapMu.Lock()
 			delete(watcherMap, key)
@@ -414,8 +418,8 @@ func (w *Watcher) wakeDependencies(ctx context.Context) error {
 
 	errs := errgroup.Group{}
 	for _, dep := range w.dependsOn {
-		if w.wakeInProgress() {
-			w.l.Debug().Str("dependency", dep.cfg.ContainerName()).Msg("dependency already starting, ignoring duplicate start event")
+		if dep.wakeInProgress() {
+			w.l.Debug().Str("dep", dep.cfg.ContainerName()).Msg("dependency already starting, ignoring duplicate start event")
 			continue
 		}
 		errs.Go(func() error {
@@ -558,7 +562,11 @@ func (w *Watcher) expires() time.Time {
 //
 // it exits only if the context is canceled, the container is destroyed,
 // errors occurred on docker client, or route provider died (mainly caused by config reload).
-func (w *Watcher) watchUntilDestroy(p idlewatcher.Provider) (returnCause error) {
+func (w *Watcher) watchUntilDestroy() (returnCause error) {
+	p := w.provider.Load()
+	if p == nil {
+		return gperr.Errorf("provider not set")
+	}
 	defer p.Close()
 	eventCh, errCh := p.Watch(w.Task().Context())
 
diff --git a/internal/net/gphttp/loadbalancer/ip_hash.go b/internal/net/gphttp/loadbalancer/ip_hash.go
index fdea282d..710757c6 100644
--- a/internal/net/gphttp/loadbalancer/ip_hash.go
+++ b/internal/net/gphttp/loadbalancer/ip_hash.go
@@ -3,6 +3,7 @@ package loadbalancer
 import (
 	"net"
 	"net/http"
+	"slices"
 	"sync"
 
 	"github.com/bytedance/gopkg/util/xxhash3"
@@ -39,16 +40,6 @@ func (impl *ipHash) OnAddServer(srv types.LoadBalancerServer) {
 	impl.mu.Lock()
 	defer impl.mu.Unlock()
 
-	for i, s := range impl.pool {
-		if s == srv {
-			return
-		}
-		if s == nil {
-			impl.pool[i] = srv
-			return
-		}
-	}
-
 	impl.pool = append(impl.pool, srv)
 }
 
@@ -58,27 +49,33 @@ func (impl *ipHash) OnRemoveServer(srv types.LoadBalancerServer) {
 
 	for i, s := range impl.pool {
 		if s == srv {
-			impl.pool[i] = nil
+			impl.pool = slices.Delete(impl.pool, i, 1)
 			return
 		}
 	}
 }
 
 func (impl *ipHash) ServeHTTP(_ types.LoadBalancerServers, rw http.ResponseWriter, r *http.Request) {
+	if impl.realIP != nil {
+		// resolve real client IP
+		if proceed := impl.realIP.TryModifyRequest(rw, r); !proceed {
+			return
+		}
+	}
+
 	srv := impl.ChooseServer(impl.pool, r)
 	if srv == nil || srv.Status().Bad() {
 		http.Error(rw, "Service unavailable", http.StatusServiceUnavailable)
 		return
 	}
 
-	if impl.realIP != nil {
-		impl.realIP.ModifyRequest(srv.ServeHTTP, rw, r)
-	} else {
-		srv.ServeHTTP(rw, r)
-	}
+	srv.ServeHTTP(rw, r)
 }
 
 func (impl *ipHash) ChooseServer(_ types.LoadBalancerServers, r *http.Request) types.LoadBalancerServer {
+	impl.mu.Lock()
+	defer impl.mu.Unlock()
+
 	if len(impl.pool) == 0 {
 		return nil
 	}
diff --git a/internal/net/gphttp/loadbalancer/least_conn.go b/internal/net/gphttp/loadbalancer/least_conn.go
index 0783e177..abc8a565 100644
--- a/internal/net/gphttp/loadbalancer/least_conn.go
+++ b/internal/net/gphttp/loadbalancer/least_conn.go
@@ -46,8 +46,8 @@ func (impl *leastConn) ServeHTTP(srvs types.LoadBalancerServers, rw http.Respons
 	}
 
 	minConn.Add(1)
+	defer minConn.Add(-1)
 	srv.ServeHTTP(rw, r)
-	minConn.Add(-1)
 }
 
 func (impl *leastConn) ChooseServer(srvs types.LoadBalancerServers, r *http.Request) types.LoadBalancerServer {
@@ -55,18 +55,15 @@ func (impl *leastConn) ChooseServer(srvs types.LoadBalancerServers, r *http.Requ
 		return nil
 	}
 
-	srv := srvs[0]
-	minConn, ok := impl.nConn.Load(srv)
-	if !ok {
-		return nil
-	}
+	var srv types.LoadBalancerServer
+	var minConn *atomic.Int64
 
-	for i := 1; i < len(srvs); i++ {
+	for i := range srvs {
 		nConn, ok := impl.nConn.Load(srvs[i])
 		if !ok {
 			continue
 		}
-		if nConn.Load() < minConn.Load() {
+		if minConn == nil || nConn.Load() < minConn.Load() {
 			minConn = nConn
 			srv = srvs[i]
 		}
diff --git a/internal/net/gphttp/loadbalancer/loadbalancer.go b/internal/net/gphttp/loadbalancer/loadbalancer.go
index b8251cc4..da398e32 100644
--- a/internal/net/gphttp/loadbalancer/loadbalancer.go
+++ b/internal/net/gphttp/loadbalancer/loadbalancer.go
@@ -364,11 +364,5 @@ func isIdlewatcherRequest(r *http.Request) bool {
 		return true
 	}
 
-	// Check if this is a page refresh after idlewatcher wake up
-	// by looking for the sticky session cookie
-	if _, err := r.Cookie("godoxy_lb_sticky"); err == nil {
-		return true
-	}
-
 	return false
 }
diff --git a/internal/net/gphttp/loadbalancer/round_robin.go b/internal/net/gphttp/loadbalancer/round_robin.go
index 31c61c9e..2f228d56 100644
--- a/internal/net/gphttp/loadbalancer/round_robin.go
+++ b/internal/net/gphttp/loadbalancer/round_robin.go
@@ -21,9 +21,6 @@ func (lb *roundRobin) ChooseServer(srvs types.LoadBalancerServers, r *http.Reque
 	if len(srvs) == 0 {
 		return nil
 	}
-	index := lb.index.Add(1) % uint32(len(srvs))
-	if lb.index.Load() >= 2*uint32(len(srvs)) {
-		lb.index.Store(0)
-	}
+	index := (lb.index.Add(1) - 1) % uint32(len(srvs))
 	return srvs[index]
 }
diff --git a/internal/net/gphttp/loadbalancer/sticky.go b/internal/net/gphttp/loadbalancer/sticky.go
index ffe7f486..6e4e0bd0 100644
--- a/internal/net/gphttp/loadbalancer/sticky.go
+++ b/internal/net/gphttp/loadbalancer/sticky.go
@@ -3,6 +3,7 @@ package loadbalancer
 import (
 	"encoding/hex"
 	"net/http"
+	"strings"
 	"time"
 	"unsafe"
 
@@ -46,5 +47,5 @@ func setStickyCookie(rw http.ResponseWriter, r *http.Request, srv types.LoadBala
 }
 
 func isSecure(r *http.Request) bool {
-	return r.TLS != nil || r.Header.Get("X-Forwarded-Proto") == "https"
+	return r.TLS != nil || strings.EqualFold(r.Header.Get("X-Forwarded-Proto"), "https")
 }
diff --git a/internal/net/gphttp/middleware/middleware.go b/internal/net/gphttp/middleware/middleware.go
index 5100b083..efd7b118 100644
--- a/internal/net/gphttp/middleware/middleware.go
+++ b/internal/net/gphttp/middleware/middleware.go
@@ -170,6 +170,13 @@ func (m *Middleware) ModifyRequest(next http.HandlerFunc, w http.ResponseWriter,
 	next(w, r)
 }
 
+func (m *Middleware) TryModifyRequest(w http.ResponseWriter, r *http.Request) (proceed bool) {
+	if exec, ok := m.impl.(RequestModifier); ok {
+		return exec.before(w, r)
+	}
+	return true
+}
+
 func (m *Middleware) ModifyResponse(resp *http.Response) error {
 	if exec, ok := m.impl.(ResponseModifier); ok {
 		return exec.modifyResponse(resp)
diff --git a/internal/watcher/health/monitor/raw.go b/internal/watcher/health/monitor/raw.go
index c9263968..72275533 100644
--- a/internal/watcher/health/monitor/raw.go
+++ b/internal/watcher/health/monitor/raw.go
@@ -1,9 +1,10 @@
 package monitor
 
 import (
+	"errors"
 	"net"
 	"net/url"
-	"strings"
+	"syscall"
 	"time"
 
 	"github.com/yusing/godoxy/internal/types"
@@ -35,10 +36,11 @@ func (mon *RawHealthMonitor) CheckHealth() (types.HealthCheckResult, error) {
 	conn, err := mon.dialer.DialContext(ctx, url.Scheme, url.Host)
 	lat := time.Since(start)
 	if err != nil {
-		errMsg := err.Error()
-		if strings.Contains(errMsg, "connection refused") ||
-			strings.Contains(errMsg, "connection reset by peer") ||
-			strings.Contains(errMsg, "connection closed") {
+		if errors.Is(err, net.ErrClosed) ||
+			errors.Is(err, syscall.ECONNREFUSED) ||
+			errors.Is(err, syscall.ECONNRESET) ||
+			errors.Is(err, syscall.ECONNABORTED) ||
+			errors.Is(err, syscall.EPIPE) {
 			return types.HealthCheckResult{
 				Latency: lat,
 				Healthy: false,