generated from cloudwego/.github
-
Notifications
You must be signed in to change notification settings - Fork 14
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(prob): support protocol probing
- Loading branch information
Showing
3 changed files
with
146 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,82 @@ | ||
package dubbo | ||
|
||
import ( | ||
"context" | ||
"errors" | ||
"net" | ||
|
||
"github.com/cloudwego/kitex/pkg/endpoint" | ||
"github.com/cloudwego/kitex/pkg/remote" | ||
"github.com/cloudwego/kitex/pkg/remote/trans/detection" | ||
"github.com/cloudwego/kitex/pkg/remote/trans/netpoll" | ||
"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2" | ||
cnetpoll "github.com/cloudwego/netpoll" | ||
|
||
"github.com/kitex-contrib/codec-dubbo/pkg/dubbo_spec" | ||
) | ||
|
||
// NewSvrTransHandlerFactory the factory expand the implementation of DetectableServerTransHandler for | ||
// hessian protocol to support the dubbo protocol probing. | ||
func NewSvrTransHandlerFactory(opts ...Option) remote.ServerTransHandlerFactory { | ||
return detection.NewSvrTransHandlerFactory(netpoll.NewSvrTransHandlerFactory(), | ||
nphttp2.NewSvrTransHandlerFactory(), | ||
&svrTransHandlerFactory{ | ||
codec: NewDubboCodec(opts...), | ||
ServerTransHandlerFactory: netpoll.NewSvrTransHandlerFactory(), | ||
}) | ||
} | ||
|
||
type svrTransHandlerFactory struct { | ||
remote.ServerTransHandlerFactory | ||
// the codec should be set with dubbo codec to keep consistent | ||
// with the function ProtocolMatch. | ||
codec remote.Codec | ||
} | ||
|
||
func (f *svrTransHandlerFactory) NewTransHandler(opt *remote.ServerOption) (remote.ServerTransHandler, error) { | ||
sourceCodec := opt.Codec | ||
opt.Codec = f.codec | ||
defer func() { | ||
// just set the handler with dubbo codec and restore the source codec | ||
opt.Codec = sourceCodec | ||
}() | ||
|
||
handler, err := f.ServerTransHandlerFactory.NewTransHandler(opt) | ||
if err != nil { | ||
return nil, err | ||
} | ||
return &svrTransHandler{ | ||
ServerTransHandler: handler, | ||
}, nil | ||
} | ||
|
||
type svrTransHandler struct { | ||
remote.ServerTransHandler | ||
} | ||
|
||
func (svr *svrTransHandler) ProtocolMatch(ctx context.Context, conn net.Conn) (err error) { | ||
// Check the validity of client preface. | ||
npReader := conn.(interface{ Reader() cnetpoll.Reader }).Reader() | ||
// read at most avoid block | ||
header, err := npReader.Peek(dubbo_spec.HEADER_SIZE) | ||
if err != nil { | ||
return err | ||
} | ||
if header[0] == dubbo_spec.MAGIC_HIGH && header[1] == dubbo_spec.MAGIC_LOW { | ||
return nil | ||
} | ||
return errors.New("error protocol not match hessian") | ||
} | ||
|
||
func (svr *svrTransHandler) GracefulShutdown(ctx context.Context) error { | ||
if g, ok := svr.ServerTransHandler.(remote.GracefulShutdown); ok { | ||
g.GracefulShutdown(ctx) | ||
} | ||
return nil | ||
} | ||
func (svr *svrTransHandler) SetInvokeHandleFunc(inkHdlFunc endpoint.Endpoint) { | ||
if s, ok := svr.ServerTransHandler.(remote.InvokeHandleFuncSetter); ok { | ||
s.SetInvokeHandleFunc(inkHdlFunc) | ||
} | ||
return | ||
} |