diff --git a/transport/internet/kcp/receiving.go b/transport/internet/kcp/receiving.go index c55b0d0ca..66caf483a 100644 --- a/transport/internet/kcp/receiving.go +++ b/transport/internet/kcp/receiving.go @@ -65,13 +65,17 @@ type ReceivingQueue struct { timeout time.Time } -func NewReceivingQueue() *ReceivingQueue { +func NewReceivingQueue(size uint32) *ReceivingQueue { return &ReceivingQueue{ - queue: make(chan *alloc.Buffer, effectiveConfig.GetReceivingQueueSize()), + queue: make(chan *alloc.Buffer, size), } } func (this *ReceivingQueue) Read(buf []byte) (int, error) { + if this.closed { + return 0, io.EOF + } + if this.cache.Len() > 0 { nBytes, err := this.cache.Read(buf) if this.cache.IsEmpty() { @@ -230,7 +234,7 @@ func NewReceivingWorker(kcp *KCP) *ReceivingWorker { windowSize := effectiveConfig.GetReceivingWindowSize() worker := &ReceivingWorker{ kcp: kcp, - queue: NewReceivingQueue(), + queue: NewReceivingQueue(effectiveConfig.GetReceivingQueueSize()), window: NewReceivingWindow(windowSize), windowSize: windowSize, } diff --git a/transport/internet/kcp/receiving_test.go b/transport/internet/kcp/receiving_test.go index 1eff3977d..9c44b1216 100644 --- a/transport/internet/kcp/receiving_test.go +++ b/transport/internet/kcp/receiving_test.go @@ -1,8 +1,11 @@ package kcp_test import ( + "io" "testing" + "time" + "github.com/v2ray/v2ray-core/common/alloc" "github.com/v2ray/v2ray-core/testing/assert" . "github.com/v2ray/v2ray-core/transport/internet/kcp" ) @@ -34,3 +37,40 @@ func TestRecivingWindow(t *testing.T) { assert.Pointer(window.Remove(1)).Equals(seg2) assert.Pointer(window.Remove(2)).Equals(seg3) } + +func TestRecivingQueue(t *testing.T) { + assert := assert.On(t) + + queue := NewReceivingQueue(2) + assert.Bool(queue.Put(alloc.NewSmallBuffer().Clear().AppendString("abcd"))).IsTrue() + assert.Bool(queue.Put(alloc.NewSmallBuffer().Clear().AppendString("efg"))).IsTrue() + assert.Bool(queue.Put(alloc.NewSmallBuffer().Clear().AppendString("more content"))).IsFalse() + + b := make([]byte, 1024) + nBytes, err := queue.Read(b) + assert.Error(err).IsNil() + assert.Int(nBytes).Equals(7) + assert.String(string(b[:nBytes])).Equals("abcdefg") + + assert.Bool(queue.Put(alloc.NewSmallBuffer().Clear().AppendString("1"))).IsTrue() + queue.Close() + nBytes, err = queue.Read(b) + assert.Error(err).Equals(io.EOF) +} + +func TestRecivingQueueTimeout(t *testing.T) { + assert := assert.On(t) + + queue := NewReceivingQueue(2) + assert.Bool(queue.Put(alloc.NewSmallBuffer().Clear().AppendString("abcd"))).IsTrue() + queue.SetReadDeadline(time.Now().Add(time.Second)) + + b := make([]byte, 1024) + nBytes, err := queue.Read(b) + assert.Error(err).IsNil() + assert.Int(nBytes).Equals(4) + assert.String(string(b[:nBytes])).Equals("abcd") + + nBytes, err = queue.Read(b) + assert.Error(err).IsNotNil() +}