vnet: fix issues (#4771)

This commit is contained in:
fatedier
2025-04-27 15:22:28 +08:00
parent 27f66baf54
commit 3c8d648ddc
18 changed files with 271 additions and 185 deletions

View File

@@ -18,9 +18,10 @@ package client
import (
"context"
"io"
"sync"
v1 "github.com/fatedier/frp/pkg/config/v1"
"github.com/fatedier/frp/pkg/util/xlog"
)
func init() {
@@ -30,6 +31,8 @@ func init() {
type VirtualNetPlugin struct {
pluginCtx PluginContext
opts *v1.VirtualNetPluginOptions
mu sync.Mutex
conns map[io.ReadWriteCloser]struct{}
}
func NewVirtualNetPlugin(pluginCtx PluginContext, options v1.ClientPluginOptions) (Plugin, error) {
@@ -43,19 +46,32 @@ func NewVirtualNetPlugin(pluginCtx PluginContext, options v1.ClientPluginOptions
}
func (p *VirtualNetPlugin) Handle(ctx context.Context, connInfo *ConnectionInfo) {
xl := xlog.FromContextSafe(ctx)
// Verify if virtual network controller is available
if p.pluginCtx.VnetController == nil {
return
}
// Register the connection with the controller
routeName := p.pluginCtx.Name
err := p.pluginCtx.VnetController.RegisterServerConn(ctx, routeName, connInfo.Conn)
if err != nil {
xl.Errorf("virtual net failed to register server connection: %v", err)
return
// Add the connection before starting the read loop to avoid race condition
// where RemoveConn might be called before the connection is added.
p.mu.Lock()
if p.conns == nil {
p.conns = make(map[io.ReadWriteCloser]struct{})
}
p.conns[connInfo.Conn] = struct{}{}
p.mu.Unlock()
// Register the connection with the controller and pass the cleanup function
p.pluginCtx.VnetController.StartServerConnReadLoop(ctx, connInfo.Conn, func() {
p.RemoveConn(connInfo.Conn)
})
}
func (p *VirtualNetPlugin) RemoveConn(conn io.ReadWriteCloser) {
p.mu.Lock()
defer p.mu.Unlock()
// Check if the map exists, as Close might have set it to nil concurrently
if p.conns != nil {
delete(p.conns, conn)
}
}
@@ -64,8 +80,13 @@ func (p *VirtualNetPlugin) Name() string {
}
func (p *VirtualNetPlugin) Close() error {
if p.pluginCtx.VnetController != nil {
p.pluginCtx.VnetController.UnregisterServerConn(p.pluginCtx.Name)
p.mu.Lock()
defer p.mu.Unlock()
// Close any remaining connections
for conn := range p.conns {
_ = conn.Close()
}
p.conns = nil
return nil
}

View File

@@ -60,7 +60,7 @@ func NewVirtualNetPlugin(pluginCtx PluginContext, options v1.VisitorPluginOption
return nil, errors.New("destinationIP is required")
}
// Parse DestinationIP as a single IP and create a host route
// Parse DestinationIP and create a host route.
ip := net.ParseIP(opts.DestinationIP)
if ip == nil {
return nil, fmt.Errorf("invalid destination IP address [%s]", opts.DestinationIP)
@@ -91,7 +91,7 @@ func (p *VirtualNetPlugin) Start() {
if len(p.routes) > 0 {
routeStr = p.routes[0].String()
}
xl.Infof("Starting VirtualNetPlugin for visitor [%s], attempting to register routes for %s", p.pluginCtx.Name, routeStr)
xl.Infof("starting VirtualNetPlugin for visitor [%s], attempting to register routes for %s", p.pluginCtx.Name, routeStr)
go p.run()
}
@@ -101,10 +101,8 @@ func (p *VirtualNetPlugin) run() {
reconnectDelay := 10 * time.Second
for {
// Create a signal channel for this connection attempt
currentCloseSignal := make(chan struct{})
// Store the signal channel under lock
p.mu.Lock()
p.closeSignal = currentCloseSignal
p.mu.Unlock()
@@ -112,7 +110,6 @@ 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)
// Ensure controllerConn from previous loop is cleaned up if necessary
p.cleanupControllerConn(xl)
return
default:
@@ -120,65 +117,43 @@ func (p *VirtualNetPlugin) run() {
controllerConn, pluginConn := net.Pipe()
// Store controllerConn under lock for cleanup purposes
p.mu.Lock()
p.controllerConn = controllerConn
p.mu.Unlock()
// Wrap pluginConn using CloseNotifyConn
pluginNotifyConn := netutil.WrapCloseNotifyConn(pluginConn, func() {
close(currentCloseSignal) // Signal the run loop
close(currentCloseSignal) // Signal the run loop on close.
})
xl.Infof("Attempting to register client route for visitor [%s]", p.pluginCtx.Name)
err := p.pluginCtx.VnetController.RegisterClientRoute(p.ctx, p.pluginCtx.Name, p.routes, controllerConn)
if err != nil {
xl.Errorf("Failed to register client route for visitor [%s]: %v. Retrying after %v", p.pluginCtx.Name, err, reconnectDelay)
p.cleanupPipePair(xl, controllerConn, pluginConn) // Close both ends on registration failure
// Wait before retrying registration, unless context is cancelled
select {
case <-time.After(reconnectDelay):
continue // Retry the loop
case <-p.ctx.Done():
xl.Infof("VirtualNetPlugin registration retry wait interrupted for visitor [%s]", p.pluginCtx.Name)
return // Exit loop if context is cancelled during wait
}
}
xl.Infof("Successfully registered client route for visitor [%s]. Starting connection handler with CloseNotifyConn.", p.pluginCtx.Name)
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 HandleConn.
// HandleConn is responsible for calling Close() on pluginNotifyConn.
p.pluginCtx.HandleConn(pluginNotifyConn)
// Wait for either the plugin context to be cancelled or the wrapper's Close() to be called via the signal channel.
// Wait for context cancellation or connection close.
select {
case <-p.ctx.Done():
xl.Infof("VirtualNetPlugin run loop stopping for visitor [%s] (context cancelled while waiting).", p.pluginCtx.Name)
// Context cancelled, ensure controller side is closed if HandleConn didn't close its side yet.
p.cleanupControllerConn(xl)
return
case <-currentCloseSignal:
xl.Infof("Detected connection closed via CloseNotifyConn for visitor [%s].", p.pluginCtx.Name)
// HandleConn closed the plugin side (pluginNotifyConn). The closeFn was called, closing currentCloseSignal.
// We still need to close the controller side.
xl.Infof("detected connection closed via CloseNotifyConn for visitor [%s].", p.pluginCtx.Name)
// HandleConn closed the plugin side. Close the controller side.
p.cleanupControllerConn(xl)
// Add a delay before attempting to reconnect, respecting context cancellation.
xl.Infof("Waiting %v before attempting reconnection for visitor [%s]...", reconnectDelay, p.pluginCtx.Name)
xl.Infof("waiting %v before attempting reconnection for visitor [%s]...", reconnectDelay, p.pluginCtx.Name)
select {
case <-time.After(reconnectDelay):
// Delay completed, loop will continue.
case <-p.ctx.Done():
xl.Infof("VirtualNetPlugin reconnection delay interrupted for visitor [%s]", p.pluginCtx.Name)
return // Exit loop if context is cancelled during wait
return
}
// Loop will continue to reconnect.
}
// Loop will restart, context check at the beginning of the loop is sufficient.
xl.Infof("Re-establishing virtual connection for visitor [%s]...", p.pluginCtx.Name)
xl.Infof("re-establishing virtual connection for visitor [%s]...", p.pluginCtx.Name)
}
}
@@ -187,46 +162,31 @@ func (p *VirtualNetPlugin) cleanupControllerConn(xl *xlog.Logger) {
p.mu.Lock()
defer p.mu.Unlock()
if p.controllerConn != nil {
xl.Debugf("Cleaning up controllerConn for visitor [%s]", p.pluginCtx.Name)
xl.Debugf("cleaning up controllerConn for visitor [%s]", p.pluginCtx.Name)
p.controllerConn.Close()
p.controllerConn = nil
}
// Also clear the closeSignal reference for the completed/cancelled connection attempt
p.closeSignal = nil
}
// cleanupPipePair closes both ends of a pipe, used typically when registration fails.
func (p *VirtualNetPlugin) cleanupPipePair(xl *xlog.Logger, controllerConn, pluginConn net.Conn) {
xl.Debugf("Cleaning up pipe pair for visitor [%s] after registration failure", p.pluginCtx.Name)
controllerConn.Close()
pluginConn.Close()
p.mu.Lock()
p.controllerConn = nil // Ensure field is nil if it was briefly set
p.closeSignal = nil // Ensure field is nil if it was briefly set
p.mu.Unlock()
}
// Close initiates the plugin shutdown.
func (p *VirtualNetPlugin) Close() error {
xl := xlog.FromContextSafe(p.pluginCtx.Ctx) // Use base context for close logging
xl.Infof("Closing VirtualNetPlugin for visitor [%s]", p.pluginCtx.Name)
xl := xlog.FromContextSafe(p.pluginCtx.Ctx)
xl.Infof("closing VirtualNetPlugin for visitor [%s]", p.pluginCtx.Name)
// 1. Signal the run loop goroutine to stop via context cancellation.
// Signal the run loop goroutine to stop.
p.cancel()
// 2. Unregister the route from the controller.
// This might implicitly cause the VnetController to close its end of the pipe (controllerConn).
// Unregister the route from the controller.
if p.pluginCtx.VnetController != nil {
p.pluginCtx.VnetController.UnregisterClientRoute(p.pluginCtx.Name)
xl.Infof("Unregistered client route for visitor [%s]", p.pluginCtx.Name)
} else {
xl.Warnf("VnetController is nil during close for visitor [%s], cannot unregister route", p.pluginCtx.Name)
xl.Infof("unregistered client route for visitor [%s]", p.pluginCtx.Name)
}
// 3. Explicitly close the controller side of the pipe managed by this plugin.
// Explicitly close the controller side of the pipe.
// This ensures the pipe is broken even if the run loop is stuck or HandleConn hasn't closed its end.
p.cleanupControllerConn(xl)
xl.Infof("Finished cleaning up connections during close for visitor [%s]", p.pluginCtx.Name)
xl.Infof("finished cleaning up connections during close for visitor [%s]", p.pluginCtx.Name)
return nil
}