From 52bbfc089cc1aba2183476be329133c85ecac71a Mon Sep 17 00:00:00 2001
From: Guy Lewin <guy.lewin@microsoft.com>
Date: Tue, 25 Feb 2020 13:01:43 -0500
Subject: [PATCH] feat: pass error over response of ping and StartWorkConn

---
 client/control.go     | 10 ++++++++++
 server/control.go     |  7 ++++++-
 server/proxy/proxy.go |  1 +
 server/service.go     |  3 +++
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/client/control.go b/client/control.go
index 170d00ff..c99bb1bc 100644
--- a/client/control.go
+++ b/client/control.go
@@ -159,6 +159,11 @@ func (ctl *Control) HandleReqWorkConn(inMsg *msg.ReqWorkConn) {
 		workConn.Close()
 		return
 	}
+	if startMsg.Error != "" {
+		xl.Error("StartWorkConn contains error: %s", startMsg.Error)
+		workConn.Close()
+		return
+	}
 
 	// dispatch this work connection to related proxy
 	ctl.pm.HandleWorkConn(startMsg.ProxyName, workConn, &startMsg)
@@ -319,6 +324,11 @@ func (ctl *Control) msgHandler() {
 			case *msg.NewProxyResp:
 				ctl.HandleNewProxyResp(m)
 			case *msg.Pong:
+				if m.Error != "" {
+					xl.Error("Pong contains error: %s", m.Error)
+					ctl.conn.Close()
+					return
+				}
 				ctl.lastPong = time.Now()
 				xl.Debug("receive heartbeat from server")
 			}
diff --git a/server/control.go b/server/control.go
index a6fc0201..b37d22de 100644
--- a/server/control.go
+++ b/server/control.go
@@ -462,11 +462,16 @@ func (ctl *Control) manager() {
 			case *msg.Ping:
 				if err := ctl.authVerifier.VerifyPing(m); err != nil {
 					xl.Warn("received invalid ping: %v", err)
+					ctl.sendCh <- &msg.Pong{
+						Error: "invalid authentication in ping",
+					}
 					return
 				}
 				ctl.lastPing = time.Now()
 				xl.Debug("receive heartbeat")
-				ctl.sendCh <- &msg.Pong{}
+				ctl.sendCh <- &msg.Pong{
+					Error: "",
+				}
 			}
 		}
 	}
diff --git a/server/proxy/proxy.go b/server/proxy/proxy.go
index 84a04e61..d046c9cb 100644
--- a/server/proxy/proxy.go
+++ b/server/proxy/proxy.go
@@ -116,6 +116,7 @@ func (pxy *BaseProxy) GetWorkConnFromPool(src, dst net.Addr) (workConn net.Conn,
 			SrcPort:   uint16(srcPort),
 			DstAddr:   dstAddr,
 			DstPort:   uint16(dstPort),
+			Error:     "",
 		})
 		if err != nil {
 			xl.Warn("failed to send message to work connection from pool: %v, times: %d", err, i)
diff --git a/server/service.go b/server/service.go
index 8f11099f..47e1ba68 100644
--- a/server/service.go
+++ b/server/service.go
@@ -438,6 +438,9 @@ func (svr *Service) RegisterWorkConn(workConn net.Conn, newMsg *msg.NewWorkConn)
 	// Check auth.
 	if err := svr.authVerifier.VerifyNewWorkConn(newMsg); err != nil {
 		xl.Warn("Invalid authentication in NewWorkConn message on run id [%s]", newMsg.RunId)
+		msg.WriteMsg(workConn, &msg.StartWorkConn{
+			Error: "invalid authentication in NewWorkConn",
+		})
 		return
 	}
 	ctl.RegisterWorkConn(workConn)