From 23ec9b4979cc656af456a1cd8a78a638d7ab31ac Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Fri, 28 Aug 2026 10:56:05 +0800
Subject: [PATCH] vnet: serialize VirtualNet route lifecycle (#5512)

* vnet: serialize virtual net route lifecycle

* test: remove fixed virtual net helper parameters
---
 pkg/plugin/visitor/virtual_net.go      |  98 ++++++++----
 pkg/plugin/visitor/virtual_net_test.go | 198 +++++++++++++++++++++++++
 pkg/vnet/controller_test.go            |   7 +
 3 files changed, 273 insertions(+), 30 deletions(-)

diff --git a/pkg/plugin/visitor/virtual_net.go b/pkg/plugin/visitor/virtual_net.go
index 87a67744..f95856c2 100644
--- a/pkg/plugin/visitor/virtual_net.go
+++ b/pkg/plugin/visitor/virtual_net.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"sync"
 	"time"
@@ -33,10 +34,16 @@ func init() {
 	Register(v1.VisitorPluginVirtualNet, NewVirtualNetPlugin)
 }
 
+type clientRouteController interface {
+	RegisterClientRoute(context.Context, string, []net.IPNet, io.ReadWriteCloser)
+	UnregisterClientRoute(string, io.Writer) bool
+}
+
 type VirtualNetPlugin struct {
 	pluginCtx PluginContext
 
-	routes []net.IPNet
+	routeController clientRouteController
+	routes          []net.IPNet
 
 	mu             sync.Mutex
 	controllerConn net.Conn
@@ -60,6 +67,9 @@ func NewVirtualNetPlugin(pluginCtx PluginContext, options v1.VisitorPluginOption
 		pluginCtx: pluginCtx,
 		routes:    make([]net.IPNet, 0),
 	}
+	if pluginCtx.VnetController != nil {
+		p.routeController = pluginCtx.VnetController
+	}
 
 	p.ctx, p.cancel = context.WithCancel(pluginCtx.Ctx)
 
@@ -90,7 +100,7 @@ func (p *VirtualNetPlugin) Name() string {
 
 func (p *VirtualNetPlugin) Start() {
 	xl := xlog.FromContextSafe(p.pluginCtx.Ctx)
-	if p.pluginCtx.VnetController == nil {
+	if p.routeController == nil {
 		return
 	}
 
@@ -116,16 +126,17 @@ func (p *VirtualNetPlugin) run() {
 		select {
 		case <-p.ctx.Done():
 			xl.Infof("VirtualNetPlugin run loop for visitor [%s] stopping (context cancelled before pipe creation).", p.pluginCtx.Name)
-			p.cleanupControllerConn(xl)
+			p.cleanupCurrentControllerConn(xl)
 			return
 		default:
 		}
 
 		controllerConn, pluginConn := net.Pipe()
-
-		p.mu.Lock()
-		p.controllerConn = controllerConn
-		p.mu.Unlock()
+		xl.Infof("attempting to register client route for visitor [%s]", p.pluginCtx.Name)
+		if !p.registerControllerConn(controllerConn, pluginConn) {
+			xl.Infof("VirtualNetPlugin run loop for visitor [%s] stopping (context cancelled before route registration).", p.pluginCtx.Name)
+			return
+		}
 
 		// Wrap with CloseNotifyConn which supports both close notification and error recording
 		var closeErr error
@@ -134,8 +145,6 @@ func (p *VirtualNetPlugin) run() {
 			close(currentCloseSignal) // Signal the run loop on close.
 		})
 
-		xl.Infof("attempting to register client route for visitor [%s]", p.pluginCtx.Name)
-		p.pluginCtx.VnetController.RegisterClientRoute(p.ctx, p.pluginCtx.Name, p.routes, controllerConn)
 		xl.Infof("successfully registered client route for visitor [%s]. Starting connection handler with CloseNotifyConn.", p.pluginCtx.Name)
 
 		// Pass the CloseNotifyConn to the visitor for handling.
@@ -146,7 +155,7 @@ func (p *VirtualNetPlugin) run() {
 		select {
 		case <-p.ctx.Done():
 			xl.Infof("VirtualNetPlugin run loop stopping for visitor [%s] (context cancelled while waiting).", p.pluginCtx.Name)
-			p.cleanupControllerConn(xl)
+			p.cleanupControllerConn(xl, controllerConn)
 			return
 		case <-currentCloseSignal:
 			// Determine reconnect delay based on error with exponential backoff
@@ -171,7 +180,7 @@ func (p *VirtualNetPlugin) run() {
 			}
 
 			// The visitor closed the plugin side. Close the controller side.
-			p.cleanupControllerConn(xl)
+			p.cleanupControllerConn(xl, controllerConn)
 
 			xl.Infof("waiting %v before attempting reconnection for visitor [%s]...", reconnectDelay, p.pluginCtx.Name)
 			select {
@@ -186,6 +195,23 @@ func (p *VirtualNetPlugin) run() {
 	}
 }
 
+// registerControllerConn publishes and registers controllerConn atomically with
+// respect to Close. A canceled plugin cannot register a new route.
+func (p *VirtualNetPlugin) registerControllerConn(controllerConn, pluginConn net.Conn) bool {
+	p.mu.Lock()
+	if p.ctx.Err() != nil || p.routeController == nil {
+		p.mu.Unlock()
+		_ = controllerConn.Close()
+		_ = pluginConn.Close()
+		return false
+	}
+
+	p.controllerConn = controllerConn
+	p.routeController.RegisterClientRoute(p.ctx, p.pluginCtx.Name, p.routes, controllerConn)
+	p.mu.Unlock()
+	return true
+}
+
 // virtualNetReconnectDelay returns a bounded reconnect delay without allowing
 // the exponential shift to overflow for large consecutive error counts.
 func virtualNetReconnectDelay(consecutiveErrors int) time.Duration {
@@ -198,16 +224,37 @@ func virtualNetReconnectDelay(consecutiveErrors int) time.Duration {
 	return virtualNetReconnectBaseDelay * time.Duration(1<<uint(consecutiveErrors-1))
 }
 
-// cleanupControllerConn closes the current controllerConn (if it exists) under lock.
-func (p *VirtualNetPlugin) cleanupControllerConn(xl *xlog.Logger) {
+// cleanupControllerConn unregisters and closes one connection round without
+// affecting a replacement route owned by another connection.
+func (p *VirtualNetPlugin) cleanupControllerConn(xl *xlog.Logger, controllerConn net.Conn) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
-	if p.controllerConn != nil {
-		xl.Debugf("cleaning up controllerConn for visitor [%s]", p.pluginCtx.Name)
-		p.controllerConn.Close()
-		p.controllerConn = nil
+	p.cleanupControllerConnLocked(xl, controllerConn)
+}
+
+func (p *VirtualNetPlugin) cleanupCurrentControllerConn(xl *xlog.Logger) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.cleanupControllerConnLocked(xl, p.controllerConn)
+}
+
+// cleanupControllerConnLocked must be called with p.mu held.
+func (p *VirtualNetPlugin) cleanupControllerConnLocked(xl *xlog.Logger, controllerConn net.Conn) {
+	if controllerConn == nil {
+		p.closeSignal = nil
+		return
+	}
+
+	if p.routeController != nil &&
+		p.routeController.UnregisterClientRoute(p.pluginCtx.Name, controllerConn) {
+		xl.Infof("unregistered client route for visitor [%s]", p.pluginCtx.Name)
+	}
+	xl.Debugf("cleaning up controllerConn for visitor [%s]", p.pluginCtx.Name)
+	_ = controllerConn.Close()
+	if p.controllerConn == controllerConn {
+		p.controllerConn = nil
+		p.closeSignal = nil
 	}
-	p.closeSignal = nil
 }
 
 // Close initiates the plugin shutdown.
@@ -218,18 +265,9 @@ func (p *VirtualNetPlugin) Close() error {
 	// Signal the run loop goroutine to stop.
 	p.cancel()
 
-	// Unregister the route only if it is still owned by this plugin instance.
-	p.mu.Lock()
-	controllerConn := p.controllerConn
-	p.mu.Unlock()
-	if p.pluginCtx.VnetController != nil && controllerConn != nil &&
-		p.pluginCtx.VnetController.UnregisterClientRoute(p.pluginCtx.Name, controllerConn) {
-		xl.Infof("unregistered client route for visitor [%s]", p.pluginCtx.Name)
-	}
-
-	// Explicitly close the controller side of the pipe.
-	// This ensures the pipe is broken even if the run loop is stuck or the visitor hasn't closed its end.
-	p.cleanupControllerConn(xl)
+	// Unregister and close the current connection while holding the same lock
+	// used to check cancellation and register a route in run.
+	p.cleanupCurrentControllerConn(xl)
 	xl.Infof("finished cleaning up connections during close for visitor [%s]", p.pluginCtx.Name)
 
 	return nil
diff --git a/pkg/plugin/visitor/virtual_net_test.go b/pkg/plugin/visitor/virtual_net_test.go
index 74f8b43a..1720dff1 100644
--- a/pkg/plugin/visitor/virtual_net_test.go
+++ b/pkg/plugin/visitor/virtual_net_test.go
@@ -17,12 +17,134 @@
 package visitor
 
 import (
+	"context"
+	"io"
+	"net"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/fatedier/frp/pkg/util/xlog"
 )
 
+const testVirtualNetVisitorName = "vnet-visitor"
+
+type fakeClientRouteController struct {
+	mu sync.Mutex
+
+	routes map[string]io.Writer
+
+	beforeRegister  func()
+	registerCalls   int
+	unregisterCalls int
+}
+
+func newFakeClientRouteController() *fakeClientRouteController {
+	return &fakeClientRouteController{
+		routes: make(map[string]io.Writer),
+	}
+}
+
+func (c *fakeClientRouteController) RegisterClientRoute(
+	_ context.Context,
+	name string,
+	_ []net.IPNet,
+	conn io.ReadWriteCloser,
+) {
+	if c.beforeRegister != nil {
+		c.beforeRegister()
+	}
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.registerCalls++
+	c.routes[name] = conn
+}
+
+func (c *fakeClientRouteController) UnregisterClientRoute(name string, conn io.Writer) bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.unregisterCalls++
+	owner, ok := c.routes[name]
+	if !ok || owner != conn {
+		return false
+	}
+	delete(c.routes, name)
+	return true
+}
+
+func (c *fakeClientRouteController) owner() io.Writer {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.routes[testVirtualNetVisitorName]
+}
+
+func (c *fakeClientRouteController) callCounts() (register, unregister int) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.registerCalls, c.unregisterCalls
+}
+
+type trackedConn struct {
+	net.Conn
+	closed atomic.Bool
+}
+
+func (c *trackedConn) Close() error {
+	c.closed.Store(true)
+	return c.Conn.Close()
+}
+
+func newTrackedPipe(t *testing.T) (*trackedConn, *trackedConn) {
+	t.Helper()
+	left, right := net.Pipe()
+	trackedLeft := &trackedConn{Conn: left}
+	trackedRight := &trackedConn{Conn: right}
+	t.Cleanup(func() {
+		_ = trackedLeft.Close()
+		_ = trackedRight.Close()
+	})
+	return trackedLeft, trackedRight
+}
+
+func newTestVirtualNetPlugin(t *testing.T, controller *fakeClientRouteController) *VirtualNetPlugin {
+	t.Helper()
+	pluginCtx := context.Background()
+	ctx, cancel := context.WithCancel(pluginCtx)
+	p := &VirtualNetPlugin{
+		pluginCtx: PluginContext{
+			Name: testVirtualNetVisitorName,
+			Ctx:  pluginCtx,
+		},
+		routeController: controller,
+		routes: []net.IPNet{{
+			IP:   net.ParseIP("10.1.0.1"),
+			Mask: net.CIDRMask(32, 32),
+		}},
+		ctx:    ctx,
+		cancel: cancel,
+	}
+	t.Cleanup(func() {
+		_ = p.Close()
+	})
+	return p
+}
+
+func waitResult[T any](t *testing.T, ch <-chan T) T {
+	t.Helper()
+	select {
+	case result := <-ch:
+		return result
+	case <-time.After(5 * time.Second):
+		t.Fatal("timed out waiting for concurrent operation")
+		var zero T
+		return zero
+	}
+}
+
 // TestVirtualNetReconnectDelay verifies the documented exponential backoff and
 // ensures large error counts remain capped instead of overflowing to zero.
 func TestVirtualNetReconnectDelay(t *testing.T) {
@@ -45,3 +167,79 @@ func TestVirtualNetReconnectDelay(t *testing.T) {
 		})
 	}
 }
+
+func TestVirtualNetPluginCloseBeforeRegisterDoesNotReplaceNewRoute(t *testing.T) {
+	controller := newFakeClientRouteController()
+	oldPlugin := newTestVirtualNetPlugin(t, controller)
+	newPlugin := newTestVirtualNetPlugin(t, controller)
+	oldControllerConn, oldPluginConn := newTrackedPipe(t)
+	newControllerConn, newPluginConn := newTrackedPipe(t)
+
+	allowOldRegister := make(chan struct{})
+	oldRegisterResult := make(chan bool, 1)
+	go func() {
+		<-allowOldRegister
+		oldRegisterResult <- oldPlugin.registerControllerConn(oldControllerConn, oldPluginConn)
+	}()
+
+	require.NoError(t, oldPlugin.Close())
+	require.True(t, newPlugin.registerControllerConn(newControllerConn, newPluginConn))
+	close(allowOldRegister)
+
+	require.False(t, waitResult(t, oldRegisterResult))
+	require.Same(t, newControllerConn, controller.owner())
+	registerCalls, _ := controller.callCounts()
+	require.Equal(t, 1, registerCalls)
+	require.True(t, oldControllerConn.closed.Load())
+	require.True(t, oldPluginConn.closed.Load())
+}
+
+func TestVirtualNetPluginRegisterBeforeCloseIsCleanedUp(t *testing.T) {
+	controller := newFakeClientRouteController()
+	p := newTestVirtualNetPlugin(t, controller)
+	controllerConn, pluginConn := newTrackedPipe(t)
+	registerEntered := make(chan struct{})
+	var registerEnteredOnce sync.Once
+	controller.beforeRegister = func() {
+		registerEnteredOnce.Do(func() {
+			close(registerEntered)
+		})
+		<-p.ctx.Done()
+	}
+
+	registerResult := make(chan bool, 1)
+	go func() {
+		registerResult <- p.registerControllerConn(controllerConn, pluginConn)
+	}()
+	waitResult(t, registerEntered)
+
+	closeResult := make(chan error, 1)
+	go func() {
+		closeResult <- p.Close()
+	}()
+
+	require.NoError(t, waitResult(t, closeResult))
+	require.True(t, waitResult(t, registerResult))
+	require.Nil(t, controller.owner())
+	registerCalls, unregisterCalls := controller.callCounts()
+	require.Equal(t, 1, registerCalls)
+	require.Equal(t, 1, unregisterCalls)
+	require.True(t, controllerConn.closed.Load())
+}
+
+func TestVirtualNetPluginOldConnectionCleanupKeepsReplacementRoute(t *testing.T) {
+	controller := newFakeClientRouteController()
+	oldPlugin := newTestVirtualNetPlugin(t, controller)
+	newPlugin := newTestVirtualNetPlugin(t, controller)
+	oldControllerConn, oldPluginConn := newTrackedPipe(t)
+	newControllerConn, newPluginConn := newTrackedPipe(t)
+
+	require.True(t, oldPlugin.registerControllerConn(oldControllerConn, oldPluginConn))
+	require.True(t, newPlugin.registerControllerConn(newControllerConn, newPluginConn))
+	require.Same(t, newControllerConn, controller.owner())
+
+	oldPlugin.cleanupControllerConn(xlog.FromContextSafe(oldPlugin.ctx), oldControllerConn)
+
+	require.Same(t, newControllerConn, controller.owner())
+	require.True(t, oldControllerConn.closed.Load())
+}
diff --git a/pkg/vnet/controller_test.go b/pkg/vnet/controller_test.go
index dd3d4225..d3b1c068 100644
--- a/pkg/vnet/controller_test.go
+++ b/pkg/vnet/controller_test.go
@@ -51,6 +51,13 @@ func TestClientRouterDeleteRouteRequiresMatchingConnection(t *testing.T) {
 	require.NoError(err)
 	require.Same(replacementConn, got)
 
+	// The read loop for an old connection can exit after a replacement route
+	// has already been registered. Its deferred cleanup must keep the new owner.
+	controller.clientRouter.removeConnRoute(oldConn)
+	got, err = controller.clientRouter.findConn(net.ParseIP("10.1.0.1"))
+	require.NoError(err)
+	require.Same(replacementConn, got)
+
 	require.True(controller.UnregisterClientRoute("vnet-visitor", replacementConn))
 	_, err = controller.clientRouter.findConn(net.ParseIP("10.1.0.1"))
 	require.Error(err)