From 2aaeeac7f6b56e7c653b92ee86247cef5946f9ce Mon Sep 17 00:00:00 2001 From: Givi Khojanashvili Date: Fri, 25 Mar 2022 16:21:04 +0400 Subject: [PATCH] Fix stop not cleaning up WireGuard interface (#286) --- client/cmd/root.go | 23 ++++----- client/cmd/service.go | 14 ++---- client/cmd/service_controller.go | 80 +++++++++++++++++--------------- client/cmd/service_installer.go | 10 ++-- client/cmd/testutil.go | 12 ++--- client/cmd/up.go | 8 +++- client/cmd/up_daemon_test.go | 9 +--- client/internal/connect.go | 14 +----- client/server/server.go | 13 ++---- client/ui/client_ui.go | 43 ++++++++--------- 10 files changed, 101 insertions(+), 125 deletions(-) diff --git a/client/cmd/root.go b/client/cmd/root.go index b4fdb2ce6..46da3a403 100644 --- a/client/cmd/root.go +++ b/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() }() } diff --git a/client/cmd/service.go b/client/cmd/service.go index f5ef50298..67df76f98 100644 --- a/client/cmd/service.go +++ b/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", } - diff --git a/client/cmd/service_controller.go b/client/cmd/service_controller.go index ee4ecc6f9..99ef1bea9 100644 --- a/client/cmd/service_controller.go +++ b/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 diff --git a/client/cmd/service_installer.go b/client/cmd/service_installer.go index 6dcd47a9a..4930dba48 100644 --- a/client/cmd/service_installer.go +++ b/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") }, } - diff --git a/client/cmd/testutil.go b/client/cmd/testutil.go index 22df9f789..9fc177564 100644 --- a/client/cmd/testutil.go +++ b/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) } diff --git a/client/cmd/up.go b/client/cmd/up.go index b22a0b661..ae870f74e 100644 --- a/client/cmd/up.go +++ b/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) diff --git a/client/cmd/up_daemon_test.go b/client/cmd/up_daemon_test.go index 3ed2482d1..2c803b7c2 100644 --- a/client/cmd/up_daemon_test.go +++ b/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() diff --git a/client/internal/connect.go b/client/internal/connect.go index 6c0a1902f..3d22249ab 100644 --- a/client/internal/connect.go +++ b/client/internal/connect.go @@ -17,9 +17,7 @@ import ( ) // RunClient with main logic. -func RunClient( - ctx context.Context, config *Config, stopCh <-chan int, cleanupCh chan<- struct{}, -) error { +func RunClient(ctx context.Context, config *Config) 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 { - case <-stopCh: - case <-ctx.Done(): - } + <-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 } - diff --git a/client/server/server.go b/client/server/server.go index d6366c5f0..4dc13fd79 100644 --- a/client/server/server.go +++ b/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( - ctx context.Context, managementURL, configPath string, - stopCh chan int, cleanupCh chan<- struct{}, -) *Server { +func New(ctx context.Context, managementURL, configPath string) *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 } diff --git a/client/ui/client_ui.go b/client/ui/client_ui.go index 6d8cff76f..7d7f436f5 100644 --- a/client/ui/client_ui.go +++ b/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") - 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") - + for { s.updateStatus() + time.Sleep(time.Second * 3) + } + }() - ticker := time.NewTicker(time.Second * 3) - defer ticker.Stop() - + go func() { 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)