Skip to content

Commit

Permalink
better enhanceUDPConn
Browse files Browse the repository at this point in the history
  • Loading branch information
wwqgtxx committed Dec 19, 2023
1 parent e1f3eb1 commit 8d74452
Show file tree
Hide file tree
Showing 5 changed files with 75 additions and 59 deletions.
10 changes: 5 additions & 5 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,12 @@ go 1.21

require (
github.com/gobwas/pool v0.2.1
github.com/gobwas/ws v1.3.0
github.com/gobwas/ws v1.3.1
github.com/gofrs/uuid/v5 v5.0.0
github.com/sagernet/tfo-go v0.0.0-20230816093905-5a5c285d44a6
golang.org/x/crypto v0.14.0
golang.org/x/net v0.16.0
golang.org/x/sys v0.13.0
github.com/sagernet/tfo-go v0.0.0-20231209031829-7b5343ac1dc6
golang.org/x/crypto v0.17.0
golang.org/x/net v0.19.0
golang.org/x/sys v0.15.0
gopkg.in/yaml.v3 v3.0.1
lukechampine.com/blake3 v1.2.1
)
Expand Down
20 changes: 10 additions & 10 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -2,21 +2,21 @@ github.com/gobwas/httphead v0.1.0 h1:exrUm0f4YX0L7EBwZHuCF4GDp8aJfVeBrlLQrs6NqWU
github.com/gobwas/httphead v0.1.0/go.mod h1:O/RXo79gxV8G+RqlR/otEwx4Q36zl9rqC5u12GKvMCM=
github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og=
github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
github.com/gobwas/ws v1.3.0 h1:sbeU3Y4Qzlb+MOzIe6mQGf7QR4Hkv6ZD0qhGkBFL2O0=
github.com/gobwas/ws v1.3.0/go.mod h1:hRKAFb8wOxFROYNsT1bqfWnhX+b5MFeJM9r2ZSwg/KY=
github.com/gobwas/ws v1.3.1 h1:Qi34dfLMWJbiKaNbDVzM9x27nZBjmkaW6i4+Ku+pGVU=
github.com/gobwas/ws v1.3.1/go.mod h1:hRKAFb8wOxFROYNsT1bqfWnhX+b5MFeJM9r2ZSwg/KY=
github.com/gofrs/uuid/v5 v5.0.0 h1:p544++a97kEL+svbcFbCQVM9KFu0Yo25UoISXGNNH9M=
github.com/gofrs/uuid/v5 v5.0.0/go.mod h1:CDOjlDMVAtN56jqyRUZh58JT31Tiw7/oQyEXZV+9bD8=
github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4=
github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
github.com/sagernet/tfo-go v0.0.0-20230816093905-5a5c285d44a6 h1:Px+hN4Vzgx+iCGVnWH5A8eR7JhNnIV3rGQmBxA7cw6Q=
github.com/sagernet/tfo-go v0.0.0-20230816093905-5a5c285d44a6/go.mod h1:zovq6vTvEM6ECiqE3Eeb9rpIylPpamPcmrJ9tv0Bt0M=
golang.org/x/crypto v0.14.0 h1:wBqGXzWJW6m1XrIKlAH0Hs1JJ7+9KBwnIO8v66Q9cHc=
golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4=
golang.org/x/net v0.16.0 h1:7eBu7KsSvFDtSXUIDbh3aqlK4DPsZ1rByC8PFfBThos=
golang.org/x/net v0.16.0/go.mod h1:NxSsAGuq816PNPmqtQdLE42eU2Fs7NoRIZrHJAlaCOE=
github.com/sagernet/tfo-go v0.0.0-20231209031829-7b5343ac1dc6 h1:z3SJQhVyU63FT26Wn/UByW6b7q8QKB0ZkPqsyqcz2PI=
github.com/sagernet/tfo-go v0.0.0-20231209031829-7b5343ac1dc6/go.mod h1:73xRZuxwkFk4aiLw28hG8W6o9cr2UPrGL9pdY2UTbvY=
golang.org/x/crypto v0.17.0 h1:r8bRNjWL3GshPW3gkd+RpvzWrZAwPS49OmTGZ/uhM4k=
golang.org/x/crypto v0.17.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4=
golang.org/x/net v0.19.0 h1:zTwKpTd2XuCqf8huc7Fo2iSy+4RHPd10s4KzeTnVr1c=
golang.org/x/net v0.19.0/go.mod h1:CfAk/cbD4CthTvqiEl8NpboMuiuOYsAr/7NOjZJtv1U=
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc=
golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
Expand Down
2 changes: 1 addition & 1 deletion udp/packet.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ type EnhancePacketConn interface {

func NewEnhancePacketConn(pc net.PacketConn) EnhancePacketConn {
if udpConn, isUDPConn := pc.(*net.UDPConn); isUDPConn {
return &enhanceUDPConn{UDPConn: udpConn}
return newEnhancePacketConn(udpConn)
}
if enhancePC, isEnhancePC := pc.(EnhancePacketConn); isEnhancePC {
return enhancePC
Expand Down
98 changes: 55 additions & 43 deletions udp/packet_posix.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,57 +11,69 @@ import (

type enhanceUDPConn struct {
*net.UDPConn
rawConn syscall.RawConn
rawConn syscall.RawConn
data []byte
put func()
addr netip.AddrPort
readErr error
boundRead func(fd uintptr) (done bool)
}

func (c *enhanceUDPConn) WaitReadFrom() (data []byte, put func(), addr netip.AddrPort, err error) {
if c.rawConn == nil {
c.rawConn, _ = c.UDPConn.SyscallConn()
func newEnhancePacketConn(udpConn *net.UDPConn) EnhancePacketConn {
c := &enhanceUDPConn{UDPConn: udpConn}
c.rawConn, _ = udpConn.SyscallConn()
c.boundRead = c.read
return c
}

func (c *enhanceUDPConn) read(fd uintptr) (done bool) {
readBuf := BufPool.Get().([]byte)
c.put = func() {
BufPool.Put(readBuf)
}
var ip netip.Addr
var port int
var readErr error
err = c.rawConn.Read(func(fd uintptr) (done bool) {
readBuf := BufPool.Get().([]byte)
put = func() {
BufPool.Put(readBuf)
}
var readFrom syscall.Sockaddr
var readN int
readN, _, _, readFrom, readErr = syscall.Recvmsg(int(fd), readBuf, nil, 0)
if readN > 0 {
data = readBuf[:readN]
} else {
put()
put = nil
data = nil
}
if readErr == syscall.EAGAIN {
return false
}
if readFrom != nil {
switch from := readFrom.(type) {
case *syscall.SockaddrInet4:
ip = netip.AddrFrom4(from.Addr)
port = from.Port
case *syscall.SockaddrInet6:
ip = netip.AddrFrom16(from.Addr).WithZone(strconv.FormatInt(int64(from.ZoneId), 10))
port = from.Port
}
var readFrom syscall.Sockaddr
var readN int
readN, _, _, readFrom, c.readErr = syscall.Recvmsg(int(fd), readBuf, nil, 0)
if readN > 0 {
c.data = readBuf[:readN]
} else {
c.put()
c.put = nil
c.data = nil
}
if c.readErr == syscall.EAGAIN {
return false
}
if readFrom != nil {
switch from := readFrom.(type) {
case *syscall.SockaddrInet4:
ip := netip.AddrFrom4(from.Addr)
port := from.Port
c.addr = netip.AddrPortFrom(ip, uint16(port))
case *syscall.SockaddrInet6:
ip := netip.AddrFrom16(from.Addr).WithZone(strconv.FormatInt(int64(from.ZoneId), 10))
port := from.Port
c.addr = netip.AddrPortFrom(ip, uint16(port))
}
// udp should not convert readN == 0 to io.EOF
//if readN == 0 {
// readErr = io.EOF
//}
return true
})
}
// udp should not convert readN == 0 to io.EOF
//if readN == 0 {
// c.readErr = io.EOF
//}
return true
}

func (c *enhanceUDPConn) WaitReadFrom() (data []byte, put func(), addr netip.AddrPort, err error) {
err = c.rawConn.Read(c.boundRead)
if err != nil {
return
}
if readErr != nil {
err = readErr
if c.readErr != nil {
err = c.readErr
return
}
addr = netip.AddrPortFrom(ip, uint16(port))
data = c.data
put = c.put
addr = c.addr
return
}
4 changes: 4 additions & 0 deletions udp/packet_windows.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,10 @@ type enhanceUDPConn struct {
*net.UDPConn
}

func newEnhancePacketConn(udpConn *net.UDPConn) EnhancePacketConn {
return &enhanceUDPConn{UDPConn: udpConn}
}

func (c *enhanceUDPConn) WaitReadFrom() (data []byte, put func(), addr netip.AddrPort, err error) {
readBuf := BufPool.Get().([]byte)
put = func() {
Expand Down

0 comments on commit 8d74452

Please sign in to comment.