feat: optimize code

cmd/kubevpn/cmds/dev.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/docker/cli/cli/command"
 	dockercomp "github.com/docker/cli/cli/command/completion"
+	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	cmdutil "k8s.io/kubectl/pkg/cmd/util"
 	"k8s.io/kubectl/pkg/util/completion"
@@ -110,9 +111,11 @@ Startup your kubernetes workloads in local Docker container with same volume、e
 			}
 
 			err = dev.DoDev(cmd.Context(), devOptions, sshConf, cmd.Flags(), f, transferImage)
-			for _, fun := range devOptions.RollbackFuncList {
+			for _, fun := range devOptions.GetRollbackFuncList() {
 				if fun != nil {
-					fun()
+					if err = fun(); err != nil {
+						log.Errorf("roll back failed, error: %s", err.Error())
+					}
 				}
 			}
 			return err

pkg/daemon/action/connect-fork.go

@@ -4,11 +4,11 @@ import (
 	"context"
 	"fmt"
 	"io"
-	"k8s.io/utils/pointer"
 	defaultlog "log"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/pflag"
+	"k8s.io/utils/pointer"
 
 	"github.com/wencaiwulue/kubevpn/pkg/config"
 	"github.com/wencaiwulue/kubevpn/pkg/daemon/rpc"
@@ -61,7 +61,10 @@ func (svr *Server) ConnectFork(req *rpc.ConnectRequest, resp rpc.Daemon_ConnectF
 	})
 
 	sshCtx, sshCancel := context.WithCancel(context.Background())
-	connect.RollbackFuncList = append(connect.RollbackFuncList, sshCancel)
+	connect.AddRolloutFunc(func() error {
+		sshCancel()
+		return nil
+	})
 	var path string
 	path, err = handler.SshJump(sshCtx, sshConf, flags, false)
 	if err != nil {
@@ -118,7 +121,10 @@ func (svr *Server) redirectConnectForkToSudoDaemon(req *rpc.ConnectRequest, resp
 		DefValue: file,
 	})
 	sshCtx, sshCancel := context.WithCancel(context.Background())
-	connect.RollbackFuncList = append(connect.RollbackFuncList, sshCancel)
+	connect.AddRolloutFunc(func() error {
+		sshCancel()
+		return nil
+	})
 	var path string
 	path, err = handler.SshJump(sshCtx, sshConf, flags, true)
 	if err != nil {

pkg/daemon/action/connect.go

@@ -76,7 +76,10 @@ func (svr *Server) Connect(req *rpc.ConnectRequest, resp rpc.Daemon_ConnectServe
 	})
 
 	sshCtx, sshCancel := context.WithCancel(context.Background())
-	svr.connect.RollbackFuncList = append(svr.connect.RollbackFuncList, sshCancel)
+	svr.connect.AddRolloutFunc(func() error {
+		sshCancel()
+		return nil
+	})
 	var path string
 	path, err = handler.SshJump(sshCtx, sshConf, flags, false)
 	if err != nil {
@@ -100,6 +103,8 @@ func (svr *Server) Connect(req *rpc.ConnectRequest, resp rpc.Daemon_ConnectServe
 	if err != nil {
 		log.Errorf("do connect error: %v", err)
 		svr.connect.Cleanup()
+		svr.connect = nil
+		svr.t = time.Time{}
 		return err
 	}
 	return nil
@@ -131,7 +136,10 @@ func (svr *Server) redirectToSudoDaemon(req *rpc.ConnectRequest, resp rpc.Daemon
 		DefValue: file,
 	})
 	sshCtx, sshCancel := context.WithCancel(context.Background())
-	connect.RollbackFuncList = append(connect.RollbackFuncList, sshCancel)
+	connect.AddRolloutFunc(func() error {
+		sshCancel()
+		return nil
+	})
 	var path string
 	path, err = handler.SshJump(sshCtx, sshConf, flags, true)
 	if err != nil {

pkg/dev/convert.go

@@ -223,8 +223,8 @@ func GetVolume(ctx context.Context, f util.Factory, ns, pod string, d *Options)
 			if volumeMount.SubPath != "" {
 				join = filepath.Join(join, volumeMount.SubPath)
 			}
-			d.RollbackFuncList = append(d.RollbackFuncList, func() {
+			d.AddRollbackFunc(func() error {
-				_ = os.RemoveAll(join)
+				return os.RemoveAll(join)
 			})
 			// pod-namespace/pod-name:path
 			remotePath := fmt.Sprintf("%s/%s:%s", ns, pod, volumeMount.MountPath)

pkg/dev/main.go

@@ -81,7 +81,7 @@ type Options struct {
 	DockerCli *command.DockerCli
 
 	// rollback
-	RollbackFuncList []func()
+	rollbackFuncList []func() error
 }
 
 func (d *Options) Main(ctx context.Context, tempContainerConfig *containerConfig) error {
@@ -217,8 +217,8 @@ func (d *Options) Main(ctx context.Context, tempContainerConfig *containerConfig
 		}
 	}
 
-	d.RollbackFuncList = append(d.RollbackFuncList, func() {
+	d.AddRollbackFunc(func() error {
-		_ = runConfigList.Remove(ctx, d.Cli)
+		return runConfigList.Remove(ctx, d.Cli)
 	})
 	err = runConfigList.Run(ctx, volume, d.Cli, d.DockerCli)
 	if err != nil {
@@ -576,8 +576,9 @@ func (d *Options) doConnect(ctx context.Context, f cmdutil.Factory, conf *util.S
 			},
 		)
 		go h.Run(func() error { select {} })
-		d.RollbackFuncList = append(d.RollbackFuncList, func() {
+		d.AddRollbackFunc(func() error {
 			h.Close()
+			return nil
 		})
 		err = runLogsWaitRunning(cancelCtx, d.DockerCli, id)
 		if err != nil {
@@ -864,3 +865,11 @@ func createKubevpnNetwork(ctx context.Context, cli *client.Client) (string, erro
 	}
 	return create.ID, nil
 }
+
+func (d *Options) AddRollbackFunc(f func() error) {
+	d.rollbackFuncList = append(d.rollbackFuncList, f)
+}
+
+func (d *Options) GetRollbackFuncList() []func() error {
+	return d.rollbackFuncList
+}

pkg/dns/dns.go

@@ -29,6 +29,8 @@ type Config struct {
 	Ns          []string
 	UseLocalDNS bool
 	TunName     string
+	// lite mode means connect to another cluster
+	Lite bool
 
 	Hosts []Entry
 }

pkg/handler/cleaner.go

@@ -64,9 +64,11 @@ func (c *ConnectOptions) Cleanup() {
 			log.Errorf("can not update ref-count: %v", err)
 		}
 	}
-	for _, function := range c.RollbackFuncList {
+	for _, function := range c.getRolloutFunc() {
 		if function != nil {
-			function()
+			if err := function(); err != nil {
+				log.Warningf("rollout function error: %v", err)
+			}
 		}
 	}
 	// leave proxy resources
@@ -77,7 +79,6 @@ func (c *ConnectOptions) Cleanup() {
 	if c.cancel != nil {
 		c.cancel()
 	}
-	c.RollbackFuncList = c.RollbackFuncList[:]
 	if c.dnsConfig != nil {
 		log.Infof("clean up dns")
 		c.dnsConfig.CancelDNS()

pkg/handler/clone.go

@@ -68,7 +68,7 @@ type CloneOptions struct {
 	config     *rest.Config
 	factory    cmdutil.Factory
 
-	RollbackFuncList []func()
+	rollbackFuncList []func() error
 }
 
 func (d *CloneOptions) InitClient(f cmdutil.Factory) (err error) {
@@ -189,8 +189,8 @@ func (d *CloneOptions) DoClone(ctx context.Context) error {
 		if err != nil {
 			return err
 		}
-		d.RollbackFuncList = append(d.RollbackFuncList, func() {
+		d.addRollbackFunc(func() error {
-			_ = client.Resource(object.Mapping.Resource).Namespace(d.TargetNamespace).Delete(context.Background(), u.GetName(), metav1.DeleteOptions{})
+			return client.Resource(object.Mapping.Resource).Namespace(d.TargetNamespace).Delete(context.Background(), u.GetName(), metav1.DeleteOptions{})
 		})
 		retryErr := retry.RetryOnConflict(retry.DefaultRetry, func() error {
 			// (1) add annotation KUBECONFIG
@@ -785,5 +785,16 @@ func (d *CloneOptions) Cleanup(workloads ...string) error {
 		}
 		log.Infof("clean up clone workload: %s successfully", workload)
 	}
+	for _, f := range d.rollbackFuncList {
+		if f != nil {
+			if err := f(); err != nil {
+				log.Warningf("exec rollback function error: %s", err.Error())
+			}
+		}
+	}
 	return nil
 }
+
+func (d *CloneOptions) addRollbackFunc(f func() error) {
+	d.rollbackFuncList = append(d.rollbackFuncList, f)
+}

pkg/handler/connect.go

@@ -92,7 +92,7 @@ type ConnectOptions struct {
 	// needs to give it back to dhcp
 	localTunIPv4     *net.IPNet
 	localTunIPv6     *net.IPNet
-	RollbackFuncList []func()
+	rollbackFuncList []func() error
 	dnsConfig        *dns.Config
 
 	apiServerIPs []net.IP
@@ -266,7 +266,7 @@ func (c *ConnectOptions) DoConnect(ctx context.Context, isLite bool) (err error)
 		log.Errorf("add extra route failed: %v", err)
 		return
 	}
-	if err = c.setupDNS(c.ctx); err != nil {
+	if err = c.setupDNS(c.ctx, isLite); err != nil {
 		log.Errorf("set up dns failed: %v", err)
 		return
 	}
@@ -614,11 +614,14 @@ func (c *ConnectOptions) deleteFirewallRule(ctx context.Context) {
 	if !util.FindAllowFirewallRule() {
 		util.AddAllowFirewallRule()
 	}
-	c.RollbackFuncList = append(c.RollbackFuncList, util.DeleteAllowFirewallRule)
+	c.AddRolloutFunc(func() error {
+		util.DeleteAllowFirewallRule()
+		return nil
+	})
 	go util.DeleteBlockFirewallRule(ctx)
 }
 
-func (c *ConnectOptions) setupDNS(ctx context.Context) error {
+func (c *ConnectOptions) setupDNS(ctx context.Context, lite bool) error {
 	const port = 53
 	pod, err := c.GetRunningPodList(ctx)
 	if err != nil {
@@ -655,6 +658,7 @@ func (c *ConnectOptions) setupDNS(ctx context.Context) error {
 		Ns:          ns.UnsortedList(),
 		UseLocalDNS: c.UseLocalDNS,
 		TunName:     tunName,
+		Lite:        lite,
 		Hosts:       c.extraHost,
 	}
 	if err = c.dnsConfig.SetupDNS(); err != nil {
@@ -1230,13 +1234,17 @@ RetryWithDNSClient:
 					for _, rr := range answer.Answer {
 						switch a := rr.(type) {
 						case *miekgdns.A:
-							addRouteFunc(domain, a.A.String())
+							if ip := net.ParseIP(a.A.String()); ip != nil && !ip.IsLoopback() {
-							c.extraHost = append(c.extraHost, dns.Entry{IP: a.A.String(), Domain: domain})
+								addRouteFunc(domain, a.A.String())
-							success = true
+								c.extraHost = append(c.extraHost, dns.Entry{IP: a.A.String(), Domain: domain})
+								success = true
+							}
 						case *miekgdns.AAAA:
-							addRouteFunc(domain, a.AAAA.String())
+							if ip := net.ParseIP(a.AAAA.String()); ip != nil && !ip.IsLoopback() {
-							c.extraHost = append(c.extraHost, dns.Entry{IP: a.AAAA.String(), Domain: domain})
+								addRouteFunc(domain, a.AAAA.String())
-							success = true
+								c.extraHost = append(c.extraHost, dns.Entry{IP: a.AAAA.String(), Domain: domain})
+								success = true
+							}
 						}
 					}
 					return nil
@@ -1517,3 +1525,11 @@ func (c *ConnectOptions) GetKubeconfigCluster() string {
 	}
 	return ""
 }
+
+func (c *ConnectOptions) AddRolloutFunc(f func() error) {
+	c.rollbackFuncList = append(c.rollbackFuncList, f)
+}
+
+func (c *ConnectOptions) getRolloutFunc() []func() error {
+	return c.rollbackFuncList
+}

pkg/handler/envoy.go

@@ -67,7 +67,7 @@ func InjectVPNAndEnvoySidecar(ctx1 context.Context, factory cmdutil.Factory, cli
 	}
 	if containerNames.HasAll(config.ContainerSidecarVPN, config.ContainerSidecarEnvoyProxy) {
 		// add rollback func to remove envoy config
-		//RollbackFuncList = append(RollbackFuncList, func() {
+		//rollbackFuncList = append(rollbackFuncList, func() {
 		//	err := UnPatchContainer(factory, clientset, namespace, workload, c.LocalTunIPv4)
 		//	if err != nil {
 		//		log.Error(err)

pkg/handler/remote.go

@@ -527,7 +527,7 @@ func InjectVPNSidecar(ctx1 context.Context, factory cmdutil.Factory, namespace,
 			return err
 		}
 
-		//RollbackFuncList = append(RollbackFuncList, func() {
+		//rollbackFuncList = append(rollbackFuncList, func() {
 		//	p2 := &v1.Pod{ObjectMeta: origin.ObjectMeta, Spec: origin.Spec}
 		//	CleanupUselessInfo(p2)
 		//	if err = CreateAfterDeletePod(factory, p2, helper); err != nil {
@@ -560,7 +560,7 @@ func InjectVPNSidecar(ctx1 context.Context, factory cmdutil.Factory, namespace,
 			return err
 		}
 
-		//RollbackFuncList = append(RollbackFuncList, func() {
+		//rollbackFuncList = append(rollbackFuncList, func() {
 		//	if err = removeInboundContainer(factory, namespace, workload); err != nil {
 		//		log.Error(err)
 		//	}