Skip to content

Commit d0cb357

Browse files
committed
Refactor read waiter interface
1 parent f22c2d3 commit d0cb357

4 files changed

Lines changed: 91 additions & 20 deletions

File tree

common/bufio/bind.go

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import (
1010

1111
type BindPacketConn struct {
1212
N.NetPacketConn
13-
Addr net.Addr
13+
addr net.Addr
1414
}
1515

1616
func NewBindPacketConn(conn net.PacketConn, addr net.Addr) *BindPacketConn {
@@ -26,23 +26,36 @@ func (c *BindPacketConn) Read(b []byte) (n int, err error) {
2626
}
2727

2828
func (c *BindPacketConn) Write(b []byte) (n int, err error) {
29-
return c.WriteTo(b, c.Addr)
29+
return c.WriteTo(b, c.addr)
30+
}
31+
32+
func (c *BindPacketConn) CreateReadWaiter() (N.ReadWaiter, bool) {
33+
readWaiter, isReadWaiter := CreatePacketReadWaiter(c.NetPacketConn)
34+
if !isReadWaiter {
35+
return nil, false
36+
}
37+
return &BindPacketReadWaiter{readWaiter}, true
3038
}
3139

3240
func (c *BindPacketConn) RemoteAddr() net.Addr {
33-
return c.Addr
41+
return c.addr
3442
}
3543

3644
func (c *BindPacketConn) Upstream() any {
3745
return c.NetPacketConn
3846
}
3947

48+
var (
49+
_ N.NetPacketConn = (*UnbindPacketConn)(nil)
50+
_ N.PacketReadWaitCreator = (*UnbindPacketConn)(nil)
51+
)
52+
4053
type UnbindPacketConn struct {
4154
N.ExtendedConn
42-
Addr M.Socksaddr
55+
addr M.Socksaddr
4356
}
4457

45-
func NewUnbindPacketConn(conn net.Conn) *UnbindPacketConn {
58+
func NewUnbindPacketConn(conn net.Conn) N.NetPacketConn {
4659
return &UnbindPacketConn{
4760
NewExtendedConn(conn),
4861
M.SocksaddrFromNet(conn.RemoteAddr()),
@@ -52,7 +65,7 @@ func NewUnbindPacketConn(conn net.Conn) *UnbindPacketConn {
5265
func (c *UnbindPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
5366
n, err = c.ExtendedConn.Read(p)
5467
if err == nil {
55-
addr = c.Addr.UDPAddr()
68+
addr = c.addr.UDPAddr()
5669
}
5770
return
5871
}
@@ -66,14 +79,22 @@ func (c *UnbindPacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksad
6679
if err != nil {
6780
return
6881
}
69-
destination = c.Addr
82+
destination = c.addr
7083
return
7184
}
7285

7386
func (c *UnbindPacketConn) WritePacket(buffer *buf.Buffer, _ M.Socksaddr) error {
7487
return c.ExtendedConn.WriteBuffer(buffer)
7588
}
7689

90+
func (c *UnbindPacketConn) CreateReadWaiter() (N.PacketReadWaiter, bool) {
91+
readWaiter, isReadWaiter := CreateReadWaiter(c.ExtendedConn)
92+
if !isReadWaiter {
93+
return nil, false
94+
}
95+
return &UnbindPacketReadWaiter{readWaiter, c.addr}, true
96+
}
97+
7798
func (c *UnbindPacketConn) Upstream() any {
7899
return c.ExtendedConn
79100
}

common/bufio/bind_wait.go

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
package bufio
2+
3+
import (
4+
"github.com/sagernet/sing/common/buf"
5+
M "github.com/sagernet/sing/common/metadata"
6+
N "github.com/sagernet/sing/common/network"
7+
)
8+
9+
var _ N.ReadWaiter = (*BindPacketReadWaiter)(nil)
10+
11+
type BindPacketReadWaiter struct {
12+
readWaiter N.PacketReadWaiter
13+
}
14+
15+
func (w *BindPacketReadWaiter) InitializeReadWaiter(newBuffer func() *buf.Buffer) {
16+
w.readWaiter.InitializeReadWaiter(newBuffer)
17+
}
18+
19+
func (w *BindPacketReadWaiter) WaitReadBuffer() (buffer *buf.Buffer, err error) {
20+
buffer, _, err = w.readWaiter.WaitReadPacket()
21+
return
22+
}
23+
24+
var _ N.PacketReadWaiter = (*UnbindPacketReadWaiter)(nil)
25+
26+
type UnbindPacketReadWaiter struct {
27+
readWaiter N.ReadWaiter
28+
addr M.Socksaddr
29+
}
30+
31+
func (w *UnbindPacketReadWaiter) InitializeReadWaiter(newBuffer func() *buf.Buffer) {
32+
w.readWaiter.InitializeReadWaiter(newBuffer)
33+
}
34+
35+
func (w *UnbindPacketReadWaiter) WaitReadPacket() (buffer *buf.Buffer, destination M.Socksaddr, err error) {
36+
buffer, err = w.readWaiter.WaitReadBuffer()
37+
if err != nil {
38+
return
39+
}
40+
destination = w.addr
41+
return
42+
}

common/bufio/copy_direct_posix.go

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ func copyWaitWithPool(originSource io.Reader, destination N.ExtendedWriter, sour
3939
})
4040
defer source.InitializeReadWaiter(nil)
4141
for {
42-
err = source.WaitReadBuffer()
42+
_, err = source.WaitReadBuffer()
4343
if err != nil {
4444
if errors.Is(err, io.EOF) {
4545
err = nil
@@ -92,7 +92,7 @@ func copyPacketWaitWithPool(originSource N.PacketReader, destinationConn N.Packe
9292
})
9393
defer source.InitializeReadWaiter(nil)
9494
for {
95-
destination, err = source.WaitReadPacket()
95+
_, destination, err = source.WaitReadPacket()
9696
if err != nil {
9797
return
9898
}
@@ -123,6 +123,7 @@ type syscallReadWaiter struct {
123123
rawConn syscall.RawConn
124124
readErr error
125125
readFunc func(fd uintptr) (done bool)
126+
buffer *buf.Buffer
126127
}
127128

128129
func createSyscallReadWaiter(reader any) (*syscallReadWaiter, bool) {
@@ -156,26 +157,29 @@ func (w *syscallReadWaiter) InitializeReadWaiter(newBuffer func() *buf.Buffer) {
156157
if readN == 0 {
157158
w.readErr = io.EOF
158159
}
160+
w.buffer = buffer
159161
return true
160162
}
161163
}
162164
}
163165

164-
func (w *syscallReadWaiter) WaitReadBuffer() error {
166+
func (w *syscallReadWaiter) WaitReadBuffer() (buffer *buf.Buffer, err error) {
165167
if w.readFunc == nil {
166-
return os.ErrInvalid
168+
return nil, os.ErrInvalid
167169
}
168-
err := w.rawConn.Read(w.readFunc)
170+
err = w.rawConn.Read(w.readFunc)
169171
if err != nil {
170-
return err
172+
return
171173
}
172174
if w.readErr != nil {
173175
if w.readErr == io.EOF {
174-
return io.EOF
176+
return nil, io.EOF
175177
}
176-
return E.Cause(w.readErr, "raw read")
178+
return nil, E.Cause(w.readErr, "raw read")
177179
}
178-
return nil
180+
buffer = w.buffer
181+
w.buffer = nil
182+
return
179183
}
180184

181185
var _ N.PacketReadWaiter = (*syscallPacketReadWaiter)(nil)
@@ -185,6 +189,7 @@ type syscallPacketReadWaiter struct {
185189
readErr error
186190
readFrom M.Socksaddr
187191
readFunc func(fd uintptr) (done bool)
192+
buffer *buf.Buffer
188193
}
189194

190195
func createSyscallPacketReadWaiter(reader any) (*syscallPacketReadWaiter, bool) {
@@ -225,14 +230,15 @@ func (w *syscallPacketReadWaiter) InitializeReadWaiter(newBuffer func() *buf.Buf
225230
w.readFrom = M.SocksaddrFrom(netip.AddrFrom16(fromAddr.Addr), uint16(fromAddr.Port)).Unwrap()
226231
}
227232
}
233+
w.buffer = buffer
228234
return true
229235
}
230236
}
231237
}
232238

233-
func (w *syscallPacketReadWaiter) WaitReadPacket() (destination M.Socksaddr, err error) {
239+
func (w *syscallPacketReadWaiter) WaitReadPacket() (buffer *buf.Buffer, destination M.Socksaddr, err error) {
234240
if w.readFunc == nil {
235-
return M.Socksaddr{}, os.ErrInvalid
241+
return nil, M.Socksaddr{}, os.ErrInvalid
236242
}
237243
err = w.rawConn.Read(w.readFunc)
238244
if err != nil {
@@ -242,6 +248,8 @@ func (w *syscallPacketReadWaiter) WaitReadPacket() (destination M.Socksaddr, err
242248
err = E.Cause(w.readErr, "raw read")
243249
return
244250
}
251+
buffer = w.buffer
252+
w.buffer = nil
245253
destination = w.readFrom
246254
return
247255
}

common/network/direct.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import (
77

88
type ReadWaiter interface {
99
InitializeReadWaiter(newBuffer func() *buf.Buffer)
10-
WaitReadBuffer() error
10+
WaitReadBuffer() (buffer *buf.Buffer, err error)
1111
}
1212

1313
type ReadWaitCreator interface {
@@ -16,7 +16,7 @@ type ReadWaitCreator interface {
1616

1717
type PacketReadWaiter interface {
1818
InitializeReadWaiter(newBuffer func() *buf.Buffer)
19-
WaitReadPacket() (destination M.Socksaddr, err error)
19+
WaitReadPacket() (buffer *buf.Buffer, destination M.Socksaddr, err error)
2020
}
2121

2222
type PacketReadWaitCreator interface {

0 commit comments

Comments
 (0)