agent 增加 SSL/TLS 选项

This commit is contained in:
Lemoe 2021-11-06 09:46:44 +08:00
parent f63d8ba613
commit e28e7f4375
No known key found for this signature in database
GPG Key ID: 2570C7D85A4A74CA

View File

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