hotfix: also cleanup in user daemon if error occurs

Author: wencaiwulue

pkg/daemon/action/connect-fork.go

@@ -79,9 +79,9 @@ func (svr *Server) ConnectFork(req *rpc.ConnectRequest, resp rpc.Daemon_ConnectF
 	}
 
 	config.Image = req.Image
-	err = connect.DoConnect(sshCtx, true)
+	err = connect.DoConnect(sshCtx, true, ctx.Done())
 	if err != nil {
-		log.Errorf("Do connect error: %v", err)
+		log.Errorf("Failed to connect: %v", err)
 		return err
 	}
 
@@ -122,6 +122,7 @@ func (svr *Server) redirectConnectForkToSudoDaemon(req *rpc.ConnectRequest, resp
 	})
 	defer func() {
 		if err != nil {
+			connect.Cleanup()
 			sshCancel()
 		}
 	}()

pkg/daemon/action/connect.go

@@ -79,6 +79,9 @@ func (svr *Server) Connect(req *rpc.ConnectRequest, resp rpc.Daemon_ConnectServe
 	})
 	defer func() {
 		if e != nil {
+			svr.connect.Cleanup()
+			svr.connect = nil
+			svr.t = time.Time{}
 			sshCancel()
 		}
 	}()
@@ -97,12 +100,9 @@ func (svr *Server) Connect(req *rpc.ConnectRequest, resp rpc.Daemon_ConnectServe
 	}
 
 	config.Image = req.Image
-	err = svr.connect.DoConnect(sshCtx, false)
+	err = svr.connect.DoConnect(sshCtx, false, ctx.Done())
 	if err != nil {
 		log.Errorf("Failed to connect: %v", err)
-		svr.connect.Cleanup()
-		svr.connect = nil
-		svr.t = time.Time{}
 		return err
 	}
 	return nil
@@ -139,6 +139,7 @@ func (svr *Server) redirectToSudoDaemon(req *rpc.ConnectRequest, resp rpc.Daemon
 	})
 	defer func() {
 		if e != nil {
+			connect.Cleanup()
 			sshCancel()
 		}
 	}()

pkg/handler/connect.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"sync/atomic"
 
 	"net"
 	"net/url"
@@ -197,8 +198,16 @@ func (c *ConnectOptions) CreateRemoteInboundPod(ctx context.Context, workloads [
 	return
 }
 
-func (c *ConnectOptions) DoConnect(ctx context.Context, isLite bool) (err error) {
+func (c *ConnectOptions) DoConnect(ctx context.Context, isLite bool, stopChan <-chan struct{}) (err error) {
 	c.ctx, c.cancel = context.WithCancel(ctx)
+	var success atomic.Bool
+	go func() {
+		// if stop chan done before current function finished, means client ctrl+c to cancel operation
+		<-stopChan
+		if !success.Load() {
+			c.cancel()
+		}
+	}()
 
 	log.Info("Starting connect")
 	m := dhcp.NewDHCPManager(c.clientset.CoreV1().ConfigMaps(c.Namespace), c.Namespace)
@@ -268,6 +277,7 @@ func (c *ConnectOptions) DoConnect(ctx context.Context, isLite bool) (err error)
 		log.Errorf("Configure DNS failed: %v", err)
 		return
 	}
+	success.Store(true)
 	log.Info("Configured DNS service")
 	return
 }

pkg/util/route.go

@@ -63,7 +63,7 @@ func ListService(ctx context.Context, lister v12.ServiceInterface, addRouteFunc
 		}
 		err = addRouteFunc(ips...)
 		if err != nil {
-			log.Errorf("Failed to add service IP: %s to route table: %v", ips, err)
+			log.Errorf("Failed to add service IP to route table: %v", err)
 		}
 		if serviceList.Continue == "" {
 			return nil
@@ -117,7 +117,7 @@ func ListPod(ctx context.Context, lister v12.PodInterface, addRouteFunc func(ipS
 		}
 		err = addRouteFunc(ips...)
 		if err != nil {
-			log.Errorf("Failed to add Pod IP: %v route table: %v", ips, err)
+			log.Errorf("Failed to add Pod IP to route table: %v", err)
 		}
 		if podList.Continue == "" {
 			return nil