service: Close listeners and tear down the netstack on shutdown (#1954)

service/listener.go

@@ -36,11 +36,13 @@ func (l *tcpListener) Accept() (net.Conn, error) {
 
 func (l *tcpListener) Close() error {
 	l.closeOnce.Do(func() {
+		// Close before leaving the map so a concurrent closeListeners that
+		// no longer sees this listener knows Accept is already released.
+		close(l.closed)
+
 		l.s.mu.Lock()
 		delete(l.s.mu.listeners, l.port)
 		l.s.mu.Unlock()
-
-		close(l.closed)
 	})
 	return nil
 }

service/service.go

@@ -32,13 +32,18 @@ import (
 const nicID = 1
 
 type Service struct {
-	eg      *errgroup.Group
 	control *nebula.Control
 	ipstack *stack.Stack
 
+	// done is closed once the forwarding goroutines have exited and the
+	// netstack is torn down, at which point err holds the first fatal error.
+	done chan struct{}
+	err  error
+
 	mu struct {
 		sync.Mutex
 
+		closed    bool
 		listeners map[uint16]*tcpListener
 	}
 }
@@ -65,7 +70,6 @@ func New(control *nebula.Control) (_ *Service, reterr error) {
 	ctx := control.Context()
 	eg, ctx := errgroup.WithContext(ctx)
 	s := Service{
-		eg:      eg,
 		control: control,
 	}
 	s.mu.listeners = map[uint16]*tcpListener{}
@@ -74,6 +78,12 @@ func New(control *nebula.Control) (_ *Service, reterr error) {
 		NetworkProtocols:   []stack.NetworkProtocolFactory{ipv4.NewProtocol, ipv6.NewProtocol},
 		TransportProtocols: []stack.TransportProtocolFactory{tcp.NewProtocol, udp.NewProtocol, icmp.NewProtocol4, icmp.NewProtocol6},
 	})
+	// The teardown goroutine only owns the stack once New succeeds
+	defer func() {
+		if reterr != nil {
+			s.ipstack.Destroy()
+		}
+	}()
 	sackEnabledOpt := tcpip.TCPSACKEnabled(true) // TCP SACK is disabled by default
 	tcpipErr := s.ipstack.SetTransportProtocolOption(tcp.ProtocolNumber, &sackEnabledOpt)
 	if tcpipErr != nil {
@@ -123,6 +133,12 @@ func New(control *nebula.Control) (_ *Service, reterr error) {
 			// this will read exactly one packet
 			n, err := reader.Read(buf)
 			if err != nil {
+				// An error after shutdown began is noise from the device pipes
+				// closing, anything else is fatal and must stop nebula too
+				if ctx.Err() != nil {
+					return nil
+				}
+				s.control.Stop()
 				return err
 			}
 			packetBuf := stack.NewPacketBuffer(stack.PacketBufferOptions{
@@ -130,8 +146,8 @@ func New(control *nebula.Control) (_ *Service, reterr error) {
 			})
 			linkEP.InjectInbound(header.IPv4ProtocolNumber, packetBuf)
 
-			if err := ctx.Err(); err != nil {
-				return err
+			if ctx.Err() != nil {
+				return nil
 			}
 		}
 	})
@@ -139,13 +155,18 @@ func New(control *nebula.Control) (_ *Service, reterr error) {
 		for {
 			packet := linkEP.ReadContext(ctx)
 			if packet == nil {
-				if err := ctx.Err(); err != nil {
-					return err
+				if ctx.Err() != nil {
+					return nil
 				}
 				continue
 			}
 			bufView := packet.ToView()
 			if _, err := bufView.WriteTo(writer); err != nil {
+				// Same shutdown noise handling as the inbound reader
+				if ctx.Err() != nil {
+					return nil
+				}
+				s.control.Stop()
 				return err
 			}
 			bufView.Release()
@@ -158,6 +179,22 @@ func New(control *nebula.Control) (_ *Service, reterr error) {
 		return control.Wait()
 	})
 
+	// Tear down here rather than in Close so it also happens when nebula stops
+	// on its own. The forwarding goroutines must exit before the stack they feed
+	// is destroyed.
+	s.done = make(chan struct{})
+	go func() {
+		s.err = eg.Wait()
+		s.closeListeners()
+		// Destroy aborts the endpoints that exist when it starts, then waits on
+		// every endpoint, so a Dial registering in between would hold Wait until
+		// its connect timed out. Without the NIC new dials fail route lookup.
+		// RemoveNIC also closes linkEP.
+		_ = s.ipstack.RemoveNIC(nicID)
+		s.ipstack.Destroy()
+		close(s.done)
+	}()
+
 	return &s, nil
 }
 
@@ -237,6 +274,9 @@ func (s *Service) Listen(network, address string) (net.Listener, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	if s.mu.closed {
+		return nil, net.ErrClosed
+	}
 	if _, ok := s.mu.listeners[port]; ok {
 		return nil, fmt.Errorf("already listening on port %d", port)
 	}
@@ -245,15 +285,38 @@ func (s *Service) Listen(network, address string) (net.Listener, error) {
 	return l, nil
 }
 
+// Wait blocks until the service has fully stopped, either via Close or a fatal
+// error, and returns the first fatal error if there was one.
 func (s *Service) Wait() error {
-	return s.eg.Wait()
+	<-s.done
+	return s.err
 }
 
+// Close closes all listeners and stops nebula. Use Wait to block until
+// shutdown has completed.
 func (s *Service) Close() error {
+	s.closeListeners()
 	s.control.Stop()
 	return nil
 }
 
+// closeListeners unblocks pending Accept calls and rejects new listeners.
+func (s *Service) closeListeners() {
+	// Every caller closes every listener rather than returning early once
+	// closed is set, so none returns while another is still closing them.
+	s.mu.Lock()
+	s.mu.closed = true
+	listeners := make([]*tcpListener, 0, len(s.mu.listeners))
+	for _, l := range s.mu.listeners {
+		listeners = append(listeners, l)
+	}
+	s.mu.Unlock()
+
+	for _, l := range listeners {
+		_ = l.Close()
+	}
+}
+
 func (s *Service) tcpHandler(r *tcp.ForwarderRequest) {
 	endpointID := r.ID()
 

service/service_test.go

@@ -327,3 +327,162 @@ func TestListenerAcceptAfterClose(t *testing.T) {
 		}
 	}
 }
+
+// Close must release blocked Accept calls and let Wait report a clean shutdown.
+func TestServiceClose(t *testing.T) {
+	ca, _, caKey, _ := cert_test.NewTestCaCert(cert.Version2, cert.Curve_CURVE25519, time.Now(), time.Now().Add(10*time.Minute), nil, nil, []string{})
+	a := newSimpleService(ca, caKey, "a", netip.MustParseAddr("10.0.0.1"), m{
+		"static_host_map": m{},
+		"listen": m{
+			"host": "127.0.0.1",
+			"port": 0,
+		},
+	})
+
+	ln, err := a.Listen("tcp", ":1234")
+	if err != nil {
+		t.Fatal(err)
+	}
+	accepted := make(chan error, 1)
+	go func() {
+		_, err := ln.Accept()
+		accepted <- err
+	}()
+
+	if err := a.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	select {
+	case err := <-accepted:
+		if !errors.Is(err, net.ErrClosed) {
+			t.Fatalf("expected net.ErrClosed from Accept, got %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("Accept still blocked after Close")
+	}
+
+	waited := make(chan error, 1)
+	go func() { waited <- a.Wait() }()
+	select {
+	case err := <-waited:
+		if err != nil {
+			t.Fatalf("expected clean shutdown from Wait, got %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("Wait still blocked after Close")
+	}
+
+	if a.ipstack.HasNIC(nicID) {
+		t.Fatal("netstack not torn down after Wait")
+	}
+	if err := a.Close(); err != nil {
+		t.Fatalf("second close: %v", err)
+	}
+	if _, err := a.Listen("tcp", ":1235"); !errors.Is(err, net.ErrClosed) {
+		t.Fatalf("expected net.ErrClosed from Listen after Close, got %v", err)
+	}
+}
+
+// If nebula stops without Close being called, the service must still release
+// Accept and tear down the netstack.
+func TestServiceNebulaStopped(t *testing.T) {
+	ca, _, caKey, _ := cert_test.NewTestCaCert(cert.Version2, cert.Curve_CURVE25519, time.Now(), time.Now().Add(10*time.Minute), nil, nil, []string{})
+	a := newSimpleService(ca, caKey, "a", netip.MustParseAddr("10.0.0.1"), m{
+		"static_host_map": m{},
+		"listen": m{
+			"host": "127.0.0.1",
+			"port": 0,
+		},
+	})
+
+	ln, err := a.Listen("tcp", ":1234")
+	if err != nil {
+		t.Fatal(err)
+	}
+	accepted := make(chan error, 1)
+	go func() {
+		_, err := ln.Accept()
+		accepted <- err
+	}()
+
+	a.control.Stop()
+
+	waited := make(chan error, 1)
+	go func() { waited <- a.Wait() }()
+	select {
+	case err := <-waited:
+		if err != nil {
+			t.Fatalf("expected clean shutdown from Wait, got %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("Wait still blocked after nebula stopped")
+	}
+
+	select {
+	case err := <-accepted:
+		if !errors.Is(err, net.ErrClosed) {
+			t.Fatalf("expected net.ErrClosed from Accept, got %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("Accept still blocked after nebula stopped")
+	}
+
+	if a.ipstack.HasNIC(nicID) {
+		t.Fatal("netstack not torn down after Wait")
+	}
+}
+
+// Dials racing Close must not hold Wait until their connect times out.
+func TestServiceDialDuringClose(t *testing.T) {
+	ca, _, caKey, _ := cert_test.NewTestCaCert(cert.Version2, cert.Curve_CURVE25519, time.Now(), time.Now().Add(10*time.Minute), nil, nil, []string{})
+	a := newSimpleService(ca, caKey, "a", netip.MustParseAddr("10.0.0.1"), m{
+		"static_host_map": m{},
+		"listen": m{
+			"host": "127.0.0.1",
+			"port": 0,
+		},
+	})
+
+	stop := make(chan struct{})
+	var eg errgroup.Group
+	for range 8 {
+		eg.Go(func() error {
+			for {
+				select {
+				case <-stop:
+					return nil
+				default:
+				}
+				// 10.0.0.2 has no tunnel, so each dial sits in SYN-SENT until aborted
+				if c, err := a.Dial("tcp", "10.0.0.2:80"); err == nil {
+					_ = c.Close()
+				}
+			}
+		})
+	}
+	time.Sleep(50 * time.Millisecond)
+
+	if err := a.Close(); err != nil {
+		t.Fatal(err)
+	}
+	waited := make(chan error, 1)
+	go func() { waited <- a.Wait() }()
+	select {
+	case err := <-waited:
+		if err != nil {
+			t.Fatalf("expected clean shutdown from Wait, got %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("Wait blocked behind a dial that raced Close")
+	}
+
+	close(stop)
+	dialersDone := make(chan error, 1)
+	go func() { dialersDone <- eg.Wait() }()
+	select {
+	case <-dialersDone:
+	case <-time.After(5 * time.Second):
+		t.Fatal("Dial still blocked after Wait")
+	}
+}