v2ray-core/transport/ray/direct.go

126 lines
1.8 KiB
Go
Raw Normal View History

2015-10-14 13:07:13 +00:00
package ray
import (
2016-04-18 16:44:10 +00:00
"io"
"sync"
2016-05-09 02:04:51 +00:00
"time"
2016-04-18 16:44:10 +00:00
2016-12-09 10:35:27 +00:00
"v2ray.com/core/common/buf"
2015-10-14 13:07:13 +00:00
)
const (
2016-12-02 13:22:46 +00:00
bufferSize = 512
2015-10-14 13:07:13 +00:00
)
2015-10-15 11:15:59 +00:00
// NewRay creates a new Ray for direct traffic transport.
2015-10-14 13:07:13 +00:00
func NewRay() Ray {
return &directRay{
2016-04-18 16:44:10 +00:00
Input: NewStream(),
Output: NewStream(),
2015-10-14 13:07:13 +00:00
}
}
type directRay struct {
2016-04-18 16:44:10 +00:00
Input *Stream
Output *Stream
2015-10-14 13:07:13 +00:00
}
2016-11-27 20:39:09 +00:00
func (v *directRay) OutboundInput() InputStream {
return v.Input
2015-10-14 13:07:13 +00:00
}
2016-11-27 20:39:09 +00:00
func (v *directRay) OutboundOutput() OutputStream {
return v.Output
2015-10-14 13:07:13 +00:00
}
2016-11-27 20:39:09 +00:00
func (v *directRay) InboundInput() OutputStream {
return v.Input
2015-10-14 13:07:13 +00:00
}
2016-11-27 20:39:09 +00:00
func (v *directRay) InboundOutput() InputStream {
return v.Output
2015-10-14 13:07:13 +00:00
}
2016-04-18 16:44:10 +00:00
type Stream struct {
access sync.RWMutex
closed bool
2016-12-09 10:35:27 +00:00
buffer chan *buf.Buffer
2016-04-18 16:44:10 +00:00
}
func NewStream() *Stream {
return &Stream{
2016-12-09 10:35:27 +00:00
buffer: make(chan *buf.Buffer, bufferSize),
2016-04-18 16:44:10 +00:00
}
}
2016-12-09 10:35:27 +00:00
func (v *Stream) Read() (*buf.Buffer, error) {
2016-11-27 20:39:09 +00:00
if v.buffer == nil {
2016-04-18 16:44:10 +00:00
return nil, io.EOF
}
2016-11-27 20:39:09 +00:00
v.access.RLock()
if v.buffer == nil {
v.access.RUnlock()
2016-04-18 16:44:10 +00:00
return nil, io.EOF
}
2016-11-27 20:39:09 +00:00
channel := v.buffer
v.access.RUnlock()
2016-04-19 12:51:52 +00:00
result, open := <-channel
2016-04-18 16:44:10 +00:00
if !open {
return nil, io.EOF
}
return result, nil
}
2016-12-09 10:35:27 +00:00
func (v *Stream) Write(data *buf.Buffer) error {
2016-11-27 20:39:09 +00:00
for !v.closed {
err := v.TryWriteOnce(data)
2016-11-17 22:21:44 +00:00
if err != io.ErrNoProgress {
2016-05-09 02:04:51 +00:00
return err
}
}
2016-11-17 22:21:44 +00:00
return io.ErrClosedPipe
2016-05-09 02:04:51 +00:00
}
2016-12-09 10:35:27 +00:00
func (v *Stream) TryWriteOnce(data *buf.Buffer) error {
2016-11-27 20:39:09 +00:00
v.access.RLock()
defer v.access.RUnlock()
if v.closed {
2016-11-17 22:21:44 +00:00
return io.ErrClosedPipe
2016-04-19 08:30:42 +00:00
}
2016-05-09 02:04:51 +00:00
select {
2016-11-27 20:39:09 +00:00
case v.buffer <- data:
2016-05-09 02:04:51 +00:00
return nil
case <-time.After(2 * time.Second):
2016-11-17 22:21:44 +00:00
return io.ErrNoProgress
2016-05-09 02:04:51 +00:00
}
2016-04-18 16:44:10 +00:00
}
2016-11-27 20:39:09 +00:00
func (v *Stream) Close() {
if v.closed {
2016-04-18 16:44:10 +00:00
return
}
2016-11-27 20:39:09 +00:00
v.access.Lock()
defer v.access.Unlock()
if v.closed {
2016-04-18 16:44:10 +00:00
return
}
2016-11-27 20:39:09 +00:00
v.closed = true
close(v.buffer)
2016-04-18 16:44:10 +00:00
}
2016-11-27 20:39:09 +00:00
func (v *Stream) Release() {
if v.buffer == nil {
2016-04-18 16:44:10 +00:00
return
}
2016-11-27 20:39:09 +00:00
v.Close()
v.access.Lock()
defer v.access.Unlock()
if v.buffer == nil {
2016-04-18 16:44:10 +00:00
return
}
2016-11-27 20:39:09 +00:00
for data := range v.buffer {
2016-04-18 16:44:10 +00:00
data.Release()
}
2016-11-27 20:39:09 +00:00
v.buffer = nil
2016-04-18 16:44:10 +00:00
}