WIP: config live-reload

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())
+		config.Rehash()
-		server.DefaultManager = server.NewManager(ctx, config.Default.Servers, config.Default.DefaultServer)
+		viper.WatchConfig()
-		runPluginServer()
+		cancel, shutdown := initApp()
-		runRestServer()
+		viper.OnConfigChange(func(in fsnotify.Event) {
-		runDockerEventListener(ctx)
+			if !in.Op.Has(fsnotify.Write) {
-		shutdown := makeShutdownFunc(cancel)
+				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)
 }
 

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

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)
 	}

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)))
+	srv := &http.Server{
-	if err != nil {
+		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)
 	}
 }

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"`