From be85633c3295a5da74894181ff441f9d33edb5ff Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Sat, 29 Mar 2025 16:44:16 +0800
Subject: [PATCH] fix(agent): fix agent host validatation and improve file path
 handling

---
 agent/pkg/agent/config.go    |  7 ++++++-
 agent/pkg/certs/zip.go       | 12 ++++++++++--
 internal/api/v1/new_agent.go |  4 ++--
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/agent/pkg/agent/config.go b/agent/pkg/agent/config.go
index 83243a2..e727905 100644
--- a/agent/pkg/agent/config.go
+++ b/agent/pkg/agent/config.go
@@ -131,7 +131,12 @@ func (cfg *AgentConfig) StartWithCerts(parent task.Parent, ca, crt, key []byte)
 }
 
 func (cfg *AgentConfig) Start(parent task.Parent) gperr.Error {
-	certData, err := os.ReadFile(certs.AgentCertsFilename(cfg.Addr))
+	filepath, ok := certs.AgentCertsFilepath(cfg.Addr)
+	if !ok {
+		return gperr.New("invalid agent host").Subject(cfg.Addr)
+	}
+
+	certData, err := os.ReadFile(filepath)
 	if err != nil {
 		return gperr.Wrap(err, "failed to read agent certs")
 	}
diff --git a/agent/pkg/certs/zip.go b/agent/pkg/certs/zip.go
index 9349499..61db6f8 100644
--- a/agent/pkg/certs/zip.go
+++ b/agent/pkg/certs/zip.go
@@ -7,6 +7,7 @@ import (
 	"path/filepath"
 
 	"github.com/yusing/go-proxy/internal/common"
+	"github.com/yusing/go-proxy/internal/utils/strutils"
 )
 
 func writeFile(zipWriter *zip.Writer, name string, data []byte) error {
@@ -50,8 +51,15 @@ func ZipCert(ca, crt, key []byte) ([]byte, error) {
 	return data.Bytes(), nil
 }
 
-func AgentCertsFilename(host string) string {
-	return filepath.Join(common.AgentCertsBasePath, host+".zip")
+func isValidAgentHost(host string) bool {
+	return strutils.IsValidFilename(host + ".zip")
+}
+
+func AgentCertsFilepath(host string) (filepathOut string, ok bool) {
+	if !isValidAgentHost(host) {
+		return "", false
+	}
+	return filepath.Join(common.AgentCertsBasePath, host+".zip"), true
 }
 
 func ExtractCert(data []byte) (ca, crt, key []byte, err error) {
diff --git a/internal/api/v1/new_agent.go b/internal/api/v1/new_agent.go
index 7c381f1..e4d67e8 100644
--- a/internal/api/v1/new_agent.go
+++ b/internal/api/v1/new_agent.go
@@ -126,8 +126,8 @@ func VerifyNewAgent(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	filename := certs.AgentCertsFilename(data.Host)
-	if !strutils.IsValidFilename(filename) {
+	filename, ok := certs.AgentCertsFilepath(data.Host)
+	if !ok {
 		gphttp.ClientError(w, gphttp.ErrInvalidKey("host"))
 		return
 	}