1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-12-23 02:26:49 -05:00
v2fly/transport/internet/udp/dispatcher_test.go

87 lines
2.0 KiB
Go
Raw Normal View History

2017-02-04 17:40:10 -05:00
package udp_test
import (
"context"
"sync/atomic"
"testing"
"time"
2021-02-16 15:31:50 -05:00
"github.com/v2fly/v2ray-core/v4/common"
"github.com/v2fly/v2ray-core/v4/common/buf"
"github.com/v2fly/v2ray-core/v4/common/net"
"github.com/v2fly/v2ray-core/v4/common/protocol/udp"
"github.com/v2fly/v2ray-core/v4/features/routing"
"github.com/v2fly/v2ray-core/v4/transport"
. "github.com/v2fly/v2ray-core/v4/transport/internet/udp"
"github.com/v2fly/v2ray-core/v4/transport/pipe"
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) {
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)
2019-02-02 16:19:30 -05:00
common.Must(err)
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()
2019-02-03 13:46:53 -05:00
if count != 1 {
t.Error("count: ", count)
}
if v := atomic.LoadUint32(&msgCount); v != 6 {
t.Error("msgCount: ", v)
}
2017-02-04 17:40:10 -05:00
}