From 1a6cbbb2d2069a307ed498c21e09a0db43f26ab6 Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Thu, 17 May 2018 00:07:56 +0800
Subject: [PATCH] return error quickly if nathole make error

---
 client/proxy.go   |  6 ++++++
 client/visitor.go |  5 +++++
 models/msg/msg.go |  1 +
 server/nathole.go | 31 +++++++++++++++++++++++--------
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/client/proxy.go b/client/proxy.go
index a4f99c6..26c9a66 100644
--- a/client/proxy.go
+++ b/client/proxy.go
@@ -272,6 +272,12 @@ func (pxy *XtcpProxy) InWorkConn(conn frpNet.Conn) {
 	}
 	clientConn.SetReadDeadline(time.Time{})
 	clientConn.Close()
+
+	if natHoleRespMsg.Error != "" {
+		pxy.Error("natHoleRespMsg get error info: %s", natHoleRespMsg.Error)
+		return
+	}
+
 	pxy.Trace("get natHoleRespMsg, sid [%s], client address [%s]", natHoleRespMsg.Sid, natHoleRespMsg.ClientAddr)
 
 	// Send sid to visitor udp address.
diff --git a/client/visitor.go b/client/visitor.go
index d07eb8e..3587d49 100644
--- a/client/visitor.go
+++ b/client/visitor.go
@@ -235,6 +235,11 @@ func (sv *XtcpVisitor) handleConn(userConn frpNet.Conn) {
 	visitorConn.SetReadDeadline(time.Time{})
 	pool.PutBuf(buf)
 
+	if natHoleRespMsg.Error != "" {
+		sv.Error("natHoleRespMsg get error info: %s", natHoleRespMsg.Error)
+		return
+	}
+
 	sv.Trace("get natHoleRespMsg, sid [%s], client address [%s]", natHoleRespMsg.Sid, natHoleRespMsg.ClientAddr)
 
 	// Close visitorConn, so we can use it's local address.
diff --git a/models/msg/msg.go b/models/msg/msg.go
index 93214fc..fd3e656 100644
--- a/models/msg/msg.go
+++ b/models/msg/msg.go
@@ -163,6 +163,7 @@ type NatHoleResp struct {
 	Sid         string `json:"sid"`
 	VisitorAddr string `json:"visitor_addr"`
 	ClientAddr  string `json:"client_addr"`
+	Error       string `json:"error"`
 }
 
 type NatHoleSid struct {
diff --git a/server/nathole.go b/server/nathole.go
index 0ad0c48..cc8339f 100644
--- a/server/nathole.go
+++ b/server/nathole.go
@@ -108,12 +108,16 @@ func (nc *NatHoleController) HandleVisitor(m *msg.NatHoleVisitor, raddr *net.UDP
 	clientCfg, ok := nc.clientCfgs[m.ProxyName]
 	if !ok {
 		nc.mu.Unlock()
-		log.Debug("xtcp server for [%s] doesn't exist", m.ProxyName)
+		errInfo := fmt.Sprintf("xtcp server for [%s] doesn't exist", m.ProxyName)
+		log.Debug(errInfo)
+		nc.listener.WriteToUDP(nc.GenNatHoleResponse(nil, errInfo), raddr)
 		return
 	}
 	if m.SignKey != util.GetAuthKey(clientCfg.Sk, m.Timestamp) {
 		nc.mu.Unlock()
-		log.Debug("xtcp connection of [%s] auth failed", m.ProxyName)
+		errInfo := fmt.Sprintf("xtcp connection of [%s] auth failed", m.ProxyName)
+		log.Debug(errInfo)
+		nc.listener.WriteToUDP(nc.GenNatHoleResponse(nil, errInfo), raddr)
 		return
 	}
 
@@ -137,7 +141,7 @@ func (nc *NatHoleController) HandleVisitor(m *msg.NatHoleVisitor, raddr *net.UDP
 	// Wait client connections.
 	select {
 	case <-session.NotifyCh:
-		resp := nc.GenNatHoleResponse(raddr, session)
+		resp := nc.GenNatHoleResponse(session, "")
 		log.Trace("send nat hole response to visitor")
 		nc.listener.WriteToUDP(resp, raddr)
 	case <-time.After(time.Duration(NatHoleTimeout) * time.Second):
@@ -156,16 +160,27 @@ func (nc *NatHoleController) HandleClient(m *msg.NatHoleClient, raddr *net.UDPAd
 	session.ClientAddr = raddr
 	session.NotifyCh <- struct{}{}
 
-	resp := nc.GenNatHoleResponse(raddr, session)
+	resp := nc.GenNatHoleResponse(session, "")
 	log.Trace("send nat hole response to client")
 	nc.listener.WriteToUDP(resp, raddr)
 }
 
-func (nc *NatHoleController) GenNatHoleResponse(raddr *net.UDPAddr, session *NatHoleSession) []byte {
+func (nc *NatHoleController) GenNatHoleResponse(session *NatHoleSession, errInfo string) []byte {
+	var (
+		sid         string
+		visitorAddr string
+		clientAddr  string
+	)
+	if session != nil {
+		sid = session.Sid
+		visitorAddr = session.VisitorAddr.String()
+		clientAddr = session.ClientAddr.String()
+	}
 	m := &msg.NatHoleResp{
-		Sid:         session.Sid,
-		VisitorAddr: session.VisitorAddr.String(),
-		ClientAddr:  session.ClientAddr.String(),
+		Sid:         sid,
+		VisitorAddr: visitorAddr,
+		ClientAddr:  clientAddr,
+		Error:       errInfo,
 	}
 	b := bytes.NewBuffer(nil)
 	err := msg.WriteMsg(b, m)