feat: add PreDialPlugin

Author: andeya

peer.go

@@ -207,11 +207,15 @@ func (p *peer) CountSession() int {
 
 // Dial connects with the peer of the destination address.
 func (p *peer) Dial(addr string, protoFunc ...ProtoFunc) (Session, *Status) {
+	stat := p.pluginContainer.preDial(p.dialer.localAddr, addr)
+	if !stat.OK() {
+		return nil, stat
+	}
 	var sess = newSession(p, nil, protoFunc)
 	_, err := p.dialer.dialWithRetry(addr, "", func(conn net.Conn) error {
 		sess.socket.Reset(conn, protoFunc...)
 		sess.socket.SetID(sess.LocalAddr().String())
-		if stat := p.pluginContainer.postDial(sess, false); !stat.OK() {
+		if stat = p.pluginContainer.postDial(sess, false); !stat.OK() {
 			conn.Close()
 			return stat.Cause()
 		}
@@ -227,23 +231,26 @@ func (p *peer) Dial(addr string, protoFunc ...ProtoFunc) (Session, *Status) {
 			oldID := sess.ID()
 			oldIP := sess.LocalAddr().String()
 			oldConn := sess.getConn()
-
-			_, err := p.dialer.dialWithRetry(addr, oldID, func(conn net.Conn) error {
-				sess.socket.Reset(conn, protoFunc...)
-				if oldIP == oldID {
-					sess.socket.SetID(sess.LocalAddr().String())
-				} else {
-					sess.socket.SetID(oldID)
-				}
-				sess.changeStatus(statusPreparing)
-				if stat := p.pluginContainer.postDial(sess, true); !stat.OK() {
-					conn.Close()
-					sess.changeStatus(statusRedialing)
-					return stat.Cause()
-				}
-				return nil
-			})
-
+			var err error
+			if stat := p.pluginContainer.preDial(p.dialer.localAddr, addr); stat.OK() {
+				_, err = p.dialer.dialWithRetry(addr, oldID, func(conn net.Conn) error {
+					sess.socket.Reset(conn, protoFunc...)
+					if oldIP == oldID {
+						sess.socket.SetID(sess.LocalAddr().String())
+					} else {
+						sess.socket.SetID(oldID)
+					}
+					sess.changeStatus(statusPreparing)
+					if stat := p.pluginContainer.postDial(sess, true); !stat.OK() {
+						conn.Close()
+						sess.changeStatus(statusRedialing)
+						return stat.Cause()
+					}
+					return nil
+				})
+			} else {
+				err = stat.Cause()
+			}
 			if err != nil {
 				sess.closeLocked()
 				sess.tryChangeStatus(statusRedialFailed, statusRedialing)

plugin.go

@@ -49,6 +49,11 @@ type (
 		Plugin
 		PostListen(net.Addr) error
 	}
+	// PreDialPlugin is executed before dialing.
+	PreDialPlugin interface {
+		Plugin
+		PreDial(localAddr net.Addr, remoteAddr string) *Status
+	}
 	// PostDialPlugin is executed after dialing.
 	PostDialPlugin interface {
 		Plugin
@@ -362,7 +367,32 @@ func (p *pluginSingleContainer) postListen(addr net.Addr) {
 	return
 }
 
-// PostDial executes the defined plugins after dialing.
+// preDial executes the defined plugins before dialing.
+func (p *pluginSingleContainer) preDial(localAddr net.Addr, remoteAddr string) (stat *Status) {
+	var pluginName string
+	defer func() {
+		if p := recover(); p != nil {
+			Errorf("[PreDialPlugin:%s] network:%s, localAddr%s ,remoteAddr:%s, panic:%v\n%s", pluginName, localAddr.Network(), localAddr.String(), remoteAddr, p, goutil.PanicTrace(2))
+			stat = statDialFailed.Copy(p)
+		}
+	}()
+	for _, plugin := range p.plugins {
+		if _plugin, ok := plugin.(PreDialPlugin); ok {
+			pluginName = plugin.Name()
+			if stat = _plugin.PreDial(localAddr, remoteAddr); !stat.OK() {
+				LazyDebugf(func() string {
+					return fmt.Sprintf("[PreDialPlugin:%s] network:%s, localAddr%s ,remoteAddr:%s, error:%s",
+						pluginName, localAddr.Network(), localAddr.String(), remoteAddr, stat.String(),
+					)
+				})
+				return stat
+			}
+		}
+	}
+	return nil
+}
+
+// postDial executes the defined plugins after dialing.
 func (p *pluginSingleContainer) postDial(sess PreSession, isRedial bool) (stat *Status) {
 	var pluginName string
 	defer func() {

plugin_impl.go

@@ -8,6 +8,7 @@ var (
 	_ PostNewPeerPlugin         = (*PluginImpl)(nil)
 	_ PostRegPlugin             = (*PluginImpl)(nil)
 	_ PostListenPlugin          = (*PluginImpl)(nil)
+	_ PreDialPlugin             = (*PluginImpl)(nil)
 	_ PostDialPlugin            = (*PluginImpl)(nil)
 	_ PostAcceptPlugin          = (*PluginImpl)(nil)
 	_ PreWriteCallPlugin        = (*PluginImpl)(nil)
@@ -41,6 +42,8 @@ type PluginImpl struct {
 	OnPostReg func(*Handler) error
 	// OnPostListen is called after a listener is created.
 	OnPostListen func(net.Addr) error
+	// OnPreDial is called before a dial is created.
+	OnPreDial func(localAddr net.Addr, remoteAddr string) *Status
 	// OnPostDial is called after a dial is created.
 	OnPostDial func(sess PreSession, isRedial bool) *Status
 	// OnPostAccept is called after a session is accepted.
@@ -118,6 +121,14 @@ func (p *PluginImpl) PostListen(addr net.Addr) error {
 	return p.OnPostListen(addr)
 }
 
+// PreDial is called before a dial is created.
+func (p *PluginImpl) PreDial(localAddr net.Addr, remoteAddr string) *Status {
+	if p.OnPreDial == nil {
+		return nil
+	}
+	return p.OnPreDial(localAddr, remoteAddr)
+}
+
 // PostDial is called after a dial is created.
 func (p *PluginImpl) PostDial(sess PreSession, isRedial bool) *Status {
 	if p.OnPostDial == nil {