From fdbf1ad787e714952a6e58bc1b7eaa3c489d0f71 Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Sun, 13 Apr 2025 06:13:17 +0800
Subject: [PATCH] refactor: simplify api code

---
 internal/api/v1/agents.go                | 14 ++------------
 internal/api/v1/health.go                | 12 +-----------
 internal/api/v1/stats.go                 | 24 ++++++------------------
 internal/net/gphttp/gpwebsocket/utils.go | 17 +++++++++++++++++
 4 files changed, 26 insertions(+), 41 deletions(-)

diff --git a/internal/api/v1/agents.go b/internal/api/v1/agents.go
index d9d2a87..81636d5 100644
--- a/internal/api/v1/agents.go
+++ b/internal/api/v1/agents.go
@@ -4,21 +4,11 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/coder/websocket"
-	"github.com/coder/websocket/wsjson"
+	"github.com/yusing/go-proxy/agent/pkg/agent"
 	config "github.com/yusing/go-proxy/internal/config/types"
-	"github.com/yusing/go-proxy/internal/net/gphttp"
 	"github.com/yusing/go-proxy/internal/net/gphttp/gpwebsocket"
-	"github.com/yusing/go-proxy/internal/net/gphttp/httpheaders"
 )
 
 func ListAgents(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
-	if httpheaders.IsWebsocket(r.Header) {
-		gpwebsocket.Periodic(w, r, 10*time.Second, func(conn *websocket.Conn) error {
-			wsjson.Write(r.Context(), conn, cfg.ListAgents())
-			return nil
-		})
-	} else {
-		gphttp.RespondJSON(w, r, cfg.ListAgents())
-	}
+	gpwebsocket.DynamicJSONHandler(w, r, agent.Agents.Slice, 10*time.Second)
 }
diff --git a/internal/api/v1/health.go b/internal/api/v1/health.go
index 3379022..fafaf93 100644
--- a/internal/api/v1/health.go
+++ b/internal/api/v1/health.go
@@ -4,20 +4,10 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/coder/websocket"
-	"github.com/coder/websocket/wsjson"
-	"github.com/yusing/go-proxy/internal/net/gphttp"
 	"github.com/yusing/go-proxy/internal/net/gphttp/gpwebsocket"
-	"github.com/yusing/go-proxy/internal/net/gphttp/httpheaders"
 	"github.com/yusing/go-proxy/internal/route/routes/routequery"
 )
 
 func Health(w http.ResponseWriter, r *http.Request) {
-	if httpheaders.IsWebsocket(r.Header) {
-		gpwebsocket.Periodic(w, r, 1*time.Second, func(conn *websocket.Conn) error {
-			return wsjson.Write(r.Context(), conn, routequery.HealthMap())
-		})
-	} else {
-		gphttp.RespondJSON(w, r, routequery.HealthMap())
-	}
+	gpwebsocket.DynamicJSONHandler(w, r, routequery.HealthMap, 1*time.Second)
 }
diff --git a/internal/api/v1/stats.go b/internal/api/v1/stats.go
index 1fbac51..80c66dd 100644
--- a/internal/api/v1/stats.go
+++ b/internal/api/v1/stats.go
@@ -4,30 +4,18 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/coder/websocket"
-	"github.com/coder/websocket/wsjson"
 	config "github.com/yusing/go-proxy/internal/config/types"
-	"github.com/yusing/go-proxy/internal/net/gphttp"
 	"github.com/yusing/go-proxy/internal/net/gphttp/gpwebsocket"
-	"github.com/yusing/go-proxy/internal/net/gphttp/httpheaders"
 	"github.com/yusing/go-proxy/internal/utils/strutils"
 )
 
 func Stats(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
-	if httpheaders.IsWebsocket(r.Header) {
-		gpwebsocket.Periodic(w, r, 1*time.Second, func(conn *websocket.Conn) error {
-			return wsjson.Write(r.Context(), conn, getStats(cfg))
-		})
-	} else {
-		gphttp.RespondJSON(w, r, getStats(cfg))
-	}
+	gpwebsocket.DynamicJSONHandler(w, r, func() map[string]any {
+		return map[string]any{
+			"proxies": cfg.Statistics(),
+			"uptime":  strutils.FormatDuration(time.Since(startTime)),
+		}
+	}, 1*time.Second)
 }
 
 var startTime = time.Now()
-
-func getStats(cfg config.ConfigInstance) map[string]any {
-	return map[string]any{
-		"proxies": cfg.Statistics(),
-		"uptime":  strutils.FormatDuration(time.Since(startTime)),
-	}
-}
diff --git a/internal/net/gphttp/gpwebsocket/utils.go b/internal/net/gphttp/gpwebsocket/utils.go
index cb67dd6..cc98592 100644
--- a/internal/net/gphttp/gpwebsocket/utils.go
+++ b/internal/net/gphttp/gpwebsocket/utils.go
@@ -6,11 +6,13 @@ import (
 	"time"
 
 	"github.com/coder/websocket"
+	"github.com/coder/websocket/wsjson"
 	"github.com/yusing/go-proxy/internal/common"
 	"github.com/yusing/go-proxy/internal/gperr"
 	"github.com/yusing/go-proxy/internal/logging"
 	"github.com/yusing/go-proxy/internal/net/gphttp"
 	"github.com/yusing/go-proxy/internal/net/gphttp/httpheaders"
+	"github.com/yusing/go-proxy/internal/utils"
 )
 
 func warnNoMatchDomains() {
@@ -84,3 +86,18 @@ func WriteText(r *http.Request, conn *websocket.Conn, msg string) bool {
 	}
 	return true
 }
+
+// DynamicJSONHandler serves a JSON response depending on the request type.
+//
+// If the request is a websocket, it serves the data for the given interval.
+//
+// Otherwise, it serves the data once.
+func DynamicJSONHandler[ResultType any](w http.ResponseWriter, r *http.Request, getter func() ResultType, interval time.Duration) {
+	if httpheaders.IsWebsocket(r.Header) {
+		Periodic(w, r, interval, func(conn *websocket.Conn) error {
+			return wsjson.Write(r.Context(), conn, utils.ToJSONObject(getter()))
+		})
+	} else {
+		gphttp.RespondJSON(w, r, getter())
+	}
+}