Skip to content

Commit 57b8a4c

Browse files
nekohasekainekohasekai
committed
Add test for copy waiter
1 parent b7a631f commit 57b8a4c

3 files changed

Lines changed: 297 additions & 4 deletions

File tree

‎common/bufio/copy_direct_test.go‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
package bufio
2+
3+
import (
4+
"net"
5+
"testing"
6+
7+
"github.com/sagernet/sing/common/buf"
8+
N "github.com/sagernet/sing/common/network"
9+
10+
"github.com/stretchr/testify/require"
11+
)
12+
13+
func TestCopyWaitTCP(t *testing.T) {
14+
t.Parallel()
15+
inputConn, outputConn := TCPPipe(t)
16+
readWaiter, created := CreateReadWaiter(outputConn)
17+
require.True(t, created)
18+
require.NotNil(t, readWaiter)
19+
readWaiter.InitializeReadWaiter(N.ReadWaitOptions{})
20+
require.NoError(t, TCPTest(t, inputConn, &readWaitWrapper{
21+
Conn: outputConn,
22+
readWaiter: readWaiter,
23+
}))
24+
}
25+
26+
type readWaitWrapper struct {
27+
net.Conn
28+
readWaiter N.ReadWaiter
29+
buffer *buf.Buffer
30+
}
31+
32+
func (r *readWaitWrapper) Read(p []byte) (n int, err error) {
33+
if r.buffer != nil {
34+
if r.buffer.Len() > 0 {
35+
return r.buffer.Read(p)
36+
}
37+
if r.buffer.IsEmpty() {
38+
r.buffer.Release()
39+
r.buffer = nil
40+
}
41+
}
42+
buffer, err := r.readWaiter.WaitReadBuffer()
43+
if err != nil {
44+
return
45+
}
46+
r.buffer = buffer
47+
return r.buffer.Read(p)
48+
}
49+
50+
func TestCopyWaitUDP(t *testing.T) {
51+
t.Parallel()
52+
inputConn, outputConn, outputAddr := UDPPipe(t)
53+
readWaiter, created := CreatePacketReadWaiter(NewPacketConn(outputConn))
54+
require.True(t, created)
55+
require.NotNil(t, readWaiter)
56+
readWaiter.InitializeReadWaiter(N.ReadWaitOptions{})
57+
require.NoError(t, UDPTest(t, inputConn, &packetReadWaitWrapper{
58+
PacketConn: outputConn,
59+
readWaiter: readWaiter,
60+
}, outputAddr))
61+
}
62+
63+
type packetReadWaitWrapper struct {
64+
net.PacketConn
65+
readWaiter N.PacketReadWaiter
66+
}
67+
68+
func (r *packetReadWaitWrapper) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
69+
buffer, destination, err := r.readWaiter.WaitReadPacket()
70+
if err != nil {
71+
return
72+
}
73+
n = copy(p, buffer.Bytes())
74+
buffer.Release()
75+
addr = destination.UDPAddr()
76+
return
77+
}

‎common/bufio/net_test.go‎

Lines changed: 219 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,19 @@ package bufio
22

33
import (
44
"context"
5+
"crypto/md5"
6+
"crypto/rand"
7+
"errors"
8+
"io"
59
"net"
10+
"sync"
611
"testing"
712
"time"
813

914
M "github.com/sagernet/sing/common/metadata"
1015
"github.com/sagernet/sing/common/task"
1116

17+
"github.com/stretchr/testify/assert"
1218
"github.com/stretchr/testify/require"
1319
)
1420

@@ -33,6 +39,10 @@ func TCPPipe(t *testing.T) (net.Conn, net.Conn) {
3339
err = group.Run()
3440
require.NoError(t, err)
3541
listener.Close()
42+
t.Cleanup(func() {
43+
serverConn.Close()
44+
clientConn.Close()
45+
})
3646
return serverConn, clientConn
3747
}
3848

@@ -56,3 +66,212 @@ func Timeout(t *testing.T) context.CancelFunc {
5666
}()
5767
return cancel
5868
}
69+
70+
type hashPair struct {
71+
sendHash map[int][]byte
72+
recvHash map[int][]byte
73+
}
74+
75+
func newLargeDataPair() (chan hashPair, chan hashPair, func(t *testing.T) error) {
76+
pingCh := make(chan hashPair)
77+
pongCh := make(chan hashPair)
78+
test := func(t *testing.T) error {
79+
defer close(pingCh)
80+
defer close(pongCh)
81+
pingOpen := false
82+
pongOpen := false
83+
var serverPair hashPair
84+
var clientPair hashPair
85+
86+
for {
87+
if pingOpen && pongOpen {
88+
break
89+
}
90+
91+
select {
92+
case serverPair, pingOpen = <-pingCh:
93+
assert.True(t, pingOpen)
94+
case clientPair, pongOpen = <-pongCh:
95+
assert.True(t, pongOpen)
96+
case <-time.After(10 * time.Second):
97+
return errors.New("timeout")
98+
}
99+
}
100+
101+
assert.Equal(t, serverPair.recvHash, clientPair.sendHash)
102+
assert.Equal(t, serverPair.sendHash, clientPair.recvHash)
103+
104+
return nil
105+
}
106+
107+
return pingCh, pongCh, test
108+
}
109+
110+
func TCPTest(t *testing.T, inputConn net.Conn, outputConn net.Conn) error {
111+
times := 100
112+
chunkSize := int64(64 * 1024)
113+
114+
pingCh, pongCh, test := newLargeDataPair()
115+
writeRandData := func(conn net.Conn) (map[int][]byte, error) {
116+
buf := make([]byte, chunkSize)
117+
hashMap := map[int][]byte{}
118+
for i := 0; i < times; i++ {
119+
if _, err := rand.Read(buf[1:]); err != nil {
120+
return nil, err
121+
}
122+
buf[0] = byte(i)
123+
124+
hash := md5.Sum(buf)
125+
hashMap[i] = hash[:]
126+
127+
if _, err := conn.Write(buf); err != nil {
128+
return nil, err
129+
}
130+
}
131+
132+
return hashMap, nil
133+
}
134+
go func() {
135+
hashMap := map[int][]byte{}
136+
buf := make([]byte, chunkSize)
137+
138+
for i := 0; i < times; i++ {
139+
_, err := io.ReadFull(outputConn, buf)
140+
if err != nil {
141+
t.Log(err.Error())
142+
return
143+
}
144+
145+
hash := md5.Sum(buf)
146+
hashMap[int(buf[0])] = hash[:]
147+
}
148+
149+
sendHash, err := writeRandData(outputConn)
150+
if err != nil {
151+
t.Log(err.Error())
152+
return
153+
}
154+
155+
pingCh <- hashPair{
156+
sendHash: sendHash,
157+
recvHash: hashMap,
158+
}
159+
}()
160+
161+
go func() {
162+
sendHash, err := writeRandData(inputConn)
163+
if err != nil {
164+
t.Log(err.Error())
165+
return
166+
}
167+
168+
hashMap := map[int][]byte{}
169+
buf := make([]byte, chunkSize)
170+
171+
for i := 0; i < times; i++ {
172+
_, err = io.ReadFull(inputConn, buf)
173+
if err != nil {
174+
t.Log(err.Error())
175+
return
176+
}
177+
178+
hash := md5.Sum(buf)
179+
hashMap[int(buf[0])] = hash[:]
180+
}
181+
182+
pongCh <- hashPair{
183+
sendHash: sendHash,
184+
recvHash: hashMap,
185+
}
186+
}()
187+
return test(t)
188+
}
189+
190+
func UDPTest(t *testing.T, inputConn net.PacketConn, outputConn net.PacketConn, outputAddr M.Socksaddr) error {
191+
rAddr := outputAddr.UDPAddr()
192+
times := 50
193+
chunkSize := 9000
194+
pingCh, pongCh, test := newLargeDataPair()
195+
writeRandData := func(pc net.PacketConn, addr net.Addr) (map[int][]byte, error) {
196+
hashMap := map[int][]byte{}
197+
mux := sync.Mutex{}
198+
for i := 0; i < times; i++ {
199+
buf := make([]byte, chunkSize)
200+
if _, err := rand.Read(buf[1:]); err != nil {
201+
t.Log(err.Error())
202+
continue
203+
}
204+
buf[0] = byte(i)
205+
206+
hash := md5.Sum(buf)
207+
mux.Lock()
208+
hashMap[i] = hash[:]
209+
mux.Unlock()
210+
211+
if _, err := pc.WriteTo(buf, addr); err != nil {
212+
t.Log(err.Error())
213+
}
214+
215+
time.Sleep(10 * time.Millisecond)
216+
}
217+
218+
return hashMap, nil
219+
}
220+
go func() {
221+
var (
222+
lAddr net.Addr
223+
err error
224+
)
225+
hashMap := map[int][]byte{}
226+
buf := make([]byte, 64*1024)
227+
228+
for i := 0; i < times; i++ {
229+
_, lAddr, err = outputConn.ReadFrom(buf)
230+
if err != nil {
231+
t.Log(err.Error())
232+
return
233+
}
234+
hash := md5.Sum(buf[:chunkSize])
235+
hashMap[int(buf[0])] = hash[:]
236+
}
237+
sendHash, err := writeRandData(outputConn, lAddr)
238+
if err != nil {
239+
t.Log(err.Error())
240+
return
241+
}
242+
243+
pingCh <- hashPair{
244+
sendHash: sendHash,
245+
recvHash: hashMap,
246+
}
247+
}()
248+
249+
go func() {
250+
sendHash, err := writeRandData(inputConn, rAddr)
251+
if err != nil {
252+
t.Log(err.Error())
253+
return
254+
}
255+
256+
hashMap := map[int][]byte{}
257+
buf := make([]byte, 64*1024)
258+
259+
for i := 0; i < times; i++ {
260+
_, _, err := inputConn.ReadFrom(buf)
261+
if err != nil {
262+
t.Log(err.Error())
263+
return
264+
}
265+
266+
hash := md5.Sum(buf[:chunkSize])
267+
hashMap[int(buf[0])] = hash[:]
268+
}
269+
270+
pongCh <- hashPair{
271+
sendHash: sendHash,
272+
recvHash: hashMap,
273+
}
274+
}()
275+
276+
return test(t)
277+
}

‎common/bufio/vectorised_test.go‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,6 @@ import (
1111
func TestWriteVectorised(t *testing.T) {
1212
t.Parallel()
1313
inputConn, outputConn := TCPPipe(t)
14-
defer inputConn.Close()
15-
defer outputConn.Close()
1614
vectorisedWriter, created := CreateVectorisedWriter(inputConn)
1715
require.True(t, created)
1816
require.NotNil(t, vectorisedWriter)
@@ -36,9 +34,8 @@ func TestWriteVectorised(t *testing.T) {
3634
}
3735

3836
func TestWriteVectorisedPacket(t *testing.T) {
37+
t.Parallel()
3938
inputConn, outputConn, outputAddr := UDPPipe(t)
40-
defer inputConn.Close()
41-
defer outputConn.Close()
4239
vectorisedWriter, created := CreateVectorisedPacketWriter(inputConn)
4340
require.True(t, created)
4441
require.NotNil(t, vectorisedWriter)

0 commit comments

Comments
 (0)