From e70d1df1861517e0858c7da4c3deb8f42eb3f21f Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Mon, 13 Feb 2023 14:26:37 +0000
Subject: [PATCH] fix: Clean up conn on NewStream error in memDRPC

---
 provisionersdk/transport.go | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/provisionersdk/transport.go b/provisionersdk/transport.go
index b444057d7ea6c..6c32257ed0fe7 100644
--- a/provisionersdk/transport.go
+++ b/provisionersdk/transport.go
@@ -110,15 +110,18 @@ func (m *memDRPC) NewStream(ctx context.Context, rpc string, enc drpc.Encoding)
 	}
 	dConn := drpcconn.New(conn)
 	stream, err := dConn.NewStream(ctx, rpc, enc)
-	if err == nil {
-		go func() {
-			select {
-			case <-stream.Context().Done():
-			case <-m.closed:
-			}
-			_ = dConn.Close()
-			_ = conn.Close()
-		}()
+	if err != nil {
+		_ = dConn.Close()
+		_ = conn.Close()
+		return nil, err
 	}
-	return stream, err
+	go func() {
+		select {
+		case <-stream.Context().Done():
+		case <-m.closed:
+		}
+		_ = dConn.Close()
+		_ = conn.Close()
+	}()
+	return stream, nil
 }