Refactor config reloading

Author: cedws

pkg/cmd/cmd.go

@@ -8,14 +8,14 @@ import (
 
 	"github.com/alecthomas/kong"
 	"github.com/cedws/sshgate/pkg/sshgate"
-	"github.com/fsnotify/fsnotify"
 )
 
 type cli struct {
-	ListenAddr string `help:"Address to listen on" default:":2222"`
-	Ruleless   bool   `help:"Run in ruleless mode"`
-	LogFormat  string `help:"Log format"`
-	Config     string `help:"Path to JSON config file"`
+	ListenAddr     string `help:"Address to listen on" default:":2222"`
+	Ruleless       bool   `help:"Run in ruleless mode"`
+	NoConfigReload bool   `help:"Disable config reload"`
+	LogFormat      string `help:"Log format"`
+	Config         string `help:"Path to JSON config file"`
 
 	Serve      serveCmd      `cmd:"" default:"1" help:"Start the server"`
 	JSONSchema jsonschemaCmd `cmd:"" name:"jsonschema" help:"Print config JSON schema"`
@@ -43,74 +43,27 @@ func (s *serveCmd) Run(cli *cli) error {
 			return err
 		}
 
-		if err := serveUntilReload(context.Background(), cli, config); err != nil && !errors.Is(err, context.Canceled) {
+		if err := serve(context.Background(), cli, config); err != nil && !errors.Is(err, context.Canceled) {
 			return err
 		}
 	}
 }
 
-func serveUntilReload(ctx context.Context, cli *cli, config *sshgate.Config) error {
-	watcher, err := fsnotify.NewWatcher()
-	if err != nil {
-		return err
-	}
-	defer watcher.Close()
-
-	if err := watcher.Add(cli.Config); err != nil {
-		return err
-	}
-
-	for _, hostKeyPath := range []string{
-		config.HostKeyPaths.ECDSA,
-		config.HostKeyPaths.ED25519,
-		config.HostKeyPaths.RSA,
-	} {
-		if hostKeyPath != "" {
-			if err := watcher.Add(hostKeyPath); err != nil {
-				return err
-			}
-		}
-	}
-
-	ctx, cancel := fsnotifyContext(ctx, watcher)
-	defer cancel()
-
-	return serve(ctx, cli, config)
-}
-
 func serve(ctx context.Context, c *cli, config *sshgate.Config) error {
 	var opts []sshgate.Option
+
 	if c.Ruleless {
 		opts = append(opts, sshgate.WithRulelessMode())
 	}
+	if !c.NoConfigReload {
+		opts = append(opts, sshgate.WithConfigReload())
+	}
 
 	server := sshgate.New(config, c.ListenAddr, opts...)
 
 	return server.ListenAndServe(ctx)
 }
 
-func fsnotifyContext(ctx context.Context, watcher *fsnotify.Watcher) (context.Context, context.CancelFunc) {
-	ctx, cancel := context.WithCancel(ctx)
-
-	go func() {
-		for {
-			select {
-			case evt := <-watcher.Events:
-				if evt.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Rename) != 0 {
-					slog.Info("config file changed, reloading")
-					cancel()
-				}
-			case <-watcher.Errors:
-				cancel()
-			case <-ctx.Done():
-				return
-			}
-		}
-	}()
-
-	return ctx, cancel
-}
-
 func Execute() {
 	var cli cli
 

pkg/sshgate/config.go

@@ -43,6 +43,7 @@ type Config struct {
 	Tsnet        Tsnet        `json:"tsnet"`
 	HostKeyPaths HostKeyPaths `json:"host_key_paths"`
 
+	path           string
 	signers        []ssh.Signer
 	parsedPolicies parsedPolicies
 }
@@ -92,6 +93,8 @@ func ReadConfig(path string) (*Config, error) {
 			Hostname: defaultTsnetHostname,
 			Port:     defaultTsnetPort,
 		},
+
+		path: path,
 	}
 	if err := json.Unmarshal(data, &config); err != nil {
 		return nil, fmt.Errorf("failed to parse config JSON: %w", err)

pkg/sshgate/sshgate.go

@@ -17,6 +17,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/fsnotify/fsnotify"
 	"golang.org/x/crypto/ssh"
 	"golang.org/x/sync/errgroup"
 	"tailscale.com/client/local"
@@ -112,7 +113,8 @@ func (d *directTCPIPExtraData) UnmarshalBinary(b []byte) error {
 }
 
 type Options struct {
-	Ruleless bool
+	Ruleless     bool
+	ConfigReload bool
 }
 
 type Option func(*Options)
@@ -123,6 +125,12 @@ func WithRulelessMode() Option {
 	}
 }
 
+func WithConfigReload() Option {
+	return func(o *Options) {
+		o.ConfigReload = true
+	}
+}
+
 type Server struct {
 	config     *Config
 	listenAddr string
@@ -158,6 +166,14 @@ func (s *Server) ListenAndServe(ctx context.Context) error {
 		slog.Warn("running in ruleless mode")
 	}
 
+	if s.options.ConfigReload {
+		var err error
+		ctx, err = notifyConfigReload(ctx, s.config)
+		if err != nil {
+			return err
+		}
+	}
+
 	errgroup, ctx := errgroup.WithContext(ctx)
 
 	errgroup.Go(func() error {
@@ -177,6 +193,59 @@ func (s *Server) ListenAndServe(ctx context.Context) error {
 	return errgroup.Wait()
 }
 
+func notifyConfigReload(ctx context.Context, config *Config) (context.Context, error) {
+	watcher, err := fsnotify.NewWatcher()
+	if err != nil {
+		return nil, err
+	}
+	go func() {
+		<-ctx.Done()
+		watcher.Close()
+	}()
+
+	if err := watcher.Add(config.path); err != nil {
+		return nil, err
+	}
+
+	for _, hostKeyPath := range []string{
+		config.HostKeyPaths.ECDSA,
+		config.HostKeyPaths.ED25519,
+		config.HostKeyPaths.RSA,
+	} {
+		if hostKeyPath != "" {
+			if err := watcher.Add(hostKeyPath); err != nil {
+				return nil, err
+			}
+		}
+	}
+
+	return fsnotifyContext(ctx, watcher), nil
+}
+
+func fsnotifyContext(ctx context.Context, watcher *fsnotify.Watcher) context.Context {
+	ctx, cancel := context.WithCancel(ctx)
+
+	go func() {
+		for {
+			select {
+			case evt := <-watcher.Events:
+				if evt.Op&(fsnotify.Write|fsnotify.Create|fsnotify.Rename) != 0 {
+					slog.Info("config file changed, reloading")
+					cancel()
+					return
+				}
+			case <-watcher.Errors:
+				cancel()
+				return
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	return ctx
+}
+
 func (s *Server) logStats(ctx context.Context) error {
 	for {
 		select {