From d7da74cb9ee92ec09c110d93b0df9233b8f24d37 Mon Sep 17 00:00:00 2001
From: Sudarsan Reddy <sreddy@cloudflare.com>
Date: Fri, 24 Sep 2021 12:56:31 +0100
Subject: [PATCH] TUN-5142: defer close rpcconn inside unregister instead of
 ServeControlStream

---
 connection/control.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/connection/control.go b/connection/control.go
index 17664a94..f586f842 100644
--- a/connection/control.go
+++ b/connection/control.go
@@ -64,9 +64,9 @@ func (c *controlStream) ServeControlStream(
 	shouldWaitForUnregister bool,
 ) error {
 	rpcClient := c.newRPCClientFunc(ctx, rw, c.observer.log)
-	defer rpcClient.Close()
 
 	if err := rpcClient.RegisterConnection(ctx, c.namedTunnelConfig, connOptions, c.connIndex, c.observer); err != nil {
+		rpcClient.Close()
 		return err
 	}
 	c.connectedFuse.Connected()
@@ -82,6 +82,7 @@ func (c *controlStream) ServeControlStream(
 
 func (c *controlStream) waitForUnregister(ctx context.Context, rpcClient NamedTunnelRPCClient) {
 	// wait for connection termination or start of graceful shutdown
+	defer rpcClient.Close()
 	select {
 	case <-ctx.Done():
 		break