fix: Clean up conn on NewStream error in memDRPC (#6182)

mafredri · web-flow · commit 2157bff13f69 · 2023-02-13T17:27:10.000+02:00
diff --git 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
 }
