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 7 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()
httpHandler := controller.ServeWeb()
mixedHandler := newHTTPandGRPCMux(httpHandler, grpcHandler)
muxHandler := newHTTPandGRPCMux(httpHandler, grpcHandler)
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()
if err := graceful.Graceful(func() error {
log.Println("NEZHA>> Dashboard::START", singleton.Conf.ListenPort)
return http1Server.Serve(l)
return muxServer.Serve(l)
}, func(c context.Context) error {
log.Println("NEZHA>> Graceful::START")
singleton.RecordTransferHourlyUsage()
log.Println("NEZHA>> Graceful::END")
return l.Close()
return muxServer.Shutdown(c)
}); err != nil {
log.Printf("NEZHA>> ERROR: %v", err)
}
@ -159,7 +159,8 @@ 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") {
if r.ProtoMajor == 2 && r.Header.Get("Content-Type") == "application/grpc" &&
strings.HasPrefix(r.URL.Path, "/"+proto.NezhaService_ServiceDesc.ServiceName) {
grpcHandler.ServeHTTP(w, r)
return
}

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

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