fix: protect shared rules across delivery and recovery (#560)

go-gost/program.go

@@ -3,6 +3,14 @@ package main
 import (
 	"context"
 	"errors"
+	"net"
+	"net/http"
+	"os"
+	"os/signal"
+	"strings"
+	"syscall"
+	"time"
+
 	"github.com/go-gost/core/auth"
 	"github.com/go-gost/core/logger"
 	"github.com/go-gost/core/service"
@@ -18,20 +26,22 @@ import (
 	xservice "github.com/go-gost/x/service"
 	"github.com/go-gost/x/socket"
 	"github.com/judwhite/go-svc"
-	"net/http"
-	"os"
-	"os/signal"
-	"strings"
-	"syscall"
-	"time"
 )
 
-type program struct {
-	srvApi       service.Service
-	srvMetrics   service.Service
-	srvProfiling *http.Server
+type reporter interface {
+	Stop()
+}
 
-	cancel context.CancelFunc
+type program struct {
+	startReporter     func() reporter
+	reporter          reporter
+	srvApi            service.Service
+	srvMetrics        service.Service
+	srvProfiling      *http.Server
+	profilingListener net.Listener
+
+	cancel  context.CancelFunc
+	stopped bool
 }
 
 func (p *program) Init(env svc.Environment) error {
@@ -48,7 +58,15 @@ func (p *program) Init(env svc.Environment) error {
 	return nil
 }
 
-func (p *program) Start() error {
+func (p *program) Start() (err error) {
+	unlock := config.LockMutation()
+	defer unlock()
+	p.stopped = false
+	defer func() {
+		if err != nil {
+			p.stopRuntime()
+		}
+	}()
 	cfg, err := parser.Parse()
 	if err != nil {
 		return err
@@ -61,23 +79,28 @@ func (p *program) Start() error {
 		os.Exit(0)
 	}
 
-	config.Set(cfg)
-
 	if err := loader.Load(cfg); err != nil {
 		return err
 	}
 
-	// Enable config persistence after initial load so runtime mutations
-	// (AddService, UpdateService, DeleteService, etc.) are saved to disk.
-	socket.EnableConfigPersist()
-
 	if err := p.run(cfg); err != nil {
 		return err
 	}
 
+	config.Set(cfg)
+	socket.EnableConfigPersist()
+
 	ctx, cancel := context.WithCancel(context.Background())
 	p.cancel = cancel
-	go p.reload(ctx)
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGHUP)
+	go p.reload(ctx, c)
+
+	// A connected panel may immediately send commands. Only expose the agent
+	// after initial config loading, runtime startup and persistence are ready.
+	if p.startReporter != nil {
+		p.reporter = p.startReporter()
+	}
 
 	go func() {
 		select {
@@ -91,7 +114,14 @@ func (p *program) Start() error {
 	return nil
 }
 
-func (p *program) run(cfg *config.Config) error {
+func (p *program) run(cfg *config.Config) (err error) {
+	defer func() {
+		if err != nil {
+			// Auxiliary listeners may occupy ports required by the rollback config.
+			// Release all resources opened by this attempt before rebuilding it.
+			p.stopRuntime()
+		}
+	}()
 	for _, svc := range registry.ServiceRegistry().GetAll() {
 		svc := svc
 		go func() {
@@ -152,6 +182,10 @@ func (p *program) run(cfg *config.Config) error {
 
 	if p.srvProfiling != nil {
 		p.srvProfiling.Close()
+		if p.profilingListener != nil {
+			p.profilingListener.Close()
+			p.profilingListener = nil
+		}
 		p.srvProfiling = nil
 	}
 	if cfg.Profiling != nil {
@@ -162,7 +196,12 @@ func (p *program) run(cfg *config.Config) error {
 		s := &http.Server{
 			Addr: addr,
 		}
+		ln, err := net.Listen("tcp", addr)
+		if err != nil {
+			return err
+		}
 		p.srvProfiling = s
+		p.profilingListener = ln
 
 		go func() {
 			defer s.Close()
@@ -170,7 +209,7 @@ func (p *program) run(cfg *config.Config) error {
 			log := logger.Default().WithFields(map[string]any{"kind": "service", "service": "@profiling"})
 
 			log.Info("listening on ", addr)
-			if err := s.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
+			if err := s.Serve(ln); !errors.Is(err, http.ErrServerClosed) {
 				log.Error(err)
 			}
 		}()
@@ -184,30 +223,45 @@ func (p *program) Stop() error {
 		p.cancel()
 	}
 
-	for name, srv := range registry.ServiceRegistry().GetAll() {
-		srv.Close()
+	if p.reporter != nil {
+		p.reporter.Stop()
+	}
+	unlock := config.LockMutation()
+	defer unlock()
+	p.stopped = true
+	p.stopRuntime()
+	return nil
+}
+
+func (p *program) stopRuntime() {
+	for name := range registry.ServiceRegistry().GetAll() {
+		registry.ServiceRegistry().Unregister(name)
 		logger.Default().Debugf("service %s shutdown", name)
 	}
 
 	if p.srvApi != nil {
 		p.srvApi.Close()
+		p.srvApi = nil
 		logger.Default().Debug("service @api shutdown")
 	}
 	if p.srvMetrics != nil {
 		p.srvMetrics.Close()
+		p.srvMetrics = nil
 		logger.Default().Debug("service @metrics shutdown")
 	}
 	if p.srvProfiling != nil {
 		p.srvProfiling.Close()
+		if p.profilingListener != nil {
+			p.profilingListener.Close()
+			p.profilingListener = nil
+		}
+		p.srvProfiling = nil
 		logger.Default().Debug("service @profiling shutdown")
 	}
-
-	return nil
 }
 
-func (p *program) reload(ctx context.Context) {
-	c := make(chan os.Signal, 1)
-	signal.Notify(c, syscall.SIGHUP)
+func (p *program) reload(ctx context.Context, c chan os.Signal) {
+	defer signal.Stop(c)
 
 	for {
 		select {
@@ -225,13 +279,16 @@ func (p *program) reload(ctx context.Context) {
 }
 
 func (p *program) reloadConfig() error {
+	unlock := config.LockMutation()
+	defer unlock()
+	if p.stopped {
+		return errors.New("agent is shutting down")
+	}
 	cfg, err := parser.Parse()
 	if err != nil {
 		return err
 	}
-	config.Set(cfg)
-
-	if err := loader.Load(cfg); err != nil {
+	if err := loader.Reload(cfg, p.run); err != nil {
 		return err
 	}
 	activeServices := make(map[string]struct{}, len(cfg.Services))
@@ -242,10 +299,6 @@ func (p *program) reloadConfig() error {
 	}
 	xservice.GetGlobalTrafficManager().RetainServices(activeServices)
 
-	if err := p.run(cfg); err != nil {
-		return err
-	}
-
 	return nil
 }