1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-10-31 00:08:08 -04:00
v2fly/transport/internet/udp/dispatcher_test.go

85 lines
1.9 KiB
Go
Raw Normal View History

2017-02-04 17:40:10 -05:00
package udp_test
import (
"context"
"sync/atomic"
"testing"
"time"
"v2ray.com/core/common/buf"
"v2ray.com/core/common/net"
2019-01-05 15:43:22 -05:00
"v2ray.com/core/common/protocol/udp"
2018-10-12 17:57:56 -04:00
"v2ray.com/core/features/routing"
2018-11-03 07:36:29 -04:00
"v2ray.com/core/transport"
2017-02-04 17:40:10 -05:00
. "v2ray.com/core/transport/internet/udp"
2018-04-16 18:31:10 -04:00
"v2ray.com/core/transport/pipe"
2017-10-24 10:15:35 -04:00
. "v2ray.com/ext/assert"
2017-02-04 17:40:10 -05:00
)
type TestDispatcher struct {
2018-11-03 07:36:29 -04:00
OnDispatch func(ctx context.Context, dest net.Destination) (*transport.Link, error)
2017-02-04 17:40:10 -05:00
}
2018-11-03 07:36:29 -04:00
func (d *TestDispatcher) Dispatch(ctx context.Context, dest net.Destination) (*transport.Link, error) {
2017-02-04 17:40:10 -05:00
return d.OnDispatch(ctx, dest)
}
2018-01-10 07:07:56 -05:00
func (d *TestDispatcher) Start() error {
return nil
}
2018-02-08 09:39:46 -05:00
func (d *TestDispatcher) Close() error {
return nil
}
2018-01-10 07:07:56 -05:00
2018-10-12 17:57:56 -04:00
func (*TestDispatcher) Type() interface{} {
return routing.DispatcherType()
}
2017-02-04 17:40:10 -05:00
func TestSameDestinationDispatching(t *testing.T) {
2017-10-24 10:15:35 -04:00
assert := With(t)
2017-02-04 17:40:10 -05:00
ctx, cancel := context.WithCancel(context.Background())
2018-05-25 06:50:33 -04:00
uplinkReader, uplinkWriter := pipe.New(pipe.WithSizeLimit(1024))
downlinkReader, downlinkWriter := pipe.New(pipe.WithSizeLimit(1024))
2018-04-16 18:31:10 -04:00
2017-02-04 17:40:10 -05:00
go func() {
for {
2018-04-16 18:31:10 -04:00
data, err := uplinkReader.ReadMultiBuffer()
2017-02-04 17:40:10 -05:00
if err != nil {
break
}
2018-04-16 18:31:10 -04:00
err = downlinkWriter.WriteMultiBuffer(data)
2017-10-24 10:15:35 -04:00
assert(err, IsNil)
2017-02-04 17:40:10 -05:00
}
}()
var count uint32
td := &TestDispatcher{
2018-11-03 07:36:29 -04:00
OnDispatch: func(ctx context.Context, dest net.Destination) (*transport.Link, error) {
2017-02-04 17:40:10 -05:00
atomic.AddUint32(&count, 1)
2018-11-03 07:36:29 -04:00
return &transport.Link{Reader: downlinkReader, Writer: uplinkWriter}, nil
2017-02-04 17:40:10 -05:00
},
}
dest := net.UDPDestination(net.LocalHostIP, 53)
2017-02-04 17:40:10 -05:00
b := buf.New()
2018-11-14 16:55:20 -05:00
b.WriteString("abcd")
2018-07-03 15:53:36 -04:00
2017-02-04 17:40:10 -05:00
var msgCount uint32
2019-01-05 15:43:22 -05:00
dispatcher := NewDispatcher(td, func(ctx context.Context, packet *udp.Packet) {
2017-02-04 17:40:10 -05:00
atomic.AddUint32(&msgCount, 1)
})
2018-07-03 15:53:36 -04:00
dispatcher.Dispatch(ctx, dest, b)
2017-02-04 17:40:10 -05:00
for i := 0; i < 5; i++ {
2018-07-03 15:53:36 -04:00
dispatcher.Dispatch(ctx, dest, b)
2017-02-04 17:40:10 -05:00
}
time.Sleep(time.Second)
cancel()
2017-10-24 10:15:35 -04:00
assert(count, Equals, uint32(1))
2018-05-27 09:06:17 -04:00
assert(atomic.LoadUint32(&msgCount), Equals, uint32(6))
2017-02-04 17:40:10 -05:00
}