Replace the eBPF WireGuard proxy with loopback endpoint addressing

2026-08-31 03:51:29 +02:00 · 2026-08-25 12:27:47 +02:00
parent 7f03a2e86f
commit 2a3b0cec08
23 changed files with 958 additions and 536 deletions
--- a/client/iface/wgproxy/loopback/addr.go
+++ b/client/iface/wgproxy/loopback/addr.go
@@ -0,0 +1,70 @@
+//go:build linux && !android
+
+package loopback
+
+import (
+	"fmt"
+	"net/netip"
+)
+
+// Peer endpoints live in the upper half of 127.0.0.0/8. Everything in that
+// range is delivered to the loopback device without any address or route being
+// configured, and staying out of 127.0.0.0/9 keeps well-known squatters such as
+// 127.0.0.53 (systemd-resolved) and 127.0.1.1 out of the way.
+const (
+	addrRangeBase   uint32 = 0x7f800000 // 127.128.0.0
+	addrRangeSize   uint32 = 1 << 23    // /9
+	addrRangePrefix        = "127.128.0.0/9"
+)
+
+// allocator hands out one loopback address per relayed connection. The address
+// is the peer's identity: WireGuard sends to it, and the proxy recovers which
+// peer a packet belongs to from the destination address.
+type allocator struct {
+	cursor uint32
+}
+
+// next returns the first free address at or after the cursor, wrapping once.
+// inUse reports whether an address is already handed out.
+func (a *allocator) next(inUse func(netip.Addr) bool) (netip.Addr, error) {
+	for i := uint32(0); i < addrRangeSize; i++ {
+		a.cursor = (a.cursor + 1) % addrRangeSize
+		addr := addrFromOffset(a.cursor)
+		if !addr.IsValid() {
+			continue
+		}
+		if inUse(addr) {
+			continue
+		}
+		return addr, nil
+	}
+	return netip.Addr{}, fmt.Errorf("no free endpoint address in %s", addrRangePrefix)
+}
+
+// addrFromOffset maps an offset in the range to an address, skipping the .0 and
+// .255 hosts. They are unremarkable on loopback, but tools and firewall rules
+// tend to treat them as network and broadcast addresses.
+func addrFromOffset(offset uint32) netip.Addr {
+	last := offset & 0xff
+	if last == 0 || last == 0xff {
+		return netip.Addr{}
+	}
+
+	v := addrRangeBase + offset
+	return netip.AddrFrom4([4]byte{
+		byte(v >> 24),
+		byte(v >> 16),
+		byte(v >> 8),
+		byte(v),
+	})
+}
+
+// inRange reports whether addr is one this proxy could have handed out.
+func inRange(addr netip.Addr) bool {
+	if !addr.Is4() {
+		return false
+	}
+	b := addr.As4()
+	v := uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
+	return v >= addrRangeBase && v < addrRangeBase+addrRangeSize
+}
--- a/client/iface/wgproxy/loopback/addr_test.go
+++ b/client/iface/wgproxy/loopback/addr_test.go
@@ -0,0 +1,103 @@
+//go:build linux && !android
+
+package loopback
+
+import (
+	"net/netip"
+	"testing"
+)
+
+func TestAllocatorHandsOutDistinctAddresses(t *testing.T) {
+	var a allocator
+	taken := make(map[netip.Addr]bool)
+
+	for i := 0; i < 1000; i++ {
+		addr, err := a.next(func(candidate netip.Addr) bool { return taken[candidate] })
+		if err != nil {
+			t.Fatalf("allocate %d: %v", i, err)
+		}
+		if taken[addr] {
+			t.Fatalf("address %s handed out twice", addr)
+		}
+		if !inRange(addr) {
+			t.Fatalf("address %s outside %s", addr, addrRangePrefix)
+		}
+		taken[addr] = true
+	}
+}
+
+func TestAllocatorSkipsNetworkAndBroadcastHosts(t *testing.T) {
+	var a allocator
+	taken := make(map[netip.Addr]bool)
+
+	// enough allocations to walk past a .255/.0 boundary
+	for i := 0; i < 600; i++ {
+		addr, err := a.next(func(candidate netip.Addr) bool { return taken[candidate] })
+		if err != nil {
+			t.Fatalf("allocate %d: %v", i, err)
+		}
+		last := addr.As4()[3]
+		if last == 0 || last == 255 {
+			t.Fatalf("address %s ends in .%d", addr, last)
+		}
+		taken[addr] = true
+	}
+}
+
+func TestAllocatorReusesReleasedAddresses(t *testing.T) {
+	var a allocator
+	taken := make(map[netip.Addr]bool)
+
+	first, err := a.next(func(candidate netip.Addr) bool { return taken[candidate] })
+	if err != nil {
+		t.Fatalf("allocate: %v", err)
+	}
+	taken[first] = true
+
+	// release it and allocate until the cursor wraps back around to it
+	delete(taken, first)
+	for i := 0; i < 10; i++ {
+		addr, err := a.next(func(candidate netip.Addr) bool { return taken[candidate] })
+		if err != nil {
+			t.Fatalf("allocate %d: %v", i, err)
+		}
+		if addr == first {
+			return
+		}
+		taken[addr] = true
+	}
+	// the cursor moves forward, so reuse only happens after a full wrap. Assert
+	// the released address is at least still considered free.
+	if inUse := taken[first]; inUse {
+		t.Fatalf("released address %s still marked in use", first)
+	}
+}
+
+func TestInRange(t *testing.T) {
+	tests := []struct {
+		addr string
+		want bool
+	}{
+		{"127.128.0.1", true},
+		{"127.255.255.254", true},
+		{"127.127.255.255", false}, // below the range, where 127.0.0.53 and friends live
+		{"127.0.0.1", false},
+		{"127.0.0.53", false},
+		{"127.0.1.1", false},
+		{"128.0.0.1", false},
+		{"10.0.0.1", false},
+	}
+
+	for _, tc := range tests {
+		addr := netip.MustParseAddr(tc.addr)
+		if got := inRange(addr); got != tc.want {
+			t.Errorf("inRange(%s) = %v, want %v", tc.addr, got, tc.want)
+		}
+	}
+}
+
+func TestInRangeIgnoresIPv6(t *testing.T) {
+	if inRange(netip.MustParseAddr("::1")) {
+		t.Error("inRange(::1) = true, want false")
+	}
+}
--- a/client/iface/wgproxy/loopback/proxy.go
+++ b/client/iface/wgproxy/loopback/proxy.go
@@ -0,0 +1,275 @@
+//go:build linux && !android
+
+package loopback
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"net/netip"
+	"sync"
+	"syscall"
+
+	"github.com/hashicorp/go-multierror"
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/net/ipv4"
+	"golang.org/x/sys/unix"
+
+	nberrors "github.com/netbirdio/netbird/client/errors"
+	"github.com/netbirdio/netbird/client/iface/bufsize"
+	"github.com/netbirdio/netbird/client/iface/wgproxy/rawsocket"
+)
+
+const (
+	loopbackDevice = "lo"
+
+	portRangeStart = 3128
+	portRangeEnd   = portRangeStart + 100
+)
+
+// Proxy forwards packets between relayed connections and a local kernel
+// WireGuard instance. Every relayed peer gets its own loopback address as its
+// WireGuard endpoint, so a single socket serves all of them: the destination
+// address of an incoming packet identifies the peer.
+type Proxy struct {
+	localWGListenPort int
+	mtu               uint16
+	proxyPort         int
+
+	conn        *net.UDPConn
+	packetConn  *ipv4.PacketConn
+	rawConnIPv4 net.PacketConn
+	rawConnIPv6 net.PacketConn
+
+	relayedConnMutex sync.Mutex
+	relayedConnStore map[netip.Addr]net.Conn
+	addrs            allocator
+
+	ctx       context.Context
+	ctxCancel context.CancelFunc
+}
+
+// NewProxy creates a proxy for the WireGuard instance listening on wgPort.
+func NewProxy(wgPort int, mtu uint16) *Proxy {
+	log.Debugf("instantiate loopback wg proxy")
+	return &Proxy{
+		localWGListenPort: wgPort,
+		mtu:               mtu,
+		relayedConnStore:  make(map[netip.Addr]net.Conn),
+	}
+}
+
+// Listen opens the shared socket and starts forwarding WireGuard packets to the
+// relayed connections.
+func (p *Proxy) Listen() error {
+	rawConnIPv4, err := rawsocket.PrepareSenderRawSocketIPv4()
+	if err != nil {
+		return fmt.Errorf("prepare IPv4 raw socket: %w", err)
+	}
+	p.rawConnIPv4 = rawConnIPv4
+
+	p.rawConnIPv6, err = rawsocket.PrepareSenderRawSocketIPv6()
+	if err != nil {
+		log.Warnf("failed to prepare IPv6 raw socket, continuing with IPv4 only: %v", err)
+	}
+
+	if err := p.listen(); err != nil {
+		if freeErr := p.Free(); freeErr != nil {
+			log.Errorf("failed to free the wgproxy: %s", freeErr)
+		}
+		return err
+	}
+
+	p.ctx, p.ctxCancel = context.WithCancel(context.Background())
+
+	go p.proxyToRemote()
+	log.Infof("local wg proxy listening on %s:%d", addrRangePrefix, p.proxyPort)
+	return nil
+}
+
+// listen binds the shared socket on the first free port of the range. The bind
+// has to be a wildcard one to receive every peer address in the range, so it is
+// restricted to the loopback device: without that the port would be reachable
+// on every interface.
+func (p *Proxy) listen() error {
+	var lastErr error
+	for port := portRangeStart; port <= portRangeEnd; port++ {
+		err := p.listenOn(port)
+		if err == nil {
+			p.proxyPort = port
+			return nil
+		}
+		lastErr = err
+	}
+	return fmt.Errorf("bind proxy port in range %d-%d: %w", portRangeStart, portRangeEnd, lastErr)
+}
+
+func (p *Proxy) listenOn(proxyPort int) error {
+	lc := net.ListenConfig{
+		Control: func(_, _ string, c syscall.RawConn) error {
+			var sockErr error
+			if err := c.Control(func(fd uintptr) {
+				if err := unix.SetsockoptString(int(fd), unix.SOL_SOCKET, unix.SO_BINDTODEVICE, loopbackDevice); err != nil {
+					sockErr = fmt.Errorf("bind to %s: %w", loopbackDevice, err)
+					return
+				}
+			}); err != nil {
+				return fmt.Errorf("control socket: %w", err)
+			}
+			return sockErr
+		},
+	}
+
+	conn, err := lc.ListenPacket(context.Background(), "udp4", fmt.Sprintf(":%d", proxyPort))
+	if err != nil {
+		return fmt.Errorf("listen on :%d: %w", proxyPort, err)
+	}
+
+	udpConn, ok := conn.(*net.UDPConn)
+	if !ok {
+		if closeErr := conn.Close(); closeErr != nil {
+			log.Errorf("failed to close proxy conn: %s", closeErr)
+		}
+		return fmt.Errorf("unexpected conn type %T", conn)
+	}
+
+	packetConn := ipv4.NewPacketConn(udpConn)
+	// the destination address carries the peer identity, the interface index is
+	// checked on receive as a second line of defense behind SO_BINDTODEVICE
+	if err := packetConn.SetControlMessage(ipv4.FlagDst|ipv4.FlagInterface, true); err != nil {
+		if closeErr := udpConn.Close(); closeErr != nil {
+			log.Errorf("failed to close proxy conn: %s", closeErr)
+		}
+		return fmt.Errorf("request destination address: %w", err)
+	}
+
+	p.conn = udpConn
+	p.packetConn = packetConn
+	return nil
+}
+
+// AddRelayedConn assigns an endpoint address to the relayed connection and
+// returns the address WireGuard should send to.
+func (p *Proxy) AddRelayedConn(relayedConn net.Conn) (*net.UDPAddr, error) {
+	addr, err := p.storeRelayedConn(relayedConn)
+	if err != nil {
+		return nil, err
+	}
+
+	log.Infof("relayed conn added to wg proxy store: %s, endpoint address: %s", relayedConn.RemoteAddr(), addr)
+
+	return &net.UDPAddr{
+		IP:   addr.AsSlice(),
+		Port: p.proxyPort,
+	}, nil
+}
+
+// Free releases the proxy resources. The relayed connections are left open.
+func (p *Proxy) Free() error {
+	log.Debugf("free up loopback wg proxy")
+	if p.ctx != nil && p.ctx.Err() != nil {
+		//nolint
+		return nil
+	}
+
+	if p.ctxCancel != nil {
+		p.ctxCancel()
+	}
+
+	var result *multierror.Error
+	if p.conn != nil {
+		if err := p.conn.Close(); err != nil {
+			result = multierror.Append(result, err)
+		}
+	}
+
+	if p.rawConnIPv4 != nil {
+		if err := p.rawConnIPv4.Close(); err != nil {
+			result = multierror.Append(result, err)
+		}
+	}
+
+	if p.rawConnIPv6 != nil {
+		if err := p.rawConnIPv6.Close(); err != nil {
+			result = multierror.Append(result, err)
+		}
+	}
+	return nberrors.FormatErrorOrNil(result)
+}
+
+// GetProxyPort returns the port every peer endpoint address is reached on.
+func (p *Proxy) GetProxyPort() uint16 {
+	return uint16(p.proxyPort)
+}
+
+// proxyToRemote reads packets from the local WireGuard instance and forwards
+// them to the relayed connection the destination address belongs to.
+func (p *Proxy) proxyToRemote() {
+	buf := make([]byte, p.mtu+bufsize.WGBufferOverhead)
+	for p.ctx.Err() == nil {
+		if err := p.readAndForwardPacket(buf); err != nil {
+			if p.ctx.Err() != nil {
+				return
+			}
+			log.Errorf("failed to proxy packet to remote conn: %s", err)
+		}
+	}
+}
+
+func (p *Proxy) readAndForwardPacket(buf []byte) error {
+	n, cm, _, err := p.packetConn.ReadFrom(buf)
+	if err != nil {
+		return fmt.Errorf("read UDP packet from WG: %w", err)
+	}
+
+	if cm == nil {
+		return fmt.Errorf("no control message on packet")
+	}
+
+	dst, ok := netip.AddrFromSlice(cm.Dst.To4())
+	if !ok || !inRange(dst) {
+		log.Tracef("dropping packet for unexpected destination %s", cm.Dst)
+		return nil
+	}
+
+	p.relayedConnMutex.Lock()
+	conn, ok := p.relayedConnStore[dst]
+	p.relayedConnMutex.Unlock()
+	if !ok {
+		if p.ctx.Err() == nil {
+			log.Debugf("relayed conn not found by address because conn already has been closed: %s", dst)
+		}
+		return nil
+	}
+
+	if _, err := conn.Write(buf[:n]); err != nil {
+		return fmt.Errorf("forward local WG packet (%s) to remote relayed conn: %w", dst, err)
+	}
+	return nil
+}
+
+func (p *Proxy) storeRelayedConn(relayedConn net.Conn) (netip.Addr, error) {
+	p.relayedConnMutex.Lock()
+	defer p.relayedConnMutex.Unlock()
+
+	addr, err := p.addrs.next(func(a netip.Addr) bool {
+		_, ok := p.relayedConnStore[a]
+		return ok
+	})
+	if err != nil {
+		return netip.Addr{}, err
+	}
+
+	p.relayedConnStore[addr] = relayedConn
+	return addr, nil
+}
+
+func (p *Proxy) removeRelayedConn(addr netip.Addr) {
+	p.relayedConnMutex.Lock()
+	defer p.relayedConnMutex.Unlock()
+
+	if _, ok := p.relayedConnStore[addr]; ok {
+		log.Debugf("remove relayed conn from store by address: %s", addr)
+	}
+	delete(p.relayedConnStore, addr)
+}
--- a/client/iface/wgproxy/loopback/proxy_privileged_test.go
+++ b/client/iface/wgproxy/loopback/proxy_privileged_test.go
@@ -0,0 +1,163 @@
+//go:build linux && !android && privileged
+
+package loopback
+
+import (
+	"net"
+	"strconv"
+	"testing"
+	"time"
+)
+
+const testWGPort = 51862
+
+// relayEnd stands in for a relayed connection: the proxy writes what it read
+// from WireGuard into it, and the test reads it back out here.
+func relayEnd(t *testing.T) (proxySide net.Conn, testSide *net.UDPConn) {
+	t.Helper()
+
+	testSide, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1")})
+	if err != nil {
+		t.Fatalf("relay listener: %v", err)
+	}
+	t.Cleanup(func() {
+		if err := testSide.Close(); err != nil {
+			t.Logf("close relay listener: %v", err)
+		}
+	})
+
+	proxySide, err = net.Dial("udp", testSide.LocalAddr().String())
+	if err != nil {
+		t.Fatalf("relay conn: %v", err)
+	}
+	t.Cleanup(func() {
+		if err := proxySide.Close(); err != nil {
+			t.Logf("close relay conn: %v", err)
+		}
+	})
+
+	return proxySide, testSide
+}
+
+// TestProxyDemuxesByDestinationAddress is the core of the design: one socket
+// serves every peer, and the destination address decides which relayed
+// connection a WireGuard packet belongs to.
+func TestProxyDemuxesByDestinationAddress(t *testing.T) {
+	proxy := NewProxy(testWGPort, 1280)
+	if err := proxy.Listen(); err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	defer func() {
+		if err := proxy.Free(); err != nil {
+			t.Errorf("free proxy: %v", err)
+		}
+	}()
+
+	const peers = 3
+	endpoints := make([]*net.UDPAddr, 0, peers)
+	readers := make([]*net.UDPConn, 0, peers)
+	for i := 0; i < peers; i++ {
+		proxySide, testSide := relayEnd(t)
+		endpoint, err := proxy.AddRelayedConn(proxySide)
+		if err != nil {
+			t.Fatalf("add relayed conn %d: %v", i, err)
+		}
+		if endpoint.Port != proxy.proxyPort {
+			t.Errorf("peer %d endpoint port = %d, want the shared proxy port %d", i, endpoint.Port, proxy.proxyPort)
+		}
+		endpoints = append(endpoints, endpoint)
+		readers = append(readers, testSide)
+	}
+
+	// every peer must have its own address, otherwise they are indistinguishable
+	seen := make(map[string]bool, peers)
+	for i, endpoint := range endpoints {
+		if seen[endpoint.IP.String()] {
+			t.Fatalf("peer %d reuses endpoint address %s", i, endpoint.IP)
+		}
+		seen[endpoint.IP.String()] = true
+	}
+
+	wgSock, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: testWGPort})
+	if err != nil {
+		t.Fatalf("wg socket: %v", err)
+	}
+	defer func() {
+		if err := wgSock.Close(); err != nil {
+			t.Logf("close wg socket: %v", err)
+		}
+	}()
+
+	for i, endpoint := range endpoints {
+		payload := []byte{byte(i), 'p', 'k', 't'}
+		if _, err := wgSock.WriteTo(payload, endpoint); err != nil {
+			t.Fatalf("write to peer %d endpoint %s: %v", i, endpoint, err)
+		}
+
+		buf := make([]byte, 1500)
+		if err := readers[i].SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil {
+			t.Fatalf("set read deadline: %v", err)
+		}
+		n, _, err := readers[i].ReadFrom(buf)
+		if err != nil {
+			t.Fatalf("peer %d did not receive its packet: %v", i, err)
+		}
+		if string(buf[:n]) != string(payload) {
+			t.Errorf("peer %d got %q, want %q", i, buf[:n], payload)
+		}
+
+		// no other peer may see it
+		for j, other := range readers {
+			if j == i {
+				continue
+			}
+			if err := other.SetReadDeadline(time.Now().Add(100 * time.Millisecond)); err != nil {
+				t.Fatalf("set read deadline: %v", err)
+			}
+			if _, _, err := other.ReadFrom(buf); err == nil {
+				t.Errorf("packet for peer %d also delivered to peer %d", i, j)
+			}
+		}
+	}
+}
+
+// TestProxyDropsPacketsOutsideTheRange guards the wildcard bind: anything that
+// is not addressed to a handed-out endpoint must not reach a relayed peer.
+func TestProxyDropsPacketsOutsideTheRange(t *testing.T) {
+	proxy := NewProxy(testWGPort+1, 1280)
+	if err := proxy.Listen(); err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	defer func() {
+		if err := proxy.Free(); err != nil {
+			t.Errorf("free proxy: %v", err)
+		}
+	}()
+
+	proxySide, testSide := relayEnd(t)
+	if _, err := proxy.AddRelayedConn(proxySide); err != nil {
+		t.Fatalf("add relayed conn: %v", err)
+	}
+
+	sender, err := net.Dial("udp", net.JoinHostPort("127.0.0.1", strconv.Itoa(proxy.proxyPort)))
+	if err != nil {
+		t.Fatalf("sender: %v", err)
+	}
+	defer func() {
+		if err := sender.Close(); err != nil {
+			t.Logf("close sender: %v", err)
+		}
+	}()
+
+	if _, err := sender.Write([]byte("stray")); err != nil {
+		t.Fatalf("write stray packet: %v", err)
+	}
+
+	buf := make([]byte, 1500)
+	if err := testSide.SetReadDeadline(time.Now().Add(500 * time.Millisecond)); err != nil {
+		t.Fatalf("set read deadline: %v", err)
+	}
+	if _, _, err := testSide.ReadFrom(buf); err == nil {
+		t.Error("packet addressed to 127.0.0.1 was forwarded to a relayed peer")
+	}
+}
--- a/client/iface/wgproxy/loopback/wrapper.go
+++ b/client/iface/wgproxy/loopback/wrapper.go
@@ -0,0 +1,328 @@
+//go:build linux && !android
+
+package loopback
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/netip"
+	"sync"
+
+	"github.com/google/gopacket"
+	"github.com/google/gopacket/layers"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/client/iface/bufsize"
+	"github.com/netbirdio/netbird/client/iface/wgproxy/listener"
+)
+
+var (
+	errIPv6ConnNotAvailable = errors.New("IPv6 endpoint but rawConnIPv6 is not available")
+	errIPv4ConnNotAvailable = errors.New("IPv4 endpoint but rawConnIPv4 is not available")
+
+	localHostNetIPv4 = net.ParseIP("127.0.0.1")
+	localHostNetIPv6 = net.ParseIP("::1")
+
+	serializeOpts = gopacket.SerializeOptions{
+		ComputeChecksums: true,
+		FixLengths:       true,
+	}
+)
+
+// PacketHeaders holds pre-created headers and buffers for efficient packet sending
+type PacketHeaders struct {
+	ipH           gopacket.SerializableLayer
+	udpH          *layers.UDP
+	layerBuffer   gopacket.SerializeBuffer
+	localHostAddr net.IP
+	isIPv4        bool
+}
+
+func NewPacketHeaders(localWGListenPort int, endpoint *net.UDPAddr) (*PacketHeaders, error) {
+	var ipH gopacket.SerializableLayer
+	var networkLayer gopacket.NetworkLayer
+	var localHostAddr net.IP
+	var isIPv4 bool
+
+	// Check if source address is IPv4 or IPv6
+	if endpoint.IP.To4() != nil {
+		// IPv4 path
+		ipv4 := &layers.IPv4{
+			DstIP:    localHostNetIPv4,
+			SrcIP:    endpoint.IP,
+			Version:  4,
+			TTL:      64,
+			Protocol: layers.IPProtocolUDP,
+		}
+		ipH = ipv4
+		networkLayer = ipv4
+		localHostAddr = localHostNetIPv4
+		isIPv4 = true
+	} else {
+		// IPv6 path
+		ipv6 := &layers.IPv6{
+			DstIP:      localHostNetIPv6,
+			SrcIP:      endpoint.IP,
+			Version:    6,
+			HopLimit:   64,
+			NextHeader: layers.IPProtocolUDP,
+		}
+		ipH = ipv6
+		networkLayer = ipv6
+		localHostAddr = localHostNetIPv6
+		isIPv4 = false
+	}
+
+	udpH := &layers.UDP{
+		SrcPort: layers.UDPPort(endpoint.Port),
+		DstPort: layers.UDPPort(localWGListenPort),
+	}
+
+	if err := udpH.SetNetworkLayerForChecksum(networkLayer); err != nil {
+		return nil, fmt.Errorf("set network layer for checksum: %w", err)
+	}
+
+	return &PacketHeaders{
+		ipH:           ipH,
+		udpH:          udpH,
+		layerBuffer:   gopacket.NewSerializeBuffer(),
+		localHostAddr: localHostAddr,
+		isIPv4:        isIPv4,
+	}, nil
+}
+
+// ProxyWrapper help to keep the remoteConn instance for net.Conn.Close function call
+type ProxyWrapper struct {
+	proxy *Proxy
+
+	remoteConn net.Conn
+	ctx        context.Context
+	cancel     context.CancelFunc
+
+	wgRelayedEndpointAddr *net.UDPAddr
+	peerAddr              netip.Addr
+	headers               *PacketHeaders
+	headerCurrentUsed     *PacketHeaders
+	rawConn               net.PacketConn
+
+	paused     bool
+	pausedCond *sync.Cond
+	isStarted  bool
+
+	closeListener *listener.CloseListener
+}
+
+func NewProxyWrapper(proxy *Proxy) *ProxyWrapper {
+	return &ProxyWrapper{
+		proxy:         proxy,
+		pausedCond:    sync.NewCond(&sync.Mutex{}),
+		closeListener: listener.NewCloseListener(),
+	}
+}
+
+func (p *ProxyWrapper) AddRelayedConn(ctx context.Context, _ *net.UDPAddr, remoteConn net.Conn) error {
+	addr, err := p.proxy.AddRelayedConn(remoteConn)
+	if err != nil {
+		return fmt.Errorf("add relayed conn: %w", err)
+	}
+
+	peerAddr, ok := netip.AddrFromSlice(addr.IP.To4())
+	if !ok {
+		return fmt.Errorf("unexpected endpoint address %s", addr.IP)
+	}
+
+	headers, err := NewPacketHeaders(p.proxy.localWGListenPort, addr)
+	if err != nil {
+		return fmt.Errorf("create packet sender: %w", err)
+	}
+
+	// Check if required raw connection is available
+	if !headers.isIPv4 && p.proxy.rawConnIPv6 == nil {
+		return errIPv6ConnNotAvailable
+	}
+	if headers.isIPv4 && p.proxy.rawConnIPv4 == nil {
+		return errIPv4ConnNotAvailable
+	}
+
+	p.remoteConn = remoteConn
+	p.ctx, p.cancel = context.WithCancel(ctx)
+	p.wgRelayedEndpointAddr = addr
+	p.peerAddr = peerAddr
+	p.headers = headers
+	p.rawConn = p.selectRawConn(headers)
+	return nil
+}
+
+func (p *ProxyWrapper) EndpointAddr() *net.UDPAddr {
+	return p.wgRelayedEndpointAddr
+}
+
+func (p *ProxyWrapper) SetDisconnectListener(disconnected func()) {
+	p.closeListener.SetCloseListener(disconnected)
+}
+
+func (p *ProxyWrapper) Work() {
+	if p.remoteConn == nil {
+		return
+	}
+
+	p.pausedCond.L.Lock()
+	p.paused = false
+
+	p.headerCurrentUsed = p.headers
+	p.rawConn = p.selectRawConn(p.headerCurrentUsed)
+
+	if !p.isStarted {
+		p.isStarted = true
+		go p.proxyToLocal(p.ctx)
+	}
+
+	p.pausedCond.Signal()
+	p.pausedCond.L.Unlock()
+}
+
+func (p *ProxyWrapper) Pause() {
+	if p.remoteConn == nil {
+		return
+	}
+
+	log.Tracef("pause proxy reading from: %s", p.remoteConn.RemoteAddr())
+	p.pausedCond.L.Lock()
+	p.paused = true
+	p.pausedCond.L.Unlock()
+}
+
+func (p *ProxyWrapper) RedirectAs(endpoint *net.UDPAddr) {
+	if endpoint == nil || endpoint.IP == nil {
+		log.Errorf("failed to start package redirection, endpoint is nil")
+		return
+	}
+
+	header, err := NewPacketHeaders(p.proxy.localWGListenPort, endpoint)
+	if err != nil {
+		log.Errorf("failed to create packet headers: %s", err)
+		return
+	}
+
+	// Check if required raw connection is available
+	if !header.isIPv4 && p.proxy.rawConnIPv6 == nil {
+		log.Error(errIPv6ConnNotAvailable)
+		return
+	}
+	if header.isIPv4 && p.proxy.rawConnIPv4 == nil {
+		log.Error(errIPv4ConnNotAvailable)
+		return
+	}
+
+	p.pausedCond.L.Lock()
+	p.paused = false
+
+	p.headerCurrentUsed = header
+	p.rawConn = p.selectRawConn(header)
+
+	p.pausedCond.Signal()
+	p.pausedCond.L.Unlock()
+}
+
+// InjectPacket writes b to the remote peer over the underlying transport.
+func (p *ProxyWrapper) InjectPacket(b []byte) error {
+	if p.remoteConn == nil {
+		return errors.New("proxy not started")
+	}
+	if _, err := p.remoteConn.Write(b); err != nil {
+		return err
+	}
+	return nil
+}
+
+// CloseConn close the remoteConn and automatically remove the conn instance from the map
+func (p *ProxyWrapper) CloseConn() error {
+	if p.cancel == nil {
+		return fmt.Errorf("proxy not started")
+	}
+
+	p.cancel()
+
+	p.closeListener.SetCloseListener(nil)
+
+	p.pausedCond.L.Lock()
+	p.paused = false
+	p.pausedCond.Signal()
+	p.pausedCond.L.Unlock()
+
+	if err := p.remoteConn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
+		return fmt.Errorf("failed to close remote conn: %w", err)
+	}
+	return nil
+}
+
+func (p *ProxyWrapper) proxyToLocal(ctx context.Context) {
+	defer p.proxy.removeRelayedConn(p.peerAddr)
+
+	buf := make([]byte, p.proxy.mtu+bufsize.WGBufferOverhead)
+	for {
+		n, err := p.readFromRemote(ctx, buf)
+		if err != nil {
+			return
+		}
+
+		p.pausedCond.L.Lock()
+		for p.paused {
+			p.pausedCond.Wait()
+		}
+
+		err = p.sendPkg(buf[:n], p.headerCurrentUsed)
+		p.pausedCond.L.Unlock()
+
+		if err != nil {
+			if ctx.Err() != nil {
+				return
+			}
+			log.Errorf("failed to write out relayed pkg to local conn: %v", err)
+		}
+	}
+}
+
+func (p *ProxyWrapper) readFromRemote(ctx context.Context, buf []byte) (int, error) {
+	n, err := p.remoteConn.Read(buf)
+	if err != nil {
+		if ctx.Err() != nil {
+			return 0, ctx.Err()
+		}
+		p.closeListener.Notify()
+		if !errors.Is(err, io.EOF) {
+			log.Errorf("failed to read from relayed conn (endpoint: %s): %s", p.wgRelayedEndpointAddr, err)
+		}
+		return 0, err
+	}
+	return n, nil
+}
+
+func (p *ProxyWrapper) sendPkg(data []byte, header *PacketHeaders) error {
+	defer func() {
+		if err := header.layerBuffer.Clear(); err != nil {
+			log.Errorf("failed to clear layer buffer: %s", err)
+		}
+	}()
+
+	payload := gopacket.Payload(data)
+
+	if err := gopacket.SerializeLayers(header.layerBuffer, serializeOpts, header.ipH, header.udpH, payload); err != nil {
+		return fmt.Errorf("serialize layers: %w", err)
+	}
+
+	if _, err := p.rawConn.WriteTo(header.layerBuffer.Bytes(), &net.IPAddr{IP: header.localHostAddr}); err != nil {
+		return fmt.Errorf("write to raw conn: %w", err)
+	}
+	return nil
+}
+
+func (p *ProxyWrapper) selectRawConn(header *PacketHeaders) net.PacketConn {
+	if header.isIPv4 {
+		return p.proxy.rawConnIPv4
+	}
+	return p.proxy.rawConnIPv6
+}