Fix stop not cleaning up WireGuard interface (#286)

client/cmd/root.go

@@ -35,10 +35,6 @@ var (
 		Short: "",
 		Long:  "",
 	}
-
-	// Execution control channel for stopCh signal
-	stopCh    chan int
-	cleanupCh chan struct{}
 )
 
 // Execute executes the root command.
@@ -47,9 +43,6 @@ func Execute() error {
 }
 
 func init() {
-	stopCh = make(chan int)
-	cleanupCh = make(chan struct{})
-
 	defaultConfigPath = "/etc/wiretrustee/config.json"
 	defaultLogFile = "/var/log/wiretrustee/client.log"
 	if runtime.GOOS == "windows" {
@@ -79,14 +72,18 @@ func init() {
 }
 
 // SetupCloseHandler handles SIGTERM signal and exits with success
-func SetupCloseHandler() {
-	c := make(chan os.Signal, 1)
-	signal.Notify(c, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+func SetupCloseHandler(ctx context.Context, cancel context.CancelFunc) {
+	termCh := make(chan os.Signal, 1)
+	signal.Notify(termCh, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 	go func() {
-		for range c {
-			log.Info("shutdown signal received")
-			stopCh <- 0
+		done := ctx.Done()
+		select {
+		case <-done:
+		case <-termCh:
 		}
+
+		log.Info("shutdown signal received")
+		cancel()
 	}()
 }
 

client/cmd/service.go

@@ -13,18 +13,13 @@ import (
 
 type program struct {
 	ctx    context.Context
-	cmd    *cobra.Command
-	args   []string
+	cancel context.CancelFunc
 	serv   *grpc.Server
 }
 
-func newProgram(cmd *cobra.Command, args []string) *program {
-	ctx := internal.CtxInitState(cmd.Context())
-	return &program{
-		ctx:  ctx,
-		cmd:  cmd,
-		args: args,
-	}
+func newProgram(ctx context.Context, cancel context.CancelFunc) *program {
+	ctx = internal.CtxInitState(ctx)
+	return &program{ctx: ctx, cancel: cancel}
 }
 
 func newSVCConfig() *service.Config {
@@ -48,4 +43,3 @@ var serviceCmd = &cobra.Command{
 	Use:   "service",
 	Short: "manages wiretrustee service",
 }
-

client/cmd/service_controller.go

@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"context"
+	"fmt"
 	"net"
 	"os"
 	"strings"
@@ -19,41 +21,40 @@ import (
 func (p *program) Start(svc service.Service) error {
 	// Start should not block. Do the actual work async.
 	log.Info("starting service") //nolint
-	go func() {
-		// in any case, even if configuration does not exists we run daemon to serve CLI gRPC API.
-		p.serv = grpc.NewServer()
+	// in any case, even if configuration does not exists we run daemon to serve CLI gRPC API.
+	p.serv = grpc.NewServer()
 
-		split := strings.Split(daemonAddr, "://")
-		switch split[0] {
-		case "unix":
-			// cleanup failed close
-			stat, err := os.Stat(split[1])
-			if err == nil && !stat.IsDir() {
-				if err := os.Remove(split[1]); err != nil {
-					log.Debugf("remove socket file: %v", err)
-				}
+	split := strings.Split(daemonAddr, "://")
+	switch split[0] {
+	case "unix":
+		// cleanup failed close
+		stat, err := os.Stat(split[1])
+		if err == nil && !stat.IsDir() {
+			if err := os.Remove(split[1]); err != nil {
+				log.Debugf("remove socket file: %v", err)
 			}
-		case "tcp":
-		default:
-			log.Errorf("unsupported daemon address protocol: %v", split[0])
-			return
 		}
+	case "tcp":
+	default:
+		return fmt.Errorf("unsupported daemon address protocol: %v", split[0])
+	}
 
-		listen, err := net.Listen(split[0], split[1])
-		if err != nil {
-			log.Fatalf("failed to listen daemon interface: %v", err)
-		}
+	listen, err := net.Listen(split[0], split[1])
+	if err != nil {
+		return fmt.Errorf("failed to listen daemon interface: %w", err)
+	}
+	go func() {
 		defer listen.Close()
 
 		if split[0] == "unix" {
-			err = os.Chmod(split[1], 0666)
+			err = os.Chmod(split[1], 0o666)
 			if err != nil {
 				log.Errorf("failed setting daemon permissions: %v", split[1])
 				return
 			}
 		}
 
-		serverInstance := server.New(p.ctx, managementURL, configPath, stopCh, cleanupCh)
+		serverInstance := server.New(p.ctx, managementURL, configPath)
 		if err := serverInstance.Start(); err != nil {
 			log.Fatalf("failed start daemon: %v", err)
 		}
@@ -67,21 +68,14 @@ func (p *program) Start(svc service.Service) error {
 	return nil
 }
 
-func (p *program) Stop(service.Service) error {
-	go func() {
-		stopCh <- 1
-	}()
+func (p *program) Stop(srv service.Service) error {
+	p.cancel()
 
-	// stop CLI daemon service
 	if p.serv != nil {
-		p.serv.GracefulStop()
+		p.serv.Stop()
 	}
 
-	select {
-	case <-cleanupCh:
-	case <-time.After(time.Second * 10):
-		log.Warnf("failed waiting for service cleanup, terminating")
-	}
+	time.Sleep(time.Second * 2)
 	log.Info("stopped Wiretrustee service") //nolint
 	return nil
 }
@@ -98,9 +92,10 @@ var runCmd = &cobra.Command{
 			return
 		}
 
-		SetupCloseHandler()
+		ctx, cancel := context.WithCancel(cmd.Context())
+		SetupCloseHandler(ctx, cancel)
 
-		s, err := newSVC(newProgram(cmd, args), newSVCConfig())
+		s, err := newSVC(newProgram(ctx, cancel), newSVCConfig())
 		if err != nil {
 			cmd.PrintErrln(err)
 			return
@@ -125,7 +120,10 @@ var startCmd = &cobra.Command{
 			log.Errorf("failed initializing log %v", err)
 			return err
 		}
-		s, err := newSVC(newProgram(cmd, args), newSVCConfig())
+
+		ctx, cancel := context.WithCancel(cmd.Context())
+
+		s, err := newSVC(newProgram(ctx, cancel), newSVCConfig())
 		if err != nil {
 			cmd.PrintErrln(err)
 			return err
@@ -150,7 +148,10 @@ var stopCmd = &cobra.Command{
 		if err != nil {
 			log.Errorf("failed initializing log %v", err)
 		}
-		s, err := newSVC(newProgram(cmd, args), newSVCConfig())
+
+		ctx, cancel := context.WithCancel(cmd.Context())
+
+		s, err := newSVC(newProgram(ctx, cancel), newSVCConfig())
 		if err != nil {
 			cmd.PrintErrln(err)
 			return
@@ -174,7 +175,10 @@ var restartCmd = &cobra.Command{
 		if err != nil {
 			log.Errorf("failed initializing log %v", err)
 		}
-		s, err := newSVC(newProgram(cmd, args), newSVCConfig())
+
+		ctx, cancel := context.WithCancel(cmd.Context())
+
+		s, err := newSVC(newProgram(ctx, cancel), newSVCConfig())
 		if err != nil {
 			cmd.PrintErrln(err)
 			return

client/cmd/service_installer.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"context"
 	"runtime"
 
 	"github.com/spf13/cobra"
@@ -28,7 +29,9 @@ var installCmd = &cobra.Command{
 			svcConfig.Dependencies = []string{"After=network.target syslog.target"}
 		}
 
-		s, err := newSVC(newProgram(cmd, args), svcConfig)
+		ctx, cancel := context.WithCancel(cmd.Context())
+
+		s, err := newSVC(newProgram(ctx, cancel), svcConfig)
 		if err != nil {
 			cmd.PrintErrln(err)
 			return err
@@ -50,7 +53,9 @@ var uninstallCmd = &cobra.Command{
 	Run: func(cmd *cobra.Command, args []string) {
 		SetFlagsFromEnvVars()
 
-		s, err := newSVC(newProgram(cmd, args), newSVCConfig())
+		ctx, cancel := context.WithCancel(cmd.Context())
+
+		s, err := newSVC(newProgram(ctx, cancel), newSVCConfig())
 		if err != nil {
 			cmd.PrintErrln(err)
 			return
@@ -64,4 +69,3 @@ var uninstallCmd = &cobra.Command{
 		cmd.Println("Wiretrustee has been uninstalled")
 	},
 }
-

client/cmd/testutil.go

@@ -2,12 +2,13 @@ package cmd
 
 import (
 	"context"
-	"github.com/wiretrustee/wiretrustee/util"
 	"net"
 	"path/filepath"
 	"testing"
 	"time"
 
+	"github.com/wiretrustee/wiretrustee/util"
+
 	clientProto "github.com/wiretrustee/wiretrustee/client/proto"
 	client "github.com/wiretrustee/wiretrustee/client/server"
 	mgmtProto "github.com/wiretrustee/wiretrustee/management/proto"
@@ -85,7 +86,6 @@ func startManagement(t *testing.T, config *mgmt.Config) (*grpc.Server, net.Liste
 
 func startClientDaemon(
 	t *testing.T, ctx context.Context, managementURL, configPath string,
-	stopCh chan int, cleanupCh chan<- struct{},
 ) (*grpc.Server, net.Listener) {
 	lis, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
@@ -93,13 +93,7 @@ func startClientDaemon(
 	}
 	s := grpc.NewServer()
 
-	server := client.New(
-		ctx,
-		managementURL,
-		configPath,
-		stopCh,
-		cleanupCh,
-	)
+	server := client.New(ctx, managementURL, configPath)
 	if err := server.Start(); err != nil {
 		t.Fatal(err)
 	}

client/cmd/up.go

@@ -1,6 +1,8 @@
 package cmd
 
 import (
+	"context"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/wiretrustee/wiretrustee/util"
@@ -38,8 +40,10 @@ var upCmd = &cobra.Command{
 				return err
 			}
 
-			SetupCloseHandler()
-			return internal.RunClient(ctx, config, stopCh, cleanupCh)
+			var cancel context.CancelFunc
+			ctx, cancel = context.WithCancel(ctx)
+			SetupCloseHandler(ctx, cancel)
+			return internal.RunClient(ctx, config)
 		}
 
 		conn, err := DialClientGRPCServer(ctx, daemonAddr)

client/cmd/up_daemon_test.go

@@ -8,9 +8,7 @@ import (
 	"github.com/wiretrustee/wiretrustee/client/internal"
 )
 
-var (
-	cliAddr string
-)
+var cliAddr string
 
 func TestUpDaemon(t *testing.T) {
 	mgmAddr := startTestingServices(t)
@@ -18,13 +16,10 @@ func TestUpDaemon(t *testing.T) {
 	tempDir := t.TempDir()
 	confPath := tempDir + "/config.json"
 
-	stopCh = make(chan int, 1)
-	cleanupCh = make(chan struct{}, 1)
-
 	ctx := internal.CtxInitState(context.Background())
 	state := internal.CtxGetState(ctx)
 
-	_, cliLis := startClientDaemon(t, ctx, "http://"+mgmAddr, confPath, stopCh, cleanupCh)
+	_, cliLis := startClientDaemon(t, ctx, "http://"+mgmAddr, confPath)
 
 	cliAddr = cliLis.Addr().String()