Hunter0x7c7
2022-08-11 b8230139fb40edea387617b6accd8371e37eda58
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
package udp_test
 
import (
    "context"
    "sync/atomic"
    "testing"
    "time"
 
    "github.com/v2fly/v2ray-core/v5/common"
    "github.com/v2fly/v2ray-core/v5/common/buf"
    "github.com/v2fly/v2ray-core/v5/common/net"
    "github.com/v2fly/v2ray-core/v5/common/protocol/udp"
    "github.com/v2fly/v2ray-core/v5/features/routing"
    "github.com/v2fly/v2ray-core/v5/transport"
    . "github.com/v2fly/v2ray-core/v5/transport/internet/udp"
    "github.com/v2fly/v2ray-core/v5/transport/pipe"
)
 
type TestDispatcher struct {
    OnDispatch func(ctx context.Context, dest net.Destination) (*transport.Link, error)
}
 
func (d *TestDispatcher) Dispatch(ctx context.Context, dest net.Destination) (*transport.Link, error) {
    return d.OnDispatch(ctx, dest)
}
 
func (d *TestDispatcher) Start() error {
    return nil
}
 
func (d *TestDispatcher) Close() error {
    return nil
}
 
func (*TestDispatcher) Type() interface{} {
    return routing.DispatcherType()
}
 
func TestSameDestinationDispatching(t *testing.T) {
    ctx, cancel := context.WithCancel(context.Background())
    uplinkReader, uplinkWriter := pipe.New(pipe.WithSizeLimit(1024))
    downlinkReader, downlinkWriter := pipe.New(pipe.WithSizeLimit(1024))
 
    go func() {
        for {
            data, err := uplinkReader.ReadMultiBuffer()
            if err != nil {
                break
            }
            err = downlinkWriter.WriteMultiBuffer(data)
            common.Must(err)
        }
    }()
 
    var count uint32
    td := &TestDispatcher{
        OnDispatch: func(ctx context.Context, dest net.Destination) (*transport.Link, error) {
            atomic.AddUint32(&count, 1)
            return &transport.Link{Reader: downlinkReader, Writer: uplinkWriter}, nil
        },
    }
    dest := net.UDPDestination(net.LocalHostIP, 53)
 
    b := buf.New()
    b.WriteString("abcd")
 
    var msgCount uint32
    dispatcher := NewSplitDispatcher(td, func(ctx context.Context, packet *udp.Packet) {
        atomic.AddUint32(&msgCount, 1)
    })
 
    dispatcher.Dispatch(ctx, dest, b)
    for i := 0; i < 5; i++ {
        dispatcher.Dispatch(ctx, dest, b)
    }
 
    time.Sleep(time.Second)
    cancel()
 
    if count != 1 {
        t.Error("count: ", count)
    }
    if v := atomic.LoadUint32(&msgCount); v != 6 {
        t.Error("msgCount: ", v)
    }
}