From 653956280933e390d2d82067c5fb9694fca1c493 Mon Sep 17 00:00:00 2001
From: Tim Allclair <tallclair@google.com>
Date: Tue, 6 Sep 2022 16:54:33 -0700
Subject: [PATCH] [konnectivity-client] Close tunnel from conn.Close()

---
 konnectivity-client/pkg/client/client.go | 9 ++++++++-
 konnectivity-client/pkg/client/conn.go   | 7 +++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/konnectivity-client/pkg/client/client.go b/konnectivity-client/pkg/client/client.go
index 145ec8911..63ef1bbf2 100644
--- a/konnectivity-client/pkg/client/client.go
+++ b/konnectivity-client/pkg/client/client.go
@@ -320,7 +320,11 @@ func (t *grpcTunnel) DialContext(requestCtx context.Context, protocol, address s
 
 	klog.V(5).Infoln("DIAL_REQ sent to proxy server")
 
-	c := &conn{stream: t.stream, random: random}
+	c := &conn{
+		stream:      t.stream,
+		random:      random,
+		closeTunnel: t.closeTunnel,
+	}
 
 	select {
 	case res := <-resCh:
@@ -366,7 +370,10 @@ func (t *grpcTunnel) closeDial(dialID int64) {
 	if err := t.stream.Send(req); err != nil {
 		klog.V(5).InfoS("Failed to send DIAL_CLS", "err", err, "dialID", dialID)
 	}
+	t.closeTunnel()
+}
 
+func (t *grpcTunnel) closeTunnel() {
 	atomic.StoreUint32(&t.closing, 1)
 	t.clientConn.Close()
 }
diff --git a/konnectivity-client/pkg/client/conn.go b/konnectivity-client/pkg/client/conn.go
index 822831b10..f76b1e37a 100644
--- a/konnectivity-client/pkg/client/conn.go
+++ b/konnectivity-client/pkg/client/conn.go
@@ -41,6 +41,9 @@ type conn struct {
 	readCh  chan []byte
 	closeCh chan string
 	rdata   []byte
+
+	// closeTunnel is an optional callback to close the underlying grpc connection.
+	closeTunnel func()
 }
 
 var _ net.Conn = &conn{}
@@ -116,6 +119,10 @@ func (c *conn) SetWriteDeadline(t time.Time) error {
 // proxy service to notify remote to drop the connection.
 func (c *conn) Close() error {
 	klog.V(4).Infoln("closing connection")
+	if c.closeTunnel != nil {
+		defer c.closeTunnel()
+	}
+
 	var req *client.Packet
 	if c.connID != 0 {
 		req = &client.Packet{