diff --git a/netLayer/dial.go b/netLayer/dial.go index 8932ec1..47af47b 100644 --- a/netLayer/dial.go +++ b/netLayer/dial.go @@ -7,12 +7,16 @@ import ( "time" ) +var ( + CustomDialerMap = make(map[string]func(address string, timeout time.Duration) (net.Conn, error)) +) + func (a *Addr) Dial() (net.Conn, error) { var istls bool var resultConn net.Conn var err error - switch a.Network { + switch n := a.Network; n { case "": a.Network = "tcp" goto tcp @@ -30,6 +34,11 @@ func (a *Addr) Dial() (net.Conn, error) { return DialUDP(ua) default: + if len(CustomDialerMap) > 0 { + if f := CustomDialerMap[n]; f != nil { + return f(a.String(), time.Second*15) + } + } goto defaultPart diff --git a/netLayer/listen.go b/netLayer/listen.go index f7cbc42..559ea8a 100644 --- a/netLayer/listen.go +++ b/netLayer/listen.go @@ -13,6 +13,10 @@ import ( "go.uber.org/zap" ) +var ( + CustomListenerMap = make(map[string]func(address string) (net.Listener, error)) +) + func loopAccept(listener net.Listener, xver int, acceptFunc func(net.Conn)) { if xver > 0 { @@ -147,6 +151,12 @@ func ListenAndAccept(network, addr string, sockopt *Sockopt, xver int, acceptFun fallthrough default: + if len(CustomListenerMap) > 0 { + if f := CustomListenerMap[network]; f != nil { + return f(addr) + } + } + listener, err = net.Listen(network, addr) if err != nil { return