add router handler option

This commit is contained in:
ginuerzh
2021-12-31 00:03:24 +08:00
parent a430384bba
commit 9769efe33c
35 changed files with 98 additions and 281 deletions

View File

@ -31,9 +31,9 @@ func init() {
}
type forwardHandler struct {
chain *chain.Chain
bypass bypass.Bypass
config *ssh.ServerConfig
router *chain.Router
logger logger.Logger
md metadata
}
@ -46,6 +46,7 @@ func NewHandler(opts ...handler.Option) handler.Handler {
return &forwardHandler{
bypass: options.Bypass,
router: options.Router,
logger: options.Logger,
}
}
@ -71,11 +72,6 @@ func (h *forwardHandler) Init(md md.Metadata) (err error) {
return nil
}
// WithChain implements chain.Chainable interface
func (h *forwardHandler) WithChain(chain *chain.Chain) {
h.chain = chain
}
func (h *forwardHandler) Handle(ctx context.Context, conn net.Conn) {
defer conn.Close()
@ -168,11 +164,7 @@ func (h *forwardHandler) directPortForwardChannel(ctx context.Context, channel s
return
}
r := (&chain.Router{}).
WithChain(h.chain).
// WithRetry(h.md.retryCount).
WithLogger(h.logger)
conn, err := r.Dial(ctx, "tcp", raddr)
conn, err := h.router.Dial(ctx, "tcp", raddr)
if err != nil {
return
}