fix orphaned subprocess leak on service Re registration and toctou

Author: IbrahimAhmed8

internal/node/service_registry.go

@@ -72,6 +72,13 @@ type ServiceRegistry struct {
 	services map[string]Service
 	dht      dhtProvider
 
+	// names claimed by a Register that has not committed yet, guarded by mu
+	pending map[string]*pendingRegistration
+
+	// bumped by TeardownAll so a Register still running cannot commit into the
+	// map it just emptied, guarded by mu
+	generation uint64
+
 	// reprovideNow asks the node to run a reprovide cycle now, so a service
 	// registered after the loop last ran does not wait a whole interval to be
 	// advertised. Optional; nil outside a running node.
@@ -83,6 +90,12 @@ type ServiceRegistry struct {
 	backendProbeTimeout time.Duration
 }
 
+// claim on a name held by a Register that has not committed yet, set when the
+// name is unregistered before the registration commits
+type pendingRegistration struct {
+	cancelled bool
+}
+
 // NewServiceRegistry constructs a registry bounding backend probes by
 // backendProbeTimeout. A zero or negative value falls back to
 // defaultDHTProbeTimeout, so callers can pass an unset
@@ -94,6 +107,7 @@ func NewServiceRegistry(d dhtProvider, backendProbeTimeout time.Duration) *Servi
 	}
 	return &ServiceRegistry{
 		services:            map[string]Service{},
+		pending:             map[string]*pendingRegistration{},
 		dht:                 d,
 		backendProbeTimeout: backendProbeTimeout,
 	}
@@ -106,6 +120,8 @@ func NewServiceRegistry(d dhtProvider, backendProbeTimeout time.Duration) *Servi
 // A backend that does not answer is registered but not advertised, rather than
 // rejected: backends routinely start after the node does, and the reprovide
 // loop picks them up once they answer.
+//
+// registering a name again replaces the instance and tears the old one down
 func (r *ServiceRegistry) Register(ctx context.Context, svc Service) error {
 	info := svc.Info()
 	if info.Type == api.ServiceType_SERVICE_TYPE_UNSPECIFIED {
@@ -123,26 +139,53 @@ func (r *ServiceRegistry) Register(ctx context.Context, svc Service) error {
 		return fmt.Errorf("invalid service name %q: %w", info.Name, err)
 	}
 
-	// Policy is evaluated on type://name and the registry is keyed on name,
-	// so a second type under the same name would let a grant for one reach
-	// the other. One name, one service.
-	r.mu.RLock()
-	existing, taken := r.services[info.Name]
-	r.mu.RUnlock()
-	if taken && existing.Info().Type != info.Type {
+	// the registry is keyed on name but policy is evaluated on type://name, so
+	// one name maps to one type. reserve under the same lock as the check, or
+	// two concurrent registrations both pass it and disagree about the owner
+	r.mu.Lock()
+	if r.pending == nil {
+		r.pending = map[string]*pendingRegistration{}
+	}
+	if existing, taken := r.services[info.Name]; taken && existing.Info().Type != info.Type {
+		r.mu.Unlock()
 		return fmt.Errorf("service name %q is already registered as %s; a name cannot serve two types", info.Name, existing.Info().Type)
 	}
+	if _, inFlight := r.pending[info.Name]; inFlight {
+		r.mu.Unlock()
+		return fmt.Errorf("service name %q is already being registered", info.Name)
+	}
+	reservation := &pendingRegistration{}
+	r.pending[info.Name] = reservation
+	generation := r.generation
+	r.mu.Unlock()
+
+	// release drops the reservation, teardown is for aborts after Init where
+	// the service has already built something to clean up
+	release := func(teardown bool) {
+		r.mu.Lock()
+		delete(r.pending, info.Name)
+		r.mu.Unlock()
+		if !teardown {
+			return
+		}
+		if err := svc.Teardown(); err != nil {
+			logger.Errorf("[ServiceRegistry] Teardown %s after abandoned registration: %v", info.Name, err)
+		}
+	}
 
 	if err := svc.Init(ctx); err != nil {
+		release(false)
 		return fmt.Errorf("init %s: %w", info.Name, err)
 	}
 
 	srvNameCID, err := serviceNameToCID(info.Type, info.Name)
 	if err != nil {
+		release(true)
 		return err
 	}
 	srvTypeCID, err := serviceTypeToCID(info.Type)
 	if err != nil {
+		release(true)
 		return err
 	}
 
@@ -163,10 +206,33 @@ func (r *ServiceRegistry) Register(ctx context.Context, svc Service) error {
 		}
 	}
 
+	// the registry may have been torn down or the name unregistered while we
+	// were initialising, so don't leave the service in the map
 	r.mu.Lock()
+	if r.generation != generation {
+		r.mu.Unlock()
+		release(true)
+		return fmt.Errorf("registry was torn down while %q was being registered", info.Name)
+	}
+	if reservation.cancelled {
+		r.mu.Unlock()
+		release(true)
+		return fmt.Errorf("service name %q was unregistered while it was being registered", info.Name)
+	}
+	delete(r.pending, info.Name)
+	replaced := r.services[info.Name]
 	r.services[info.Name] = svc
 	r.mu.Unlock()
 
+	// teardown outside the lock, it kills a subprocess
+	if replaced != nil && replaced != svc {
+		if err := replaced.Teardown(); err != nil {
+			logger.Errorf("[ServiceRegistry] Teardown replaced %s: %v", info.Name, err)
+		} else {
+			logger.Infof("[ServiceRegistry] Replaced %s/%s, previous instance torn down", info.Type, info.Name)
+		}
+	}
+
 	if probeErr == nil {
 		logger.Infof("[ServiceRegistry] Registered %s/%s (name CID: %s, type CID: %s)", info.Type, info.Name, srvNameCID, srvTypeCID)
 	} else if r.reprovideNow != nil {
@@ -181,6 +247,11 @@ func (r *ServiceRegistry) Unregister(ctx context.Context, name string) error {
 	r.mu.Lock()
 	svc, ok := r.services[name]
 	delete(r.services, name)
+	// cancel a registration of this name that has not committed, else it lands
+	// after we return and the name comes back
+	if p, inFlight := r.pending[name]; inFlight {
+		p.cancelled = true
+	}
 	r.mu.Unlock()
 	if !ok {
 		return nil
@@ -238,10 +309,14 @@ func (r *ServiceRegistry) insertService(svc Service) {
 
 // TeardownAll calls Teardown on every registered service and clears the
 // map. Per-service errors are logged; iteration continues.
+//
+// the generation bump voids registrations still running, which would otherwise
+// commit into the map this just emptied
 func (r *ServiceRegistry) TeardownAll() {
 	r.mu.Lock()
 	svcs := r.services
 	r.services = map[string]Service{}
+	r.generation++
 	r.mu.Unlock()
 	for name, svc := range svcs {
 		if err := svc.Teardown(); err != nil {