diff --git a/cmd/agent/main.go b/cmd/agent/main.go index 71510c5..ffe5c1f 100644 --- a/cmd/agent/main.go +++ b/cmd/agent/main.go @@ -2,6 +2,7 @@ package main import ( "context" + "crypto/tls" "encoding/json" "errors" "fmt" @@ -18,6 +19,7 @@ import ( "github.com/p14yground/go-github-selfupdate/selfupdate" flag "github.com/spf13/pflag" "google.golang.org/grpc" + "google.golang.org/grpc/credentials" "github.com/naiba/nezha/cmd/agent/monitor" "github.com/naiba/nezha/cmd/agent/processgroup" @@ -38,6 +40,7 @@ type AgentConfig struct { Server string ClientSecret string ReportDelay int + TLS bool } var ( @@ -80,6 +83,7 @@ func main() { flag.BoolVar(&agentConf.DisableCommandExecute, "disable-command-execute", false, "禁止在此机器上执行命令") flag.BoolVar(&agentConf.DisableAutoUpdate, "disable-auto-update", false, "禁用自动升级") flag.BoolVar(&agentConf.DisableForceUpdate, "disable-force-update", false, "禁用强制升级") + flag.BoolVar(&agentConf.TLS, "tls", false, "启用SSL/TLS加密") flag.Parse() if agentConf.ClientSecret == "" { @@ -138,7 +142,13 @@ func run() { for { timeOutCtx, cancel := context.WithTimeout(context.Background(), networkTimeOut) - conn, err = grpc.DialContext(timeOutCtx, agentConf.Server, grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth)) + var securityOption grpc.DialOption + if agentConf.TLS { + securityOption = grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{})) + } else { + securityOption = grpc.WithInsecure() + } + conn, err = grpc.DialContext(timeOutCtx, agentConf.Server, securityOption, grpc.WithPerRPCCredentials(&auth)) if err != nil { println("与面板建立连接失败:", err) cancel()