tcjlj/nezha
1
0
Fork 0
mirror of https://github.com/nezhahq/nezha.git synced 2025-01-22 12:48:14 -05:00
Code Issues Actions Packages Projects Releases Wiki Activity

add path check for multiplexer (#451)

Browse Source
This commit is contained in:
UUBulb 2024-10-23 12:55:10 +08:00 committed by GitHub
parent cb9436a8f7
commit 8d0f6fb7c4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 7 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
cmd/dashboard/main.go
View File

@ -123,20 +123,20 @@ func main() {
grpcHandler := rpc.ServeRPC() grpcHandler := rpc.ServeRPC()
httpHandler := controller.ServeWeb() httpHandler := controller.ServeWeb()
mixedHandler := newHTTPandGRPCMux(httpHandler, grpcHandler) muxHandler := newHTTPandGRPCMux(httpHandler, grpcHandler)
http2Server := &http2.Server{} http2Server := &http2.Server{}
http1Server := &http.Server{Handler: h2c.NewHandler(mixedHandler, http2Server), ReadHeaderTimeout: time.Second * 5} muxServer := &http.Server{Handler: h2c.NewHandler(muxHandler, http2Server), ReadHeaderTimeout: time.Second * 5}
go dispatchReportInfoTask() go dispatchReportInfoTask()
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 muxServer.Serve(l)
}, 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() return muxServer.Shutdown(c)
}); err != nil { }); err != nil {
log.Printf("NEZHA>> ERROR: %v", err) log.Printf("NEZHA>> ERROR: %v", err)
} }
@ -159,7 +159,8 @@ func dispatchReportInfoTask() {
func newHTTPandGRPCMux(httpHandler http.Handler, grpcHandler http.Handler) http.Handler { func newHTTPandGRPCMux(httpHandler http.Handler, grpcHandler http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("content-type"), "application/grpc") { if r.ProtoMajor == 2 && r.Header.Get("Content-Type") == "application/grpc" &&
strings.HasPrefix(r.URL.Path, "/"+proto.NezhaService_ServiceDesc.ServiceName) {
grpcHandler.ServeHTTP(w, r) grpcHandler.ServeHTTP(w, r)
return return
} }

4
cmd/dashboard/rpc/rpc.go
View File

@ -1,8 +1,6 @@
package rpc package rpc
import ( import (
"net/http"
"google.golang.org/grpc" "google.golang.org/grpc"
"github.com/naiba/nezha/model" "github.com/naiba/nezha/model"
@ -11,7 +9,7 @@ import (
"github.com/naiba/nezha/service/singleton" "github.com/naiba/nezha/service/singleton"
) )
func ServeRPC() http.Handler { func ServeRPC() *grpc.Server {
server := grpc.NewServer() server := grpc.NewServer()
rpcService.NezhaHandlerSingleton = rpcService.NewNezhaHandler() rpcService.NezhaHandlerSingleton = rpcService.NewNezhaHandler()
pb.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton) pb.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton)