From adc6496c4c8eb78681d9450b7d86d982c91c87c8 Mon Sep 17 00:00:00 2001 From: Neur0toxine Date: Wed, 29 Nov 2023 19:21:53 +0300 Subject: [PATCH] WIP: config live-reload --- cmd/root.go | 47 +++++++++++++++++++++++++++-------- go.mod | 2 +- internal/api/plugin/server.go | 7 +++++- internal/api/rest/server.go | 19 +++++++++++--- internal/config/model.go | 22 +++++++++++++++- 5 files changed, 81 insertions(+), 16 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index f626974..be1a781 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -5,6 +5,7 @@ import ( "os" "os/signal" "syscall" + "time" "github.com/Neur0toxine/sshpoke/internal/api/plugin" "github.com/Neur0toxine/sshpoke/internal/api/rest" @@ -14,6 +15,7 @@ import ( "github.com/Neur0toxine/sshpoke/internal/server" "github.com/Neur0toxine/sshpoke/pkg/dto" plugin2 "github.com/Neur0toxine/sshpoke/pkg/plugin" + "github.com/fsnotify/fsnotify" "github.com/go-playground/validator/v10" "github.com/spf13/cobra" "github.com/spf13/viper" @@ -30,12 +32,28 @@ var rootCmd = &cobra.Command{ Short: "Expose your Docker services to the Internet via SSH.", Long: `sshpoke is a CLI application that listens to the docker socket and automatically exposes relevant services to the Internet.`, Run: func(cmd *cobra.Command, args []string) { - ctx, cancel := context.WithCancel(context.Background()) - server.DefaultManager = server.NewManager(ctx, config.Default.Servers, config.Default.DefaultServer) - runPluginServer() - runRestServer() - runDockerEventListener(ctx) - shutdown := makeShutdownFunc(cancel) + config.Rehash() + viper.WatchConfig() + cancel, shutdown := initApp() + viper.OnConfigChange(func(in fsnotify.Event) { + if !in.Op.Has(fsnotify.Write) { + return + } + initConfig() + if !config.HasBeenUpdated() { + return + } + config.Rehash() + logger.Default.Info("configuration has been updated, restarting the app...") + cancel() + go server.DefaultManager.WaitForShutdown() + ctx, innerCancel := context.WithTimeout(context.Background(), 2*time.Second) + defer innerCancel() + select { + case <-ctx.Done(): + } + initApp() + }) linuxSig := make(chan os.Signal, 1) signal.Notify(linuxSig) @@ -50,6 +68,15 @@ var rootCmd = &cobra.Command{ }, } +func initApp() (func(), func(os.Signal)) { + ctx, cancel := context.WithCancel(context.Background()) + server.DefaultManager = server.NewManager(ctx, config.Default.Servers, config.Default.DefaultServer) + runPluginServer(ctx) + runRestServer(ctx) + runDockerEventListener(ctx) + return cancel, makeShutdownFunc(cancel) +} + func Execute() { err := rootCmd.Execute() if err != nil { @@ -98,20 +125,20 @@ func initConfig() { logger.Sugar.Debugw("configuration loaded", "config", config.Default) } -func runPluginServer() { +func runPluginServer(ctx context.Context) { port := config.Default.API.Plugin.Port if port == 0 { port = plugin2.DefaultPort } - go plugin.StartServer(port, logger.Sugar.With("component", "pluginServer")) + go plugin.StartServer(ctx, port, logger.Sugar.With("component", "pluginServer")) } -func runRestServer() { +func runRestServer(ctx context.Context) { port := config.Default.API.Rest.Port if port == 0 { port = rest.DefaultPort } - go rest.StartServer( + go rest.StartServer(ctx, config.Default.API.Rest.Token, port, logger.Sugar.With("component", "webServer"), config.Default.Debug) } diff --git a/go.mod b/go.mod index 5fa5d08..45e6d1c 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.21.4 require ( github.com/docker/docker v24.0.7+incompatible github.com/docker/go-connections v0.4.0 + github.com/fsnotify/fsnotify v1.6.0 github.com/function61/gokit v0.0.0-20231117065306-355fe206d542 github.com/gin-contrib/secure v0.0.1 github.com/gin-gonic/gin v1.9.1 @@ -39,7 +40,6 @@ require ( github.com/distribution/reference v0.5.0 // indirect github.com/docker/distribution v2.8.3+incompatible // indirect github.com/docker/go-units v0.5.0 // indirect - github.com/fsnotify/fsnotify v1.6.0 // indirect github.com/gabriel-vasile/mimetype v1.4.2 // indirect github.com/gin-contrib/sse v0.1.0 // indirect github.com/go-openapi/jsonpointer v0.19.5 // indirect diff --git a/internal/api/plugin/server.go b/internal/api/plugin/server.go index 00b6e40..1ff31ae 100644 --- a/internal/api/plugin/server.go +++ b/internal/api/plugin/server.go @@ -59,7 +59,7 @@ func (p *pluginAPI) receiverForContext(ctx context.Context) plugin.Plugin { return server.DefaultManager.PluginByToken(tokens[0]) } -func StartServer(port int, log *zap.SugaredLogger) { +func StartServer(ctx context.Context, port int, log *zap.SugaredLogger) { socket, err := net.Listen("tcp", fmt.Sprintf(":%d", port)) if err != nil { log.Errorf("cannot start plugin API server on port %d: %s", port, err) @@ -68,6 +68,11 @@ func StartServer(port int, log *zap.SugaredLogger) { s := grpc.NewServer() pb.RegisterPluginServiceServer(s, &pluginAPI{log: log}) log.Debugf("starting plugin server on :%d", port) + go func() { + <-ctx.Done() + s.GracefulStop() + socket.Close() + }() if err := s.Serve(socket); err != nil { log.Fatalf("cannot start plugin server on :%d: %s", port, err) } diff --git a/internal/api/rest/server.go b/internal/api/rest/server.go index d8dbb6b..34dc016 100644 --- a/internal/api/rest/server.go +++ b/internal/api/rest/server.go @@ -2,7 +2,10 @@ package rest import ( + "context" + "errors" "net" + "net/http" "strconv" "github.com/Neur0toxine/sshpoke/internal/api/rest/handler" @@ -27,7 +30,7 @@ const ( // @name Authorization // @description Rest API token (leave empty if it's not provided in config). -func StartServer(token string, port int, log *zap.SugaredLogger, debug bool) { +func StartServer(ctx context.Context, token string, port int, log *zap.SugaredLogger, debug bool) { if !debug { gin.SetMode(gin.ReleaseMode) } @@ -45,8 +48,18 @@ func StartServer(token string, port int, log *zap.SugaredLogger, debug bool) { })) g.Use(secure.New(secureConfig())) router(g, token) - err := g.Run(net.JoinHostPort("127.0.0.1", strconv.Itoa(port))) - if err != nil { + srv := &http.Server{ + Addr: net.JoinHostPort("127.0.0.1", strconv.Itoa(port)), + Handler: g, + } + go func() { + <-ctx.Done() + if err := srv.Shutdown(ctx); err != nil { + logger.Sugar.Errorf("failed to stop the server gracefully: %s", err) + } + }() + err := srv.ListenAndServe() + if err != nil && !errors.Is(err, http.ErrServerClosed) { logger.Sugar.Errorf("cannot start Rest API server on port %d: %s", port, err) } } diff --git a/internal/config/model.go b/internal/config/model.go index bfbd2a7..ab8ab22 100644 --- a/internal/config/model.go +++ b/internal/config/model.go @@ -1,6 +1,9 @@ package config import ( + "crypto/sha1" + "encoding/gob" + "encoding/hex" "net/http" "path/filepath" @@ -9,7 +12,10 @@ import ( "github.com/docker/go-connections/tlsconfig" ) -var Default Config +var ( + Default Config + configHash string +) type Config struct { Debug bool `mapstructure:"debug" json:"debug"` @@ -20,6 +26,20 @@ type Config struct { Servers []Server `mapstructure:"servers" json:"servers"` } +func HasBeenUpdated() bool { + return generateConfigHash() != configHash +} + +func Rehash() { + configHash = generateConfigHash() +} + +func generateConfigHash() string { + h := sha1.New() + _ = gob.NewEncoder(h).Encode(Default) + return hex.EncodeToString(h.Sum(nil)) +} + type Service struct { Name smarttypes.MatchableString `mapstructure:"name" json:"name"` Params ServiceLabels `mapstructure:"params" json:"params"`