From 1457b37c7da80cfde1232f95d4bd3013c91eb7a7 Mon Sep 17 00:00:00 2001
From: Yousong Zhou <zhouyousong@yunionyun.com>
Date: Mon, 10 May 2021 13:36:18 +0800
Subject: [PATCH] cloudproxy: agent: ssh: wait ssh conn state while passing it
 out

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

diff --git a/pkg/cloudproxy/agent/ssh/client.go b/pkg/cloudproxy/agent/ssh/client.go
index 75e4560ea0..0e337a80f7 100644
--- a/pkg/cloudproxy/agent/ssh/client.go
+++ b/pkg/cloudproxy/agent/ssh/client.go
@@ -178,7 +178,7 @@ func (c *Client) Start(ctx context.Context) {
 	}
 }
 
-func (c *Client) runClientState(ctx context.Context, sshClientC chan *ssh.Client) {
+func (c *Client) runClientState(ctx context.Context, sshClientC chan<- *ssh.Client) {
 	for {
 		select {
 		case <-ctx.Done():
@@ -203,12 +203,6 @@ func (c *Client) runClientState(ctx context.Context, sshClientC chan *ssh.Client
 		func() {
 			defer sshc.Conn.Close()
 
-			select {
-			case sshClientC <- sshc:
-			case <-ctx.Done():
-				return
-			}
-
 			closeC := make(chan struct{})
 			go func() {
 				defer close(closeC)
@@ -219,6 +213,13 @@ func (c *Client) runClientState(ctx context.Context, sshClientC chan *ssh.Client
 				}
 			}()
 
+			select {
+			case sshClientC <- sshc:
+			case <-closeC:
+			case <-ctx.Done():
+				return
+			}
+
 			select {
 			case <-closeC:
 			case <-ctx.Done():