diff --git a/conn/bind_std.go b/conn/bind_std.go index 69789b3..6d004d3 100644 --- a/conn/bind_std.go +++ b/conn/bind_std.go @@ -42,6 +42,14 @@ type StdNetBind struct { blackhole4 bool blackhole6 bool + + receiverCreator ReceiverCreator +} + +func NewStdNetBindWithReceiverCreator(receiverCreator ReceiverCreator) *StdNetBind { + b, _ := NewStdNetBind().(*StdNetBind) + b.receiverCreator = receiverCreator + return b } func NewStdNetBind() Bind { @@ -59,7 +67,7 @@ func NewStdNetBind() Bind { msgs := make([]ipv4.Message, IdealBatchSize) for i := range msgs { msgs[i].Buffers = make(net.Buffers, 1) - msgs[i].OOB = make([]byte, srcControlSize) + msgs[i].OOB = make([]byte, SrcControlSize) } return &msgs }, @@ -70,7 +78,7 @@ func NewStdNetBind() Bind { msgs := make([]ipv6.Message, IdealBatchSize) for i := range msgs { msgs[i].Buffers = make(net.Buffers, 1) - msgs[i].OOB = make([]byte, srcControlSize) + msgs[i].OOB = make([]byte, SrcControlSize) } return &msgs }, @@ -192,7 +200,11 @@ again: v4pc = ipv4.NewPacketConn(v4conn) s.ipv4PC = v4pc } - fns = append(fns, s.makeReceiveIPv4(v4pc, v4conn)) + if s.receiverCreator != nil { + fns = append(fns, s.receiverCreator.CreateIPv4ReceiverFn(&s.ipv4MsgsPool, v4pc, v4conn)) + } else { + fns = append(fns, s.makeReceiveIPv4(v4pc, v4conn)) + } s.ipv4 = v4conn } if v6conn != nil { @@ -236,7 +248,7 @@ func (s *StdNetBind) makeReceiveIPv4(pc *ipv4.PacketConn, conn *net.UDPConn) Rec sizes[i] = msg.N addrPort := msg.Addr.(*net.UDPAddr).AddrPort() ep := &StdNetEndpoint{AddrPort: addrPort} // TODO: remove allocation - getSrcFromControl(msg.OOB[:msg.NN], ep) + GetSrcFromControl(msg.OOB[:msg.NN], ep) eps[i] = ep } return numMsgs, nil @@ -269,7 +281,7 @@ func (s *StdNetBind) makeReceiveIPv6(pc *ipv6.PacketConn, conn *net.UDPConn) Rec sizes[i] = msg.N addrPort := msg.Addr.(*net.UDPAddr).AddrPort() ep := &StdNetEndpoint{AddrPort: addrPort} // TODO: remove allocation - getSrcFromControl(msg.OOB[:msg.NN], ep) + GetSrcFromControl(msg.OOB[:msg.NN], ep) eps[i] = ep } return numMsgs, nil diff --git a/conn/receiver_creator.go b/conn/receiver_creator.go new file mode 100644 index 0000000..fd99148 --- /dev/null +++ b/conn/receiver_creator.go @@ -0,0 +1,12 @@ +package conn + +import ( + "net" + "sync" + + "golang.org/x/net/ipv4" +) + +type ReceiverCreator interface { + CreateIPv4ReceiverFn(msgPool *sync.Pool, pc *ipv4.PacketConn, conn *net.UDPConn) ReceiveFunc +} diff --git a/conn/sticky_default.go b/conn/sticky_default.go index 05f00ea..c1651d2 100644 --- a/conn/sticky_default.go +++ b/conn/sticky_default.go @@ -10,9 +10,9 @@ package conn // TODO: macOS, FreeBSD and other BSDs likely do support this feature set, but // use alternatively named flags and need ports and require testing. -// getSrcFromControl parses the control for PKTINFO and if found updates ep with +// GetSrcFromControl parses the control for PKTINFO and if found updates ep with // the source information found. -func getSrcFromControl(control []byte, ep *StdNetEndpoint) { +func GetSrcFromControl(control []byte, ep *StdNetEndpoint) { } // setSrcControl parses the control for PKTINFO and if found updates ep with @@ -20,8 +20,8 @@ func getSrcFromControl(control []byte, ep *StdNetEndpoint) { func setSrcControl(control *[]byte, ep *StdNetEndpoint) { } -// srcControlSize returns the recommended buffer size for pooling sticky control +// SrcControlSize returns the recommended buffer size for pooling sticky control // data. -const srcControlSize = 0 +const SrcControlSize = 0 const StdNetSupportsStickySockets = false diff --git a/conn/sticky_linux.go b/conn/sticky_linux.go index 274fa38..75b690e 100644 --- a/conn/sticky_linux.go +++ b/conn/sticky_linux.go @@ -14,9 +14,9 @@ import ( "golang.org/x/sys/unix" ) -// getSrcFromControl parses the control for PKTINFO and if found updates ep with +// GetSrcFromControl parses the control for PKTINFO and if found updates ep with // the source information found. -func getSrcFromControl(control []byte, ep *StdNetEndpoint) { +func GetSrcFromControl(control []byte, ep *StdNetEndpoint) { ep.ClearSrc() var ( @@ -80,7 +80,7 @@ func setSrcControl(control *[]byte, ep *StdNetEndpoint) { return } - if len(*control) < srcControlSize { + if len(*control) < SrcControlSize { *control = (*control)[:0] return } @@ -112,6 +112,6 @@ func setSrcControl(control *[]byte, ep *StdNetEndpoint) { } -var srcControlSize = unix.CmsgSpace(unix.SizeofInet6Pktinfo) +var SrcControlSize = unix.CmsgSpace(unix.SizeofInet6Pktinfo) const StdNetSupportsStickySockets = true diff --git a/conn/sticky_linux_test.go b/conn/sticky_linux_test.go index 0219ac3..ba36140 100644 --- a/conn/sticky_linux_test.go +++ b/conn/sticky_linux_test.go @@ -26,7 +26,7 @@ func Test_setSrcControl(t *testing.T) { ep.src.Addr = netip.MustParseAddr("127.0.0.1") ep.src.ifidx = 5 - control := make([]byte, srcControlSize) + control := make([]byte, SrcControlSize) setSrcControl(&control, ep) @@ -56,7 +56,7 @@ func Test_setSrcControl(t *testing.T) { ep.src.Addr = netip.MustParseAddr("::1") ep.src.ifidx = 5 - control := make([]byte, srcControlSize) + control := make([]byte, SrcControlSize) setSrcControl(&control, ep) @@ -80,7 +80,7 @@ func Test_setSrcControl(t *testing.T) { }) t.Run("ClearOnNoSrc", func(t *testing.T) { - control := make([]byte, srcControlSize) + control := make([]byte, SrcControlSize) hdr := (*unix.Cmsghdr)(unsafe.Pointer(&control[0])) hdr.Level = 1 hdr.Type = 2 @@ -96,7 +96,7 @@ func Test_setSrcControl(t *testing.T) { func Test_getSrcFromControl(t *testing.T) { t.Run("IPv4", func(t *testing.T) { - control := make([]byte, srcControlSize) + control := make([]byte, SrcControlSize) hdr := (*unix.Cmsghdr)(unsafe.Pointer(&control[0])) hdr.Level = unix.IPPROTO_IP hdr.Type = unix.IP_PKTINFO @@ -106,7 +106,7 @@ func Test_getSrcFromControl(t *testing.T) { info.Ifindex = 5 ep := &StdNetEndpoint{} - getSrcFromControl(control, ep) + GetSrcFromControl(control, ep) if ep.src.Addr != netip.MustParseAddr("127.0.0.1") { t.Errorf("unexpected address: %v", ep.src.Addr) @@ -116,7 +116,7 @@ func Test_getSrcFromControl(t *testing.T) { } }) t.Run("IPv6", func(t *testing.T) { - control := make([]byte, srcControlSize) + control := make([]byte, SrcControlSize) hdr := (*unix.Cmsghdr)(unsafe.Pointer(&control[0])) hdr.Level = unix.IPPROTO_IPV6 hdr.Type = unix.IPV6_PKTINFO @@ -126,7 +126,7 @@ func Test_getSrcFromControl(t *testing.T) { info.Ifindex = 5 ep := &StdNetEndpoint{} - getSrcFromControl(control, ep) + GetSrcFromControl(control, ep) if ep.SrcIP() != netip.MustParseAddr("::1") { t.Errorf("unexpected address: %v", ep.SrcIP()) @@ -136,12 +136,12 @@ func Test_getSrcFromControl(t *testing.T) { } }) t.Run("ClearOnEmpty", func(t *testing.T) { - control := make([]byte, srcControlSize) + control := make([]byte, SrcControlSize) ep := &StdNetEndpoint{} ep.src.Addr = netip.MustParseAddr("::1") ep.src.ifidx = 5 - getSrcFromControl(control, ep) + GetSrcFromControl(control, ep) if ep.SrcIP().IsValid() { t.Errorf("unexpected address: %v", ep.src.Addr) } @@ -154,7 +154,7 @@ func Test_getSrcFromControl(t *testing.T) { zeroHdr := (*unix.Cmsghdr)(unsafe.Pointer(&zeroControl[0])) zeroHdr.SetLen(unix.CmsgLen(0)) - control := make([]byte, srcControlSize) + control := make([]byte, SrcControlSize) hdr := (*unix.Cmsghdr)(unsafe.Pointer(&control[0])) hdr.Level = unix.IPPROTO_IP hdr.Type = unix.IP_PKTINFO @@ -168,7 +168,7 @@ func Test_getSrcFromControl(t *testing.T) { combined = append(combined, control...) ep := &StdNetEndpoint{} - getSrcFromControl(combined, ep) + GetSrcFromControl(combined, ep) if ep.src.Addr != netip.MustParseAddr("127.0.0.1") { t.Errorf("unexpected address: %v", ep.src.Addr)