fix agent connect

cmd/dashboard/controller/controller.go

@@ -23,7 +23,7 @@ import (
 	"github.com/naiba/nezha/service/singleton"
 )
 
-func ServeWeb() *http.Server {
+func ServeWeb() http.Handler {
 	gin.SetMode(gin.ReleaseMode)
 	r := gin.Default()
 	docs.SwaggerInfo.BasePath = "/api/v1"
@@ -39,10 +39,7 @@ func ServeWeb() *http.Server {
 	r.Use(recordPath)
 	routers(r)
 
-	return &http.Server{
+	return r
-		ReadHeaderTimeout: time.Second * 5,
-		Handler:           r,
-	}
 }
 
 func routers(r *gin.Engine) {

cmd/dashboard/main.go

@@ -5,13 +5,16 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"net/http"
+	"strings"
 	"time"
 	_ "time/tzdata"
 
 	"github.com/ory/graceful"
-	"github.com/soheilhy/cmux"
 	flag "github.com/spf13/pflag"
 	"golang.org/x/crypto/bcrypt"
+	"golang.org/x/net/http2"
+	"golang.org/x/net/http2/h2c"
 
 	"github.com/naiba/nezha/cmd/dashboard/controller"
 	"github.com/naiba/nezha/cmd/dashboard/rpc"
@@ -110,37 +113,33 @@ func main() {
 		log.Fatal(err)
 	}
 
-	m := cmux.New(l)
-	grpcL := m.Match(cmux.HTTP2HeaderField("content-type", "application/grpc"))
-	httpL := m.Match(cmux.HTTP1Fast())
-
 	singleton.CleanMonitorHistory()
-	go rpc.ServeRPC(grpcL)
 	serviceSentinelDispatchBus := make(chan model.Monitor) // 用于传递服务监控任务信息的channel
 	go rpc.DispatchTask(serviceSentinelDispatchBus)
 	go rpc.DispatchKeepalive()
 	go singleton.AlertSentinelStart()
 	singleton.NewServiceSentinel(serviceSentinelDispatchBus)
-	srv := controller.ServeWeb()
+
+	grpcHandler := rpc.ServeRPC()
+	httpHandler := controller.ServeWeb()
+
+	mixedHandler := newHTTPandGRPCMux(httpHandler, grpcHandler)
+	http2Server := &http2.Server{}
+	http1Server := &http.Server{Handler: h2c.NewHandler(mixedHandler, http2Server)}
 
 	go dispatchReportInfoTask()
 
-	go func() {
+	if err := graceful.Graceful(func() error {
-		if err := graceful.Graceful(func() error {
+		log.Println("NEZHA>> Dashboard::START", singleton.Conf.ListenPort)
-			log.Println("NEZHA>> Dashboard::START", singleton.Conf.ListenPort)
+		return http1Server.Serve(l)
-			return srv.Serve(httpL)
+	}, func(c context.Context) error {
-		}, func(c context.Context) error {
+		log.Println("NEZHA>> Graceful::START")
-			log.Println("NEZHA>> Graceful::START")
+		singleton.RecordTransferHourlyUsage()
-			singleton.RecordTransferHourlyUsage()
+		log.Println("NEZHA>> Graceful::END")
-			log.Println("NEZHA>> Graceful::END")
+		return l.Close()
-			m.Close()
+	}); err != nil {
-			return nil
+		log.Printf("NEZHA>> ERROR: %v", err)
-		}); err != nil {
+	}
-			log.Printf("NEZHA>> ERROR: %v", err)
-		}
-	}()
-
-	m.Serve()
 }
 
 func dispatchReportInfoTask() {
@@ -157,3 +156,13 @@ func dispatchReportInfoTask() {
 		})
 	}
 }
+
+func newHTTPandGRPCMux(httpHandler http.Handler, grpcHandler http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("content-type"), "application/grpc") {
+			grpcHandler.ServeHTTP(w, r)
+			return
+		}
+		httpHandler.ServeHTTP(w, r)
+	})
+}

cmd/dashboard/rpc/rpc.go

@@ -1,7 +1,7 @@
 package rpc
 
 import (
-	"net"
+	"net/http"
 
 	"google.golang.org/grpc"
 
@@ -11,11 +11,11 @@ import (
 	"github.com/naiba/nezha/service/singleton"
 )
 
-func ServeRPC(l net.Listener) {
+func ServeRPC() http.Handler {
 	server := grpc.NewServer()
 	rpcService.NezhaHandlerSingleton = rpcService.NewNezhaHandler()
 	pb.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton)
-	server.Serve(l)
+	return server
 }
 
 func DispatchTask(serviceSentinelDispatchBus <-chan model.Monitor) {

go.mod

@@ -20,7 +20,6 @@ require (
 	github.com/oschwald/maxminddb-golang v1.13.1
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/robfig/cron/v3 v3.0.1
-	github.com/soheilhy/cmux v0.1.5
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.18.2
 	github.com/swaggo/files v1.0.1

go.sum

@@ -154,8 +154,6 @@ github.com/sagikazarmark/locafero v0.4.0 h1:HApY1R9zGo4DBgr7dqsTH/JJxLTTsOt7u6ke
 github.com/sagikazarmark/locafero v0.4.0/go.mod h1:Pe1W6UlPYUk/+wc/6KFhbORCfqzgYEpgQ3O5fPuL3H4=
 github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE=
 github.com/sagikazarmark/slog-shim v0.1.0/go.mod h1:SrcSrq8aKtyuqEI1uvTDTK1arOWRIczQRv+GVI1AkeQ=
-github.com/soheilhy/cmux v0.1.5 h1:jjzc5WVemNEDTLwv9tlmemhC73tI08BNOIGwBOo10Js=
-github.com/soheilhy/cmux v0.1.5/go.mod h1:T7TcVDs9LWfQgPlPsdngu6I6QIoyIFZDDC6sNE1GqG0=
 github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=
 github.com/sourcegraph/conc v0.3.0/go.mod h1:Sdozi7LEKbFPqYX2/J+iBAM6HpqSLTASQIKqDmF7Mt0=
 github.com/spf13/afero v1.11.0 h1:WJQKhtpdm3v2IzqG8VMqrr6Rf3UYpEF239Jy9wNepM8=
@@ -208,7 +206,6 @@ go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTV
 golang.org/x/arch v0.9.0 h1:ub9TgUInamJ8mrZIGlBG6/4TqWeMszd4N8lNorbrr6k=
 golang.org/x/arch v0.9.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
-golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/crypto v0.26.0 h1:RrRspgV4mU+YwB4FYnuBoKsUapNIL5cohGAmSH3azsw=
@@ -218,9 +215,7 @@ golang.org/x/exp v0.0.0-20230905200255-921286631fa9/go.mod h1:S2oDrQGGwySpoQPVqR
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/mod v0.18.0 h1:5+9lSbEzPSdWkH32vYPBwEpX8KwDbM52Ud9xBUvNlb0=
 golang.org/x/mod v0.18.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
-golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
-golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20210421230115-4e50805a0758/go.mod h1:72T/g9IO56b78aLF+1Kcs5dz7/ng1VjMUvfKvpfy+jM=
 golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
@@ -232,8 +227,6 @@ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ=
 golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
-golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210420072515-93ed5bcd2bfe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

model/config.go

@@ -48,28 +48,28 @@ const (
 )
 
 type Config struct {
-	Debug bool // debug模式开关
+	Debug bool `mapstructure:"debug" json:"debug,omitempty"` // debug模式开关
 
-	Language       string // 系统语言，默认 zh-CN
+	Language       string `mapstructure:"language" json:"language,omitempty"` // 系统语言，默认 zh-CN
-	SiteName       string
+	SiteName       string `mapstructure:"site_name" json:"site_name,omitempty"`
-	JWTSecretKey   string
+	JWTSecretKey   string `mapstructure:"jwt_secret_key" json:"jwt_secret_key,omitempty"`
-	AgentSecretKey string
+	AgentSecretKey string `mapstructure:"agent_secret_key" json:"agent_secret_key,omitempty"`
-	ListenPort     uint
+	ListenPort     uint   `mapstructure:"listen_port" json:"listen_port,omitempty"`
-	InstallHost    string
+	InstallHost    string `mapstructure:"install_host" json:"install_host,omitempty"`
-	TLS            bool
+	TLS            bool   `mapstructure:"tls" json:"tls,omitempty"`
-	Location       string // 时区，默认为 Asia/Shanghai
+	Location       string `mapstructure:"location" json:"location,omitempty"` // 时区，默认为 Asia/Shanghai
 
-	EnablePlainIPInNotification bool // 通知信息IP不打码
+	EnablePlainIPInNotification bool `mapstructure:"enable_plain_ip_in_notification" json:"enable_plain_ip_in_notification,omitempty"` // 通知信息IP不打码
 
 	// IP变更提醒
-	EnableIPChangeNotification bool
+	EnableIPChangeNotification bool   `mapstructure:"enable_ip_change_notification" json:"enable_ip_change_notification,omitempty"`
-	IPChangeNotificationTag    string
+	IPChangeNotificationTag    string `mapstructure:"ip_change_notification_tag" json:"ip_change_notification_tag,omitempty"`
-	Cover                      uint8  // 覆盖范围（0:提醒未被 IgnoredIPNotification 包含的所有服务器; 1:仅提醒被 IgnoredIPNotification 包含的服务器;）
+	Cover                      uint8  `mapstructure:"cover" json:"cover,omitempty"`                                     // 覆盖范围（0:提醒未被 IgnoredIPNotification 包含的所有服务器; 1:仅提醒被 IgnoredIPNotification 包含的服务器;）
-	IgnoredIPNotification      string // 特定服务器IP（多个服务器用逗号分隔）
+	IgnoredIPNotification      string `mapstructure:"ignored_ip_notification" json:"ignored_ip_notification,omitempty"` // 特定服务器IP（多个服务器用逗号分隔）
 
-	IgnoredIPNotificationServerIDs map[uint64]bool // [ServerID] -> bool(值为true代表当前ServerID在特定服务器列表内）
+	IgnoredIPNotificationServerIDs map[uint64]bool `mapstructure:"ignored_ip_notification_server_i_ds" json:"ignored_ip_notification_server_i_ds,omitempty"` // [ServerID] -> bool(值为true代表当前ServerID在特定服务器列表内）
-	AvgPingCount                   int
+	AvgPingCount                   int             `mapstructure:"avg_ping_count" json:"avg_ping_count,omitempty"`
-	DNSServers                     string
+	DNSServers                     string          `mapstructure:"dns_servers" json:"dns_servers,omitempty"`
 
 	v *viper.Viper
 }

service/rpc/auth.go

@@ -8,6 +8,7 @@ import (
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 
+	"github.com/hashicorp/go-uuid"
 	"github.com/naiba/nezha/model"
 	"github.com/naiba/nezha/service/singleton"
 )
@@ -17,14 +18,6 @@ type authHandler struct {
 	ClientUUID   string
 }
 
-func (a *authHandler) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
-	return map[string]string{"client_secret": a.ClientSecret, "client_uuid": a.ClientUUID}, nil
-}
-
-func (a *authHandler) RequireTransportSecurity() bool {
-	return false
-}
-
 func (a *authHandler) Check(ctx context.Context) (uint64, error) {
 	md, ok := metadata.FromIncomingContext(ctx)
 	if !ok {
@@ -45,6 +38,10 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
 		clientUUID = value[0]
 	}
 
+	if _, err := uuid.ParseUUID(clientUUID); err != nil {
+		return 0, status.Errorf(codes.Unauthenticated, "客户端 UUID 不合法")
+	}
+
 	singleton.ServerLock.RLock()
 	defer singleton.ServerLock.RUnlock()
 	clientID, hasID := singleton.ServerUUIDToID[clientUUID]
@@ -58,6 +55,7 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) {
 		s.TaskCloseLock = new(sync.Mutex)
 		singleton.ServerList[s.ID] = &s
 		singleton.ServerUUIDToID[clientUUID] = s.ID
+		clientID = s.ID
 	}
 
 	return clientID, nil