-
Notifications
You must be signed in to change notification settings - Fork 16
/
stream.go
117 lines (99 loc) · 4.15 KB
/
stream.go
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package direct
import (
"context"
"io"
"github.com/database64128/shadowsocks-go/conn"
"github.com/database64128/shadowsocks-go/socks5"
"github.com/database64128/shadowsocks-go/zerocopy"
)
var (
_ zerocopy.DirectReader = (*DirectStreamReadWriter)(nil)
_ zerocopy.DirectWriter = (*DirectStreamReadWriter)(nil)
)
// DirectStreamReadWriter implements the zerocopy ReadWriter interface and reads/writes everything
// directly from/to the wrapped io.ReadWriter.
type DirectStreamReadWriter struct {
rw zerocopy.DirectReadWriteCloser
}
// WriterInfo implements the Writer WriterInfo method.
func (rw *DirectStreamReadWriter) WriterInfo() zerocopy.WriterInfo {
return zerocopy.WriterInfo{}
}
// WriteZeroCopy implements the Writer WriteZeroCopy method.
func (rw *DirectStreamReadWriter) WriteZeroCopy(b []byte, payloadStart, payloadLen int) (payloadWritten int, err error) {
payloadWritten, err = rw.rw.Write(b[payloadStart : payloadStart+payloadLen])
return
}
// DirectWriter implements the DirectWriter DirectWriter method.
func (rw *DirectStreamReadWriter) DirectWriter() io.Writer {
return rw.rw
}
// ReaderInfo implements the Reader ReaderInfo method.
func (rw *DirectStreamReadWriter) ReaderInfo() zerocopy.ReaderInfo {
return zerocopy.ReaderInfo{}
}
// ReadZeroCopy implements the Reader ReadZeroCopy method.
func (rw *DirectStreamReadWriter) ReadZeroCopy(b []byte, payloadBufStart, payloadBufLen int) (payloadLen int, err error) {
payloadLen, err = rw.rw.Read(b[payloadBufStart : payloadBufStart+payloadBufLen])
return
}
// DirectReader implements the DirectReader DirectReader method.
func (rw *DirectStreamReadWriter) DirectReader() io.Reader {
return rw.rw
}
// CloseWrite implements the ReadWriter CloseWrite method.
func (rw *DirectStreamReadWriter) CloseWrite() error {
return rw.rw.CloseWrite()
}
// CloseRead implements the ReadWriter CloseRead method.
func (rw *DirectStreamReadWriter) CloseRead() error {
return rw.rw.CloseRead()
}
// Close implements the ReadWriter Close method.
func (rw *DirectStreamReadWriter) Close() error {
return rw.rw.Close()
}
// NewDirectStreamReadWriter returns a ReadWriter that passes all reads and writes directly to the underlying stream.
func NewDirectStreamReadWriter(rw zerocopy.DirectReadWriteCloser) *DirectStreamReadWriter {
return &DirectStreamReadWriter{
rw: rw,
}
}
// NewShadowsocksNoneStreamClientReadWriter creates a ReadWriter that acts as a Shadowsocks none method client.
func NewShadowsocksNoneStreamClientReadWriter(ctx context.Context, rwo zerocopy.DirectReadWriteCloserOpener, targetAddr conn.Addr, payload []byte) (*DirectStreamReadWriter, zerocopy.DirectReadWriteCloser, error) {
targetAddrLen := socks5.LengthOfAddrFromConnAddr(targetAddr)
writeBuf := make([]byte, targetAddrLen+len(payload))
socks5.WriteAddrFromConnAddr(writeBuf, targetAddr)
copy(writeBuf[targetAddrLen:], payload)
rawRW, err := rwo.Open(ctx, writeBuf)
if err != nil {
return nil, nil, err
}
return &DirectStreamReadWriter{rw: rawRW}, rawRW, nil
}
// NewShadowsocksNoneStreamServerReadWriter creates a ReadWriter that acts as a Shadowsocks none method server.
func NewShadowsocksNoneStreamServerReadWriter(rw zerocopy.DirectReadWriteCloser) (*DirectStreamReadWriter, conn.Addr, error) {
addr, err := socks5.ConnAddrFromReader(rw)
if err != nil {
return nil, addr, err
}
return &DirectStreamReadWriter{rw: rw}, addr, nil
}
// NewSocks5StreamClientReadWriter writes a SOCKS5 CONNECT request to rw and wraps rw into a ReadWriter ready for use.
func NewSocks5StreamClientReadWriter(rw zerocopy.DirectReadWriteCloser, targetAddr conn.Addr) (*DirectStreamReadWriter, error) {
if err := socks5.ClientConnect(rw, targetAddr); err != nil {
return nil, err
}
return &DirectStreamReadWriter{rw: rw}, nil
}
// NewSocks5StreamServerReadWriter handles a SOCKS5 request from rw and wraps rw into a ReadWriter ready for use.
// conn must be provided when UDP is enabled.
func NewSocks5StreamServerReadWriter(rw zerocopy.DirectReadWriteCloser, enableTCP, enableUDP bool) (dsrw *DirectStreamReadWriter, addr conn.Addr, err error) {
addr, err = socks5.ServerAccept(rw, enableTCP, enableUDP)
if err == nil {
dsrw = &DirectStreamReadWriter{
rw: rw,
}
}
return
}