starred/godoxy-yusing
1
0
Fork 0
mirror of https://github.com/yusing/godoxy.git synced 2026-03-18 23:33:51 +01:00
Code Issues 18 Packages Projects Releases 10 Wiki Activity

refactor(agent): extract agent pool and HTTP utilities to dedicated package

Browse Source 
Moved non-agent-specific logic from agent/pkg/agent/ to internal/agentpool/:
- pool.go: Agent pool management (Get, Add, Remove, List, Iter, etc.)
- http_requests.go: HTTP utilities (health checks, forwarding, websockets, reverse proxy)
- agent.go: Agent struct with HTTP client management

This separates general-purpose pool management from agent-specific configuration,
improving code organization and making the agent package focused on agent config only.
This commit is contained in:
yusing
2026-01-08 12:02:21 +08:00
parent 751d73da7c
commit db2eda49f1
21 changed files with 380 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
agent/pkg/agent/agent_pool.go
View File

@@ -1,68 +0,0 @@
package agent
import (
"iter"
"os"
"strings"
"github.com/puzpuzpuz/xsync/v4"
)
var agentPool = xsync.NewMap[string, *AgentConfig](xsync.WithPresize(10))
func init() {
if strings.HasSuffix(os.Args[0], ".test") {
agentPool.Store("test-agent", &AgentConfig{
Addr: "test-agent",
})
}
}
func GetAgent(agentAddrOrDockerHost string) (*AgentConfig, bool) {
if !IsDockerHostAgent(agentAddrOrDockerHost) {
return getAgentByAddr(agentAddrOrDockerHost)
}
return getAgentByAddr(GetAgentAddrFromDockerHost(agentAddrOrDockerHost))
}
func GetAgentByName(name string) (*AgentConfig, bool) {
for _, agent := range agentPool.Range {
if agent.Name == name {
return agent, true
}
}
return nil, false
}
func AddAgent(agent *AgentConfig) {
agentPool.Store(agent.Addr, agent)
}
func RemoveAgent(agent *AgentConfig) {
agentPool.Delete(agent.Addr)
}
func RemoveAllAgents() {
agentPool.Clear()
}
func ListAgents() []*AgentConfig {
agents := make([]*AgentConfig, 0, agentPool.Size())
for _, agent := range agentPool.Range {
agents = append(agents, agent)
}
return agents
}
func IterAgents() iter.Seq2[string, *AgentConfig] {
return agentPool.Range
}
func NumAgents() int {
return agentPool.Size()
}
func getAgentByAddr(addr string) (agent *AgentConfig, ok bool) {
agent, ok = agentPool.Load(addr)
return agent, ok
}

143
agent/pkg/agent/config.go
View File

@@ -4,9 +4,11 @@ import (
"context"
"crypto/tls"
"crypto/x509"
"encoding/json"
"encoding/pem"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/url"
@@ -16,11 +18,11 @@ import (
"github.com/rs/zerolog"
"github.com/rs/zerolog/log"
"github.com/valyala/fasthttp"
"github.com/yusing/godoxy/agent/pkg/agent/common"
agentstream "github.com/yusing/godoxy/agent/pkg/agent/stream"
"github.com/yusing/godoxy/agent/pkg/certs"
gperr "github.com/yusing/goutils/errs"
httputils "github.com/yusing/goutils/http"
"github.com/yusing/goutils/version"
)
@@ -31,9 +33,7 @@ type AgentConfig struct {
IsTCPStreamSupported bool `json:"supports_tcp_stream"`
IsUDPStreamSupported bool `json:"supports_udp_stream"`
httpClient *http.Client
fasthttpClientHealthCheck *fasthttp.Client
tlsConfig tls.Config
tlsConfig tls.Config
// for stream
caCert *x509.Certificate
@@ -106,7 +106,8 @@ func (cfg *AgentConfig) Parse(addr string) error {
var serverVersion = version.Get()
func (cfg *AgentConfig) StartWithCerts(ctx context.Context, ca, crt, key []byte) error {
// InitWithCerts initializes the agent config with the given CA, certificate, and key.
func (cfg *AgentConfig) InitWithCerts(ctx context.Context, ca, crt, key []byte) error {
clientCert, err := tls.X509KeyPair(crt, key)
if err != nil {
return err
@@ -134,12 +135,6 @@ func (cfg *AgentConfig) StartWithCerts(ctx context.Context, ca, crt, key []byte)
ServerName: common.CertsDNSName,
}
// create transport and http client
cfg.httpClient = cfg.NewHTTPClient()
applyNormalTransportConfig(cfg.httpClient)
cfg.fasthttpClientHealthCheck = cfg.NewFastHTTPHealthCheckClient()
ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
defer cancel()
@@ -233,6 +228,25 @@ func (cfg *AgentConfig) StartWithCerts(ctx context.Context, ca, crt, key []byte)
return nil
}
func (cfg *AgentConfig) Init(ctx context.Context) error {
filepath, ok := certs.AgentCertsFilepath(cfg.Addr)
if !ok {
return fmt.Errorf("invalid agent host: %s", cfg.Addr)
}
certData, err := os.ReadFile(filepath)
if err != nil {
return fmt.Errorf("failed to read agent certs: %w", err)
}
ca, crt, key, err := certs.ExtractCert(certData)
if err != nil {
return fmt.Errorf("failed to extract agent certs: %w", err)
}
return cfg.InitWithCerts(ctx, ca, crt, key)
}
// NewTCPClient creates a new TCP client for the agent.
//
// It returns an error if
@@ -265,50 +279,6 @@ func (cfg *AgentConfig) NewUDPClient(targetAddress string) (net.Conn, error) {
return agentstream.NewUDPClient(cfg.Addr, targetAddress, cfg.caCert, cfg.clientCert)
}
func (cfg *AgentConfig) Start(ctx context.Context) error {
filepath, ok := certs.AgentCertsFilepath(cfg.Addr)
if !ok {
return fmt.Errorf("invalid agent host: %s", cfg.Addr)
}
certData, err := os.ReadFile(filepath)
if err != nil {
return fmt.Errorf("failed to read agent certs: %w", err)
}
ca, crt, key, err := certs.ExtractCert(certData)
if err != nil {
return fmt.Errorf("failed to extract agent certs: %w", err)
}
return cfg.StartWithCerts(ctx, ca, crt, key)
}
func (cfg *AgentConfig) NewHTTPClient() *http.Client {
return &http.Client{
Transport: cfg.Transport(),
}
}
func (cfg *AgentConfig) NewFastHTTPHealthCheckClient() *fasthttp.Client {
return &fasthttp.Client{
Dial: func(addr string) (net.Conn, error) {
if addr != AgentHost+":443" {
return nil, &net.AddrError{Err: "invalid address", Addr: addr}
}
return net.Dial("tcp", cfg.Addr)
},
TLSConfig: &cfg.tlsConfig,
ReadTimeout: 5 * time.Second,
WriteTimeout: 3 * time.Second,
DisableHeaderNamesNormalizing: true,
DisablePathNormalizing: true,
NoDefaultUserAgentHeader: true,
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
}
func (cfg *AgentConfig) Transport() *http.Transport {
return &http.Transport{
DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
@@ -324,6 +294,10 @@ func (cfg *AgentConfig) Transport() *http.Transport {
}
}
func (cfg *AgentConfig) TLSConfig() *tls.Config {
return &cfg.tlsConfig
}
var dialer = &net.Dialer{Timeout: 5 * time.Second}
func (cfg *AgentConfig) DialContext(ctx context.Context) (net.Conn, error) {
@@ -334,10 +308,57 @@ func (cfg *AgentConfig) String() string {
return cfg.Name + "@" + cfg.Addr
}
func applyNormalTransportConfig(client *http.Client) {
transport := client.Transport.(*http.Transport)
transport.MaxIdleConns = 100
transport.MaxIdleConnsPerHost = 100
transport.ReadBufferSize = 16384
transport.WriteBufferSize = 16384
func (cfg *AgentConfig) do(ctx context.Context, method, endpoint string, body io.Reader) (*http.Response, error) {
req, err := http.NewRequestWithContext(ctx, method, APIBaseURL+endpoint, body)
if err != nil {
return nil, err
}
client := http.Client{
Transport: cfg.Transport(),
}
return client.Do(req)
}
func (cfg *AgentConfig) fetchString(ctx context.Context, endpoint string) (string, int, error) {
resp, err := cfg.do(ctx, "GET", endpoint, nil)
if err != nil {
return "", 0, err
}
defer resp.Body.Close()
data, release, err := httputils.ReadAllBody(resp)
if err != nil {
return "", 0, err
}
ret := string(data)
release(data)
return ret, resp.StatusCode, nil
}
// fetchJSON fetches a JSON response from the agent and unmarshals it into the provided struct
//
// It will return the status code of the response, and error if any.
// If the status code is not http.StatusOK, out will be unchanged but error will still be nil.
func (cfg *AgentConfig) fetchJSON(ctx context.Context, endpoint string, out any) (int, error) {
resp, err := cfg.do(ctx, "GET", endpoint, nil)
if err != nil {
return 0, err
}
defer resp.Body.Close()
data, release, err := httputils.ReadAllBody(resp)
if err != nil {
return 0, err
}
defer release(data)
if resp.StatusCode != http.StatusOK {
return resp.StatusCode, nil
}
err = json.Unmarshal(data, out)
if err != nil {
return 0, err
}
return resp.StatusCode, nil
}

140
agent/pkg/agent/http_requests.go
View File

@@ -1,140 +0,0 @@
package agent
import (
"context"
"fmt"
"io"
"net/http"
"time"
"github.com/bytedance/sonic"
"github.com/gorilla/websocket"
"github.com/valyala/fasthttp"
httputils "github.com/yusing/goutils/http"
"github.com/yusing/goutils/http/reverseproxy"
)
func (cfg *AgentConfig) Do(ctx context.Context, method, endpoint string, body io.Reader) (*http.Response, error) {
req, err := http.NewRequestWithContext(ctx, method, APIBaseURL+endpoint, body)
if err != nil {
return nil, err
}
return cfg.httpClient.Do(req)
}
func (cfg *AgentConfig) Forward(req *http.Request, endpoint string) (*http.Response, error) {
req.URL.Host = AgentHost
req.URL.Scheme = "https"
req.URL.Path = APIEndpointBase + endpoint
req.RequestURI = ""
resp, err := cfg.httpClient.Do(req)
if err != nil {
return nil, err
}
return resp, nil
}
type HealthCheckResponse struct {
Healthy bool `json:"healthy"`
Detail string `json:"detail"`
Latency time.Duration `json:"latency"`
}
func (cfg *AgentConfig) DoHealthCheck(timeout time.Duration, query string) (ret HealthCheckResponse, err error) {
req := fasthttp.AcquireRequest()
defer fasthttp.ReleaseRequest(req)
resp := fasthttp.AcquireResponse()
defer fasthttp.ReleaseResponse(resp)
req.SetRequestURI(APIBaseURL + EndpointHealth + "?" + query)
req.Header.SetMethod(fasthttp.MethodGet)
req.Header.Set("Accept-Encoding", "identity")
req.SetConnectionClose()
start := time.Now()
err = cfg.fasthttpClientHealthCheck.DoTimeout(req, resp, timeout)
ret.Latency = time.Since(start)
if err != nil {
return ret, err
}
if status := resp.StatusCode(); status != http.StatusOK {
ret.Detail = fmt.Sprintf("HTTP %d %s", status, resp.Body())
return ret, nil
} else {
err = sonic.Unmarshal(resp.Body(), &ret)
if err != nil {
return ret, err
}
}
return ret, nil
}
func (cfg *AgentConfig) fetchString(ctx context.Context, endpoint string) (string, int, error) {
resp, err := cfg.Do(ctx, "GET", endpoint, nil)
if err != nil {
return "", 0, err
}
defer resp.Body.Close()
data, release, err := httputils.ReadAllBody(resp)
if err != nil {
return "", 0, err
}
ret := string(data)
release(data)
return ret, resp.StatusCode, nil
}
// fetchJSON fetches a JSON response from the agent and unmarshals it into the provided struct
//
// It will return the status code of the response, and error if any.
// If the status code is not http.StatusOK, out will be unchanged but error will still be nil.
func (cfg *AgentConfig) fetchJSON(ctx context.Context, endpoint string, out any) (int, error) {
resp, err := cfg.Do(ctx, "GET", endpoint, nil)
if err != nil {
return 0, err
}
defer resp.Body.Close()
data, release, err := httputils.ReadAllBody(resp)
if err != nil {
return 0, err
}
defer release(data)
if resp.StatusCode != http.StatusOK {
return resp.StatusCode, nil
}
err = sonic.Unmarshal(data, out)
if err != nil {
return 0, err
}
return resp.StatusCode, nil
}
func (cfg *AgentConfig) Websocket(ctx context.Context, endpoint string) (*websocket.Conn, *http.Response, error) {
transport := cfg.Transport()
dialer := websocket.Dialer{
NetDialContext: transport.DialContext,
NetDialTLSContext: transport.DialTLSContext,
}
return dialer.DialContext(ctx, APIBaseURL+endpoint, http.Header{
"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) {
rp := reverseproxy.NewReverseProxy("agent", AgentURL, cfg.Transport())
req.URL.Host = AgentHost
req.URL.Scheme = "https"
req.URL.Path = endpoint
req.RequestURI = ""
rp.ServeHTTP(w, req)
}