From cd14aa11188021ebcf599a1f2ac4ecc74d2565e2 Mon Sep 17 00:00:00 2001
From: Yousong Zhou <zhouyousong@yunionyun.com>
Date: Sat, 17 Apr 2021 23:11:53 +0800
Subject: [PATCH] cloudproxy: agent: return when ctx is Done

---
 pkg/cloudproxy/agent/ssh/client.go | 34 +++++++++++++++++++++++-------
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/pkg/cloudproxy/agent/ssh/client.go b/pkg/cloudproxy/agent/ssh/client.go
index c7d785cf71..5f19a11036 100644
--- a/pkg/cloudproxy/agent/ssh/client.go
+++ b/pkg/cloudproxy/agent/ssh/client.go
@@ -180,6 +180,12 @@ func (c *Client) Start(ctx context.Context) {
 
 func (c *Client) runClientState(ctx context.Context, sshClientC chan *ssh.Client) {
 	for {
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+
 		tmoCtx, _ := context.WithTimeout(ctx, 31*time.Second)
 		sshc, err := c.cc.ConnectContext(tmoCtx)
 		if err != nil {
@@ -187,17 +193,29 @@ func (c *Client) runClientState(ctx context.Context, sshClientC chan *ssh.Client
 			continue
 		}
 
-		select {
-		case sshClientC <- sshc:
-		case <-ctx.Done():
-		}
-
 		func() {
 			defer sshc.Conn.Close()
 
-			err := sshc.Conn.Wait()
-			if err != nil {
-				log.Infof("ssh client conn: %v", err)
+			select {
+			case sshClientC <- sshc:
+			case <-ctx.Done():
+				return
+			}
+
+			closeC := make(chan struct{})
+			go func() {
+				defer close(closeC)
+
+				err := sshc.Conn.Wait()
+				if err != nil {
+					log.Infof("ssh client conn: %v", err)
+				}
+			}()
+
+			select {
+			case <-closeC:
+			case <-ctx.Done():
+				return
 			}
 		}()
 	}