refactor header utils to httpheader package, cleanup api endpoints

This commit is contained in:
yusing
2025-02-13 07:32:59 +08:00
parent 5c9083a5df
commit 02d1c9ce98
19 changed files with 237 additions and 177 deletions

View File

@@ -8,11 +8,16 @@ import (
"github.com/coder/websocket/wsjson"
U "github.com/yusing/go-proxy/internal/api/v1/utils"
config "github.com/yusing/go-proxy/internal/config/types"
"github.com/yusing/go-proxy/internal/net/http/httpheaders"
)
func AgentsWS(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
U.PeriodicWS(cfg.Value().MatchDomains, w, r, 10*time.Second, func(conn *websocket.Conn) error {
wsjson.Write(r.Context(), conn, cfg.ListAgents())
return nil
})
if httpheaders.IsWebsocket(r.Header) {
U.PeriodicWS(w, r, 10*time.Second, func(conn *websocket.Conn) error {
wsjson.Write(r.Context(), conn, cfg.ListAgents())
return nil
})
} else {
U.RespondJSON(w, r, cfg.ListAgents())
}
}

View File

@@ -7,12 +7,16 @@ import (
"github.com/coder/websocket"
"github.com/coder/websocket/wsjson"
U "github.com/yusing/go-proxy/internal/api/v1/utils"
config "github.com/yusing/go-proxy/internal/config/types"
"github.com/yusing/go-proxy/internal/net/http/httpheaders"
"github.com/yusing/go-proxy/internal/route/routes/routequery"
)
func HealthWS(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
U.PeriodicWS(cfg.Value().MatchDomains, w, r, 1*time.Second, func(conn *websocket.Conn) error {
return wsjson.Write(r.Context(), conn, routequery.HealthMap())
})
func Health(w http.ResponseWriter, r *http.Request) {
if httpheaders.IsWebsocket(r.Header) {
U.PeriodicWS(w, r, 1*time.Second, func(conn *websocket.Conn) error {
return wsjson.Write(r.Context(), conn, routequery.HealthMap())
})
} else {
U.RespondJSON(w, r, routequery.HealthMap())
}
}

View File

@@ -8,17 +8,18 @@ import (
"github.com/coder/websocket/wsjson"
U "github.com/yusing/go-proxy/internal/api/v1/utils"
config "github.com/yusing/go-proxy/internal/config/types"
"github.com/yusing/go-proxy/internal/net/http/httpheaders"
"github.com/yusing/go-proxy/internal/utils/strutils"
)
func Stats(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
U.RespondJSON(w, r, getStats(cfg))
}
func StatsWS(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
U.PeriodicWS(cfg.Value().MatchDomains, w, r, 1*time.Second, func(conn *websocket.Conn) error {
return wsjson.Write(r.Context(), conn, getStats(cfg))
})
if httpheaders.IsWebsocket(r.Header) {
U.PeriodicWS(w, r, 1*time.Second, func(conn *websocket.Conn) error {
return wsjson.Write(r.Context(), conn, getStats(cfg))
})
} else {
U.RespondJSON(w, r, getStats(cfg))
}
}
var startTime = time.Now()

View File

@@ -2,24 +2,22 @@ package v1
import (
"net/http"
"strings"
"github.com/coder/websocket/wsjson"
agentPkg "github.com/yusing/go-proxy/agent/pkg/agent"
U "github.com/yusing/go-proxy/internal/api/v1/utils"
config "github.com/yusing/go-proxy/internal/config/types"
E "github.com/yusing/go-proxy/internal/error"
"github.com/yusing/go-proxy/internal/metrics/systeminfo"
"github.com/yusing/go-proxy/internal/net/http/httpheaders"
)
func SystemInfo(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Request) {
isWS := strings.HasSuffix(r.URL.Path, "/ws")
agentName := r.URL.Query().Get("agent_name")
query := r.URL.Query()
agentName := query.Get("agent_name")
query.Del("agent_name")
if agentName == "" {
if isWS {
systeminfo.Poller.ServeWS(cfg.Value().MatchDomains, w, r)
} else {
systeminfo.Poller.ServeHTTP(w, r)
}
systeminfo.Poller.ServeHTTP(w, r)
return
}
@@ -28,10 +26,12 @@ func SystemInfo(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Reques
U.HandleErr(w, r, U.ErrInvalidKey("agent_name"), http.StatusNotFound)
return
}
isWS := httpheaders.IsWebsocket(r.Header)
if !isWS {
respData, status, err := agent.Forward(r, agentPkg.EndpointSystemInfo)
if err != nil {
U.HandleErr(w, r, err)
U.HandleErr(w, r, E.Wrap(err, "failed to forward request to agent"))
return
}
if status != http.StatusOK {
@@ -40,14 +40,16 @@ func SystemInfo(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Reques
}
U.WriteBody(w, respData)
} else {
clientConn, err := U.InitiateWS(cfg.Value().MatchDomains, w, r)
r = r.WithContext(r.Context())
clientConn, err := U.InitiateWS(w, r)
if err != nil {
U.HandleErr(w, r, err)
U.HandleErr(w, r, E.Wrap(err, "failed to initiate websocket"))
return
}
agentConn, _, err := agent.Websocket(r.Context(), agentPkg.EndpointSystemInfo)
defer clientConn.CloseNow()
agentConn, _, err := agent.Websocket(r.Context(), agentPkg.EndpointSystemInfo+"?"+query.Encode())
if err != nil {
U.HandleErr(w, r, err)
U.HandleErr(w, r, E.Wrap(err, "failed to connect to agent with websocket"))
return
}
//nolint:errcheck
@@ -63,7 +65,7 @@ func SystemInfo(cfg config.ConfigInstance, w http.ResponseWriter, r *http.Reques
err = wsjson.Write(r.Context(), clientConn, data)
}
if err != nil {
U.HandleErr(w, r, err)
U.HandleErr(w, r, E.Wrap(err, "failed to write data to client"))
return
}
}

View File

@@ -8,6 +8,7 @@ import (
"github.com/coder/websocket"
"github.com/yusing/go-proxy/internal/common"
"github.com/yusing/go-proxy/internal/logging"
"github.com/yusing/go-proxy/internal/net/http/httpheaders"
)
func warnNoMatchDomains() {
@@ -16,11 +17,12 @@ func warnNoMatchDomains() {
var warnNoMatchDomainOnce sync.Once
func InitiateWS(allowedDomains []string, w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
func InitiateWS(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
var originPats []string
localAddresses := []string{"127.0.0.1", "10.0.*.*", "172.16.*.*", "192.168.*.*"}
allowedDomains := httpheaders.WebsocketAllowedDomains(r.Header)
if len(allowedDomains) == 0 || common.IsDebug {
warnNoMatchDomainOnce.Do(warnNoMatchDomains)
originPats = []string{"*"}
@@ -40,8 +42,8 @@ func InitiateWS(allowedDomains []string, w http.ResponseWriter, r *http.Request)
})
}
func PeriodicWS(allowedDomains []string, w http.ResponseWriter, r *http.Request, interval time.Duration, do func(conn *websocket.Conn) error) {
conn, err := InitiateWS(allowedDomains, w, r)
func PeriodicWS(w http.ResponseWriter, r *http.Request, interval time.Duration, do func(conn *websocket.Conn) error) {
conn, err := InitiateWS(w, r)
if err != nil {
HandleErr(w, r, err)
return