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() {
+func SetupCloseHandler(ctx context.Context, cancel context.CancelFunc) {
-	c := make(chan os.Signal, 1)
+	termCh := make(chan os.Signal, 1)
-	signal.Notify(c, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+	signal.Notify(termCh, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 	go func() {
-		for range c {
+		done := ctx.Done()
-			log.Info("shutdown signal received")
+		select {
-			stopCh <- 0
+		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
+	cancel context.CancelFunc
-	args   []string
 	serv   *grpc.Server
 }
 
-func newProgram(cmd *cobra.Command, args []string) *program {
+func newProgram(ctx context.Context, cancel context.CancelFunc) *program {
-	ctx := internal.CtxInitState(cmd.Context())
+	ctx = internal.CtxInitState(ctx)
-	return &program{
+	return &program{ctx: ctx, cancel: cancel}
-		ctx:  ctx,
-		cmd:  cmd,
-		args: args,
-	}
 }
 
 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.
-		// in any case, even if configuration does not exists we run daemon to serve CLI gRPC API.
+	p.serv = grpc.NewServer()
-		p.serv = grpc.NewServer()
 
-		split := strings.Split(daemonAddr, "://")
+	split := strings.Split(daemonAddr, "://")
-		switch split[0] {
+	switch split[0] {
-		case "unix":
+	case "unix":
-			// cleanup failed close
+		// cleanup failed close
-			stat, err := os.Stat(split[1])
+		stat, err := os.Stat(split[1])
-			if err == nil && !stat.IsDir() {
+		if err == nil && !stat.IsDir() {
-				if err := os.Remove(split[1]); err != nil {
+			if err := os.Remove(split[1]); err != nil {
-					log.Debugf("remove socket file: %v", err)
+				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])
+	listen, err := net.Listen(split[0], split[1])
-		if err != nil {
+	if err != nil {
-			log.Fatalf("failed to listen daemon interface: %v", err)
+		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 {
+func (p *program) Stop(srv service.Service) error {
-	go func() {
+	p.cancel()
-		stopCh <- 1
-	}()
 
-	// stop CLI daemon service
 	if p.serv != nil {
-		p.serv.GracefulStop()
+		p.serv.Stop()
 	}
 
-	select {
+	time.Sleep(time.Second * 2)
-	case <-cleanupCh:
-	case <-time.After(time.Second * 10):
-		log.Warnf("failed waiting for service cleanup, terminating")
-	}
 	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(
+	server := client.New(ctx, managementURL, configPath)
-		ctx,
-		managementURL,
-		configPath,
-		stopCh,
-		cleanupCh,
-	)
 	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()
+			var cancel context.CancelFunc
-			return internal.RunClient(ctx, config, stopCh, cleanupCh)
+			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 (
+var cliAddr string
-	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()
 

client/internal/connect.go

@@ -17,9 +17,7 @@ import (
 )
 
 // RunClient with main logic.
-func RunClient(
+func RunClient(ctx context.Context, config *Config) error {
-	ctx context.Context, config *Config, stopCh <-chan int, cleanupCh chan<- struct{},
-) error {
 	backOff := &backoff.ExponentialBackOff{
 		InitialInterval:     time.Second,
 		RandomizationFactor: backoff.DefaultRandomizationFactor,
@@ -90,10 +88,7 @@ func RunClient(
 		log.Print("Wiretrustee engine started, my IP is: ", peerConfig.Address)
 		state.Set(StatusConnected)
 
-		select {
+		<-ctx.Done()
-		case <-stopCh:
-		case <-ctx.Done():
-		}
 
 		backOff.Reset()
 
@@ -114,10 +109,6 @@ func RunClient(
 			return wrapErr(err)
 		}
 
-		go func() {
-			cleanupCh <- struct{}{}
-		}()
-
 		log.Info("stopped Wiretrustee client")
 
 		if _, err := state.Status(); err == ErrResetConnection {
@@ -207,4 +198,3 @@ func connectToManagement(ctx context.Context, managementAddr string, ourPrivateK
 
 	return client, loginResp, nil
 }
-

client/server/server.go

@@ -18,8 +18,6 @@ type Server struct {
 
 	managementURL string
 	configPath    string
-	stopCh        chan int
-	cleanupCh     chan<- struct{}
 
 	mutex  sync.Mutex
 	config *internal.Config
@@ -27,16 +25,11 @@ type Server struct {
 }
 
 // New server instance constructor.
-func New(
+func New(ctx context.Context, managementURL, configPath string) *Server {
-	ctx context.Context, managementURL, configPath string,
-	stopCh chan int, cleanupCh chan<- struct{},
-) *Server {
 	return &Server{
 		rootCtx:       ctx,
 		managementURL: managementURL,
 		configPath:    configPath,
-		stopCh:        stopCh,
-		cleanupCh:     cleanupCh,
 	}
 }
 
@@ -67,7 +60,7 @@ func (s *Server) Start() error {
 	s.config = config
 
 	go func() {
-		if err := internal.RunClient(ctx, config, s.stopCh, s.cleanupCh); err != nil {
+		if err := internal.RunClient(ctx, config); err != nil {
 			log.Errorf("init connections: %v", err)
 		}
 	}()
@@ -131,7 +124,7 @@ func (s *Server) Up(_ context.Context, msg *proto.UpRequest) (*proto.UpResponse,
 	}
 
 	go func() {
-		if err := internal.RunClient(ctx, s.config, s.stopCh, s.cleanupCh); err != nil {
+		if err := internal.RunClient(ctx, s.config); err != nil {
 			log.Errorf("run client connection: %v", state.Wrap(err))
 			return
 		}

client/ui/client_ui.go

@@ -151,28 +151,31 @@ func (s *serviceClient) updateStatus() {
 
 func (s *serviceClient) onTrayReady() {
 	systray.SetTemplateIcon(iconDisconnected, iconDisconnected)
+
+	s.mStatus = systray.AddMenuItem("Disconnected", "Disconnected")
+	s.mStatus.Disable()
+
+	systray.AddSeparator()
+
+	s.mUp = systray.AddMenuItem("Up", "Up")
+
+	s.mDown = systray.AddMenuItem("Down", "Down")
+	s.mDown.Disable()
+
+	mURL := systray.AddMenuItem("Open UI", "wiretrustee website")
+
+	systray.AddSeparator()
+
+	mQuit := systray.AddMenuItem("Quit", "Quit the whole app")
+
 	go func() {
-		s.mStatus = systray.AddMenuItem("Disconnected", "Disconnected")
+		for {
-		s.mStatus.Disable()
-
-		systray.AddSeparator()
-
-		s.mUp = systray.AddMenuItem("Up", "Up")
-
-		s.mDown = systray.AddMenuItem("Down", "Down")
-		s.mDown.Disable()
-
-		mURL := systray.AddMenuItem("Open UI", "wiretrustee website")
-
-		systray.AddSeparator()
-
-		mQuit := systray.AddMenuItem("Quit", "Quit the whole app")
-
 		s.updateStatus()
+			time.Sleep(time.Second * 3)
+		}
+	}()
 
-		ticker := time.NewTicker(time.Second * 3)
+	go func() {
-		defer ticker.Stop()
-
 		var err error
 		for {
 			select {
@@ -191,8 +194,6 @@ func (s *serviceClient) onTrayReady() {
 			case <-mQuit.ClickedCh:
 				systray.Quit()
 				return
-			case <-ticker.C:
-				s.updateStatus()
 			}
 			if err != nil {
 				log.Errorf("process connection: %v", err)