fix: optimize memory usage, fix agent and code refactor (#118)

* refactor: simplify io code and make utils module independent

* fix(docker): agent and socket-proxy docker event flushing with modified reverse proxy handler

* refactor: remove unused code

* refactor: remove the use of logging module in most code

* refactor: streamline domain mismatch check in certState function

* tweak: use ecdsa p-256 for autocert

* fix(tests): update health check tests for invalid host and add case for port in host

* feat(acme): custom acme directory

* refactor: code refactor and improved context and error handling

* tweak: optimize memory usage under load

* fix(oidc): restore old user matching behavior

* docs: add ChatGPT assistant to README

---------

Co-authored-by: yusing <yusing@6uo.me>

agent/pkg/agent/config.go

@@ -15,8 +15,8 @@ import (
 	"time"
 
 	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
 	"github.com/yusing/go-proxy/agent/pkg/certs"
-	"github.com/yusing/go-proxy/internal/logging"
 	"github.com/yusing/go-proxy/pkg"
 )
 
@@ -115,7 +115,7 @@ func (cfg *AgentConfig) StartWithCerts(ctx context.Context, ca, crt, key []byte)
 
 	cfg.name = string(name)
 
-	cfg.l = logging.With().Str("agent", cfg.name).Logger()
+	cfg.l = log.With().Str("agent", cfg.name).Logger()
 
 	// check agent version
 	agentVersionBytes, _, err := cfg.Fetch(ctx, EndpointVersion)
@@ -127,10 +127,10 @@ func (cfg *AgentConfig) StartWithCerts(ctx context.Context, ca, crt, key []byte)
 	agentVersion := pkg.ParseVersion(cfg.version)
 
 	if serverVersion.IsNewerMajorThan(agentVersion) {
-		logging.Warn().Msgf("agent %s major version mismatch: server: %s, agent: %s", cfg.name, serverVersion, agentVersion)
+		log.Warn().Msgf("agent %s major version mismatch: server: %s, agent: %s", cfg.name, serverVersion, agentVersion)
 	}
 
-	logging.Info().Msgf("agent %q initialized", cfg.name)
+	log.Info().Msgf("agent %q initialized", cfg.name)
 	return nil
 }
 

agent/pkg/handler/check_health_test.go

@@ -172,9 +172,9 @@ func TestCheckHealthTCPUDP(t *testing.T) {
 		{
 			name:            "InvalidHost",
 			scheme:          "tcp",
-			host:            "invalid",
+			host:            "",
 			port:            8080,
-			expectedStatus:  http.StatusOK,
+			expectedStatus:  http.StatusBadRequest,
 			expectedHealthy: false,
 		},
 		{
@@ -188,9 +188,17 @@ func TestCheckHealthTCPUDP(t *testing.T) {
 		{
 			name:            "InvalidHost",
 			scheme:          "udp",
-			host:            "invalid",
+			host:            "",
 			port:            8080,
-			expectedStatus:  http.StatusOK,
+			expectedStatus:  http.StatusBadRequest,
+			expectedHealthy: false,
+		},
+		{
+			name:            "Port in both host and port",
+			scheme:          "tcp",
+			host:            "localhost:1234",
+			port:            1234,
+			expectedStatus:  http.StatusBadRequest,
 			expectedHealthy: false,
 		},
 	}
@@ -208,9 +216,11 @@ func TestCheckHealthTCPUDP(t *testing.T) {
 
 			require.Equal(t, recorder.Code, tt.expectedStatus)
 
-			var result health.HealthCheckResult
-			require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &result))
-			require.Equal(t, result.Healthy, tt.expectedHealthy)
+			if tt.expectedStatus == http.StatusOK {
+				var result health.HealthCheckResult
+				require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &result))
+				require.Equal(t, result.Healthy, tt.expectedHealthy)
+			}
 		})
 	}
 }

agent/pkg/handler/handler.go

@@ -1,17 +1,14 @@
 package handler
 
 import (
-	"context"
 	"fmt"
-	"net"
 	"net/http"
-	"net/http/httputil"
-	"time"
 
 	"github.com/yusing/go-proxy/agent/pkg/agent"
 	"github.com/yusing/go-proxy/agent/pkg/env"
 	"github.com/yusing/go-proxy/internal/metrics/systeminfo"
 	"github.com/yusing/go-proxy/pkg"
+	socketproxy "github.com/yusing/go-proxy/socketproxy/pkg"
 )
 
 type ServeMux struct{ *http.ServeMux }
@@ -24,26 +21,6 @@ func (mux ServeMux) HandleFunc(endpoint string, handler http.HandlerFunc) {
 	mux.ServeMux.HandleFunc(agent.APIEndpointBase+endpoint, handler)
 }
 
-var dialer = &net.Dialer{KeepAlive: 1 * time.Second}
-
-func dialDockerSocket(ctx context.Context, _, _ string) (net.Conn, error) {
-	return dialer.DialContext(ctx, "unix", env.DockerSocket)
-}
-
-func dockerSocketHandler() http.HandlerFunc {
-	rp := httputil.ReverseProxy{
-		Director: func(r *http.Request) {
-			r.URL.Scheme = "http"
-			r.URL.Host = "api.moby.localhost"
-			r.RequestURI = r.URL.String()
-		},
-		Transport: &http.Transport{
-			DialContext: dialDockerSocket,
-		},
-	}
-	return rp.ServeHTTP
-}
-
 func NewAgentHandler() http.Handler {
 	mux := ServeMux{http.NewServeMux()}
 
@@ -54,6 +31,6 @@ func NewAgentHandler() http.Handler {
 	})
 	mux.HandleEndpoint("GET", agent.EndpointHealth, CheckHealth)
 	mux.HandleEndpoint("GET", agent.EndpointSystemInfo, systeminfo.Poller.ServeHTTP)
-	mux.ServeMux.HandleFunc("/", dockerSocketHandler())
+	mux.ServeMux.HandleFunc("/", socketproxy.DockerSocketHandler(env.DockerSocket))
 	return mux
 }

agent/pkg/server/server.go

@@ -6,9 +6,9 @@ import (
 	"fmt"
 	"net/http"
 
+	"github.com/rs/zerolog/log"
 	"github.com/yusing/go-proxy/agent/pkg/env"
 	"github.com/yusing/go-proxy/agent/pkg/handler"
-	"github.com/yusing/go-proxy/internal/logging"
 	"github.com/yusing/go-proxy/internal/net/gphttp/server"
 	"github.com/yusing/go-proxy/internal/task"
 )
@@ -33,12 +33,11 @@ func StartAgentServer(parent task.Parent, opt Options) {
 		tlsConfig.ClientAuth = tls.NoClientCert
 	}
 
-	logger := logging.GetLogger()
 	agentServer := &http.Server{
 		Addr:      fmt.Sprintf(":%d", opt.Port),
 		Handler:   handler.NewAgentHandler(),
 		TLSConfig: tlsConfig,
 	}
 
-	server.Start(parent, agentServer, nil, logger)
+	server.Start(parent, agentServer, nil, &log.Logger)
 }