[konnectivity-client] Close tunnel from conn.Close()

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()
 }

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{