diff --git a/agent/pkg/agent/http_requests.go b/agent/pkg/agent/http_requests.go index 6dc655d9..ba0d9693 100644 --- a/agent/pkg/agent/http_requests.go +++ b/agent/pkg/agent/http_requests.go @@ -6,6 +6,8 @@ import ( "net/http" "github.com/gorilla/websocket" + "github.com/yusing/go-proxy/internal/net/gphttp/reverseproxy" + nettypes "github.com/yusing/go-proxy/internal/net/types" ) func (cfg *AgentConfig) Do(ctx context.Context, method, endpoint string, body io.Reader) (*http.Response, error) { @@ -16,7 +18,7 @@ func (cfg *AgentConfig) Do(ctx context.Context, method, endpoint string, body io return cfg.httpClient.Do(req) } -func (cfg *AgentConfig) Forward(req *http.Request, endpoint string) ([]byte, int, error) { +func (cfg *AgentConfig) Forward(req *http.Request, endpoint string) (*http.Response, error) { req = req.WithContext(req.Context()) req.URL.Host = AgentHost req.URL.Scheme = "https" @@ -24,11 +26,9 @@ func (cfg *AgentConfig) Forward(req *http.Request, endpoint string) ([]byte, int req.RequestURI = "" resp, err := cfg.httpClient.Do(req) if err != nil { - return nil, 0, err + return nil, err } - defer resp.Body.Close() - data, _ := io.ReadAll(resp.Body) - return data, resp.StatusCode, nil + return resp, nil } func (cfg *AgentConfig) Fetch(ctx context.Context, endpoint string) ([]byte, int, error) { @@ -51,3 +51,22 @@ func (cfg *AgentConfig) Websocket(ctx context.Context, endpoint string) (*websoc "Host": {AgentHost}, }) } + +// ReverseProxy reverse proxies the request to the agent +// +// It will create a new request with the same context, method, and body, but with the agent host and scheme, and the endpoint +// If the request has a query, it will be added to the proxy request's URL +func (cfg *AgentConfig) ReverseProxy(w http.ResponseWriter, req *http.Request, endpoint string) error { + rp := reverseproxy.NewReverseProxy("agent", nettypes.NewURL(AgentURL), cfg.Transport()) + uri := APIEndpointBase + endpoint + if req.URL.RawQuery != "" { + uri += "?" + req.URL.RawQuery + } + r, err := http.NewRequestWithContext(req.Context(), req.Method, uri, req.Body) + if err != nil { + return err + } + r.Header = req.Header + rp.ServeHTTP(w, r) + return nil +} diff --git a/internal/api/v1/metrics/system_info.go b/internal/api/v1/metrics/system_info.go index a898faf0..5cd47c87 100644 --- a/internal/api/v1/metrics/system_info.go +++ b/internal/api/v1/metrics/system_info.go @@ -11,17 +11,16 @@ import ( "github.com/yusing/go-proxy/internal/metrics/period" "github.com/yusing/go-proxy/internal/metrics/systeminfo" "github.com/yusing/go-proxy/internal/net/gphttp/httpheaders" - "github.com/yusing/go-proxy/internal/net/gphttp/reverseproxy" - nettypes "github.com/yusing/go-proxy/internal/net/types" ) type SystemInfoRequest struct { AgentAddr string `query:"agent_addr"` + AgentName string `query:"agent_name"` Aggregate systeminfo.SystemInfoAggregateMode `query:"aggregate"` Period period.Filter `query:"period"` } // @name SystemInfoRequest -type SystemInfoAggregate period.ResponseType[systeminfo.Aggregated] // @name SystemInfoAggregate +type SystemInfoAggregate period.ResponseType[systeminfo.AggregatedJSON] // @name SystemInfoAggregate // @x-id "system_info" // @BasePath /api/v1 @@ -40,15 +39,20 @@ type SystemInfoAggregate period.ResponseType[systeminfo.Aggregated] // @name Sys func SystemInfo(c *gin.Context) { query := c.Request.URL.Query() agentAddr := query.Get("agent_addr") + agentName := query.Get("agent_name") query.Del("agent_addr") - if agentAddr == "" { + query.Del("agent_name") + if agentAddr == "" && agentName == "" { systeminfo.Poller.ServeHTTP(c) return } agent, ok := agentPkg.GetAgent(agentAddr) if !ok { - c.JSON(http.StatusNotFound, apitypes.Error("agent_addr not found")) + agent, ok = agentPkg.GetAgentByName(agentName) + } + if !ok { + c.JSON(http.StatusNotFound, apitypes.Error("agent_addr or agent_name not found")) return } @@ -59,17 +63,16 @@ func SystemInfo(c *gin.Context) { c.Error(apitypes.InternalServerError(err, "failed to forward request to agent")) return } + defer resp.Body.Close() + maps.Copy(c.Writer.Header(), resp.Header) c.Status(resp.StatusCode) io.Copy(c.Writer, resp.Body) } else { - rp := reverseproxy.NewReverseProxy("agent", nettypes.NewURL(agentPkg.AgentURL), agent.Transport()) - r, err := http.NewRequestWithContext(c.Request.Context(), c.Request.Method, agentPkg.EndpointSystemInfo+"?"+query.Encode(), c.Request.Body) + err := agent.ReverseProxy(c.Writer, c.Request, agentPkg.EndpointSystemInfo+"?"+query.Encode()) if err != nil { - c.Error(apitypes.InternalServerError(err, "failed to create request")) + c.Error(apitypes.InternalServerError(err, "failed to reverse proxy")) return } - r.Header = c.Request.Header - rp.ServeHTTP(c.Writer, r) } }