Add pause function for proxies

client/internal/wgproxy/ebpf/proxy.go

@@ -5,7 +5,6 @@ package ebpf
 import (
 	"context"
 	"fmt"
-	"io"
 	"net"
 	"os"
 	"sync"
@@ -94,13 +93,12 @@ func (p *WGEBPFProxy) Listen() error {
 }
 
 // AddTurnConn add new turn connection for the proxy
-func (p *WGEBPFProxy) AddTurnConn(ctx context.Context, turnConn net.Conn) (net.Addr, error) {
+func (p *WGEBPFProxy) AddTurnConn(turnConn net.Conn) (*net.UDPAddr, error) {
 	wgEndpointPort, err := p.storeTurnConn(turnConn)
 	if err != nil {
 		return nil, err
 	}
 
-	go p.proxyToLocal(ctx, wgEndpointPort, turnConn)
 	log.Infof("turn conn added to wg proxy store: %s, endpoint port: :%d", turnConn.RemoteAddr(), wgEndpointPort)
 
 	wgEndpoint := &net.UDPAddr{
@@ -137,35 +135,6 @@ func (p *WGEBPFProxy) Free() error {
 	return nberrors.FormatErrorOrNil(result)
 }
 
-func (p *WGEBPFProxy) proxyToLocal(ctx context.Context, endpointPort uint16, remoteConn net.Conn) {
-	defer p.removeTurnConn(endpointPort)
-
-	var (
-		err error
-		n   int
-	)
-	buf := make([]byte, 1500)
-	for ctx.Err() == nil {
-		n, err = remoteConn.Read(buf)
-		if err != nil {
-			if ctx.Err() != nil {
-				return
-			}
-			if err != io.EOF {
-				log.Errorf("failed to read from turn conn (endpoint: :%d): %s", endpointPort, err)
-			}
-			return
-		}
-
-		if err := p.sendPkg(buf[:n], endpointPort); err != nil {
-			if ctx.Err() != nil || p.ctx.Err() != nil {
-				return
-			}
-			log.Errorf("failed to write out turn pkg to local conn: %v", err)
-		}
-	}
-}
-
 // proxyToRemote read messages from local WireGuard interface and forward it to remote conn
 // From this go routine has only one instance.
 func (p *WGEBPFProxy) proxyToRemote() {

client/internal/wgproxy/ebpf/wrapper.go

@@ -5,7 +5,11 @@ package ebpf
 import (
 	"context"
 	"fmt"
+	"io"
 	"net"
+	"sync"
+
+	log "github.com/sirupsen/logrus"
 )
 
 // ProxyWrapper help to keep the remoteConn instance for net.Conn.Close function call
@@ -13,22 +17,62 @@ type ProxyWrapper struct {
 	WgeBPFProxy *WGEBPFProxy
 
 	remoteConn net.Conn
-	cancel     context.CancelFunc // with thic cancel function, we stop remoteToLocal thread
+	ctx        context.Context
+	cancel     context.CancelFunc
+
+	wgEndpointPort uint16
+
+	pausedMu  sync.Mutex
+	paused    bool
+	isStarted bool
 }
 
-func (e *ProxyWrapper) AddTurnConn(ctx context.Context, remoteConn net.Conn) (net.Addr, error) {
-	ctxConn, cancel := context.WithCancel(ctx)
-	addr, err := e.WgeBPFProxy.AddTurnConn(ctxConn, remoteConn)
-
+func (p *ProxyWrapper) AddTurnConn(ctx context.Context, remoteConn net.Conn) (net.Addr, error) {
+	addr, err := p.WgeBPFProxy.AddTurnConn(remoteConn)
 	if err != nil {
-		cancel()
 		return nil, fmt.Errorf("add turn conn: %w", err)
 	}
-	e.remoteConn = remoteConn
-	e.cancel = cancel
+	p.remoteConn = remoteConn
+	p.ctx, p.cancel = context.WithCancel(ctx)
+	p.wgEndpointPort = uint16(addr.Port)
 	return addr, err
 }
 
+func (p *ProxyWrapper) Work() {
+	if p.remoteConn == nil {
+		return
+	}
+
+	p.pausedMu.Lock()
+	p.paused = false
+	p.pausedMu.Unlock()
+
+	if !p.isStarted {
+		p.isStarted = true
+		go p.proxyToLocal(p.ctx)
+	}
+}
+
+func (p *ProxyWrapper) Pause() {
+	if p.remoteConn == nil {
+		return
+	}
+
+	p.pausedMu.Lock()
+	p.paused = true
+	p.pausedMu.Unlock()
+}
+
+func (p *ProxyWrapper) Resume() {
+	if p.remoteConn == nil {
+		return
+	}
+
+	p.pausedMu.Lock()
+	p.paused = false
+	p.pausedMu.Unlock()
+}
+
 // CloseConn close the remoteConn and automatically remove the conn instance from the map
 func (e *ProxyWrapper) CloseConn() error {
 	if e.cancel == nil {
@@ -42,3 +86,32 @@ func (e *ProxyWrapper) CloseConn() error {
 	}
 	return nil
 }
+
+func (p *ProxyWrapper) proxyToLocal(ctx context.Context) {
+	defer p.WgeBPFProxy.removeTurnConn(p.wgEndpointPort)
+
+	var (
+		err error
+		n   int
+	)
+	buf := make([]byte, 1500)
+	for ctx.Err() == nil {
+		n, err = p.remoteConn.Read(buf)
+		if err != nil {
+			if ctx.Err() != nil {
+				return
+			}
+			if err != io.EOF {
+				log.Errorf("failed to read from turn conn (endpoint: :%d): %s", p.wgEndpointPort, err)
+			}
+			return
+		}
+
+		if err := p.WgeBPFProxy.sendPkg(buf[:n], p.wgEndpointPort); err != nil {
+			if ctx.Err() != nil {
+				return
+			}
+			log.Errorf("failed to write out turn pkg to local conn: %v", err)
+		}
+	}
+}