From c6832bb477d5116d79786a92763f423bd4298725 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= <zoltan.pmail@gmail.com>
Date: Mon, 22 Apr 2024 18:56:16 +0200
Subject: [PATCH] Add POC code

---
 conn/bind_std.go         | 31 +++++++++++++++++++++++++++++--
 conn/conn.go             |  2 ++
 conn/receiver_creator.go |  1 +
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/conn/bind_std.go b/conn/bind_std.go
index 7381069..0b17804 100644
--- a/conn/bind_std.go
+++ b/conn/bind_std.go
@@ -84,6 +84,7 @@ func NewStdNetBind() Bind {
 type StdNetEndpoint struct {
 	// AddrPort is the endpoint destination.
 	netip.AddrPort
+	Conn net.PacketConn
 	// src is the current sticky source address and interface index, if
 	// supported. Typically this is a PKTINFO structure from/for control
 	// messages, see unix.PKTINFO for an example.
@@ -127,6 +128,10 @@ func (e *StdNetEndpoint) DstToString() string {
 	return e.AddrPort.String()
 }
 
+func (e *StdNetEndpoint) GetConn() net.PacketConn {
+	return e.Conn
+}
+
 func listenNet(network string, port int) (*net.UDPConn, int, error) {
 	conn, err := listenConfig().ListenPacket(context.Background(), network, ":"+strconv.Itoa(port))
 	if err != nil {
@@ -190,6 +195,10 @@ again:
 		if s.receiverCreator != nil {
 			// Todo: check if this still works
 			fns = append(fns, s.receiverCreator.CreateIPv4ReceiverFn(&s.msgsPool, v4pc, v4conn))
+			turnFn := s.receiverCreator.CreateRelayReceiverFn(&s.msgsPool)
+			if turnFn != nil {
+				fns = append(fns, s.receiverCreator.CreateRelayReceiverFn(&s.msgsPool))
+			}
 		} else {
 			fns = append(fns, s.makeReceiveIPv4(v4pc, v4conn, s.ipv4RxOffload))
 		}
@@ -395,7 +404,11 @@ func (s *StdNetBind) Send(bufs [][]byte, endpoint Endpoint) error {
 retry:
 	if offload {
 		n := coalesceMessages(ua, endpoint.(*StdNetEndpoint), bufs, *msgs, setGSOSize)
-		err = s.send(conn, br, (*msgs)[:n])
+		if endpoint.GetConn() != nil {
+			err = s.sendPacketConn(endpoint.GetConn(), (*msgs)[:n])
+		} else {
+			err = s.send(conn, br, (*msgs)[:n])
+		}
 		if err != nil && offload && errShouldDisableUDPGSO(err) {
 			offload = false
 			s.mu.Lock()
@@ -414,7 +427,11 @@ retry:
 			(*msgs)[i].Buffers[0] = bufs[i]
 			setSrcControl(&(*msgs)[i].OOB, endpoint.(*StdNetEndpoint))
 		}
-		err = s.send(conn, br, (*msgs)[:len(bufs)])
+		if endpoint.GetConn() != nil {
+			err = s.sendPacketConn(endpoint.GetConn(), (*msgs)[:len(bufs)])
+		} else {
+			err = s.send(conn, br, (*msgs)[:len(bufs)])
+		}
 	}
 	if retried {
 		return ErrUDPGSODisabled{onLaddr: conn.LocalAddr().String(), RetryErr: err}
@@ -422,6 +439,16 @@ retry:
 	return err
 }
 
+func (s *StdNetBind) sendPacketConn(conn net.PacketConn, msgs []ipv6.Message) error {
+	for _, msg := range msgs {
+		_, err := conn.WriteTo(msg.Buffers[0], msg.Addr.(*net.UDPAddr))
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (s *StdNetBind) send(conn *net.UDPConn, pc batchWriter, msgs []ipv6.Message) error {
 	var (
 		n     int
diff --git a/conn/conn.go b/conn/conn.go
index a1f57d2..5c2a757 100644
--- a/conn/conn.go
+++ b/conn/conn.go
@@ -9,6 +9,7 @@ package conn
 import (
 	"errors"
 	"fmt"
+	"net"
 	"net/netip"
 	"reflect"
 	"runtime"
@@ -82,6 +83,7 @@ type Endpoint interface {
 	DstToBytes() []byte  // used for mac2 cookie calculations
 	DstIP() netip.Addr
 	SrcIP() netip.Addr
+	GetConn() net.PacketConn
 }
 
 var (
diff --git a/conn/receiver_creator.go b/conn/receiver_creator.go
index fd99148..7c4e084 100644
--- a/conn/receiver_creator.go
+++ b/conn/receiver_creator.go
@@ -9,4 +9,5 @@ import (
 
 type ReceiverCreator interface {
 	CreateIPv4ReceiverFn(msgPool *sync.Pool, pc *ipv4.PacketConn, conn *net.UDPConn) ReceiveFunc
+	CreateRelayReceiverFn(msgPool *sync.Pool) ReceiveFunc
 }