mirror of
https://github.com/fatedier/frp.git
synced 2026-01-11 22:23:12 +00:00
vnet: fix issues (#4771)
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user