add tests for shadowsocks UDP relay

This commit is contained in:
ginuerzh 2018-12-27 19:59:21 +08:00
parent a7d49f0b37
commit 72d7850c2c
4 changed files with 197 additions and 47 deletions

View File

@ -617,15 +617,15 @@ func (c *udpServerConn) RemoteAddr() net.Addr {
}
func (c *udpServerConn) SetDeadline(t time.Time) error {
return nil
return c.conn.SetDeadline(t)
}
func (c *udpServerConn) SetReadDeadline(t time.Time) error {
return nil
return c.conn.SetReadDeadline(t)
}
func (c *udpServerConn) SetWriteDeadline(t time.Time) error {
return nil
return c.conn.SetWriteDeadline(t)
}
type tcpRemoteForwardListener struct {

163
ss.go
View File

@ -16,46 +16,6 @@ import (
ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
)
// Due to in/out byte length is inconsistent of the shadowsocks.Conn.Write,
// we wrap around it to make io.Copy happy.
type shadowConn struct {
conn net.Conn
}
func (c *shadowConn) Read(b []byte) (n int, err error) {
return c.conn.Read(b)
}
func (c *shadowConn) Write(b []byte) (n int, err error) {
n = len(b) // force byte length consistent
_, err = c.conn.Write(b)
return
}
func (c *shadowConn) Close() error {
return c.conn.Close()
}
func (c *shadowConn) LocalAddr() net.Addr {
return c.conn.LocalAddr()
}
func (c *shadowConn) RemoteAddr() net.Addr {
return c.conn.RemoteAddr()
}
func (c *shadowConn) SetDeadline(t time.Time) error {
return c.conn.SetDeadline(t)
}
func (c *shadowConn) SetReadDeadline(t time.Time) error {
return c.conn.SetReadDeadline(t)
}
func (c *shadowConn) SetWriteDeadline(t time.Time) error {
return c.conn.SetWriteDeadline(t)
}
type shadowConnector struct {
Cipher *url.Userinfo
}
@ -102,7 +62,7 @@ func (c *shadowConnector) Connect(conn net.Conn, addr string, options ...Connect
if _, err := sc.Write(rawaddr); err != nil {
return nil, err
}
return &shadowConn{conn: sc}, nil
return &shadowConn{sc}, nil
}
type shadowHandler struct {
@ -142,7 +102,7 @@ func (h *shadowHandler) Handle(conn net.Conn) {
conn.RemoteAddr(), conn.LocalAddr(), err)
return
}
conn = &shadowConn{conn: ss.NewConn(conn, cipher)}
conn = &shadowConn{ss.NewConn(conn, cipher)}
conn.SetReadDeadline(time.Now().Add(ReadTimeout))
host, err := h.getRequest(conn)
@ -284,6 +244,55 @@ func (h *shadowHandler) getRequest(r io.Reader) (host string, err error) {
return
}
type shadowUDPConnector struct {
Cipher *url.Userinfo
}
// ShadowUDPConnector creates a Connector for shadowsocks UDP client.
// It accepts a cipher info for shadowsocks data encryption/decryption.
// The cipher must not be nil.
func ShadowUDPConnector(cipher *url.Userinfo) Connector {
return &shadowUDPConnector{Cipher: cipher}
}
func (c *shadowUDPConnector) Connect(conn net.Conn, addr string, options ...ConnectOption) (net.Conn, error) {
opts := &ConnectOptions{}
for _, option := range options {
option(opts)
}
timeout := opts.Timeout
if timeout <= 0 {
timeout = ConnectTimeout
}
conn.SetDeadline(time.Now().Add(timeout))
defer conn.SetDeadline(time.Time{})
rawaddr, err := ss.RawAddr(addr)
if err != nil {
return nil, err
}
var method, password string
if c.Cipher != nil {
method = c.Cipher.Username()
password, _ = c.Cipher.Password()
}
cipher, err := ss.NewCipher(method, password)
if err != nil {
return nil, err
}
sc := ss.NewSecurePacketConn(&shadowPacketConn{conn}, cipher, false)
return &shadowUDPConn{
PacketConn: sc,
raddr: conn.RemoteAddr(),
header: rawaddr,
}, nil
}
type shadowUDPListener struct {
ln net.PacketConn
conns map[string]*udpServerConn
@ -434,7 +443,9 @@ func (h *shadowUDPdHandler) transportUDP(sc net.Conn, cc net.PacketConn) error {
errc := make(chan error, 1)
go func() {
for {
b := make([]byte, mediumBufferSize)
b := mPool.Get().([]byte)
defer mPool.Put(b)
n, err := sc.Read(b[3:]) // add rsv and frag fields to make it the standard SOCKS5 UDP datagram
if err != nil {
// log.Logf("[ssu] %s - %s : %s", sc.RemoteAddr(), sc.LocalAddr(), err)
@ -468,7 +479,9 @@ func (h *shadowUDPdHandler) transportUDP(sc net.Conn, cc net.PacketConn) error {
go func() {
for {
b := make([]byte, mediumBufferSize)
b := mPool.Get().([]byte)
defer mPool.Put(b)
n, addr, err := cc.ReadFrom(b)
if err != nil {
errc <- err
@ -501,3 +514,65 @@ func (h *shadowUDPdHandler) transportUDP(sc net.Conn, cc net.PacketConn) error {
}
return err
}
// Due to in/out byte length is inconsistent of the shadowsocks.Conn.Write,
// we wrap around it to make io.Copy happy.
type shadowConn struct {
net.Conn
}
func (c *shadowConn) Write(b []byte) (n int, err error) {
n = len(b) // force byte length consistent
_, err = c.Conn.Write(b)
return
}
type shadowUDPConn struct {
net.PacketConn
raddr net.Addr
header []byte
}
func (c *shadowUDPConn) Write(b []byte) (n int, err error) {
n = len(b) // force byte length consistent
if len(c.header) > 0 {
b = append(c.header, b...)
}
_, err = c.PacketConn.WriteTo(b, c.raddr)
return
}
func (c *shadowUDPConn) Read(b []byte) (n int, err error) {
buf := mPool.Get().([]byte)
defer mPool.Put(buf)
n, _, err = c.PacketConn.ReadFrom(buf[3:])
if err != nil {
return
}
dgram, err := gosocks5.ReadUDPDatagram(bytes.NewReader(buf[:n+3]))
if err != nil {
return
}
n = copy(b, dgram.Data)
return
}
func (c *shadowUDPConn) RemoteAddr() net.Addr {
return c.raddr
}
type shadowPacketConn struct {
net.Conn
}
func (c *shadowPacketConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
n, err = c.Conn.Read(b)
addr = c.Conn.RemoteAddr()
return
}
func (c *shadowPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
return c.Conn.Write(b)
}

View File

@ -5,8 +5,13 @@ import (
"net/http/httptest"
"net/url"
"testing"
"time"
)
func init() {
// ss.Debug = true
}
var ssTests = []struct {
clientCipher *url.Userinfo
serverCipher *url.Userinfo
@ -289,3 +294,72 @@ func BenchmarkSSProxyParallel(b *testing.B) {
}
})
}
func shadowUDPRoundtrip(t *testing.T, host string, data []byte) error {
ln, err := ShadowUDPListener("localhost:0", url.UserPassword("chacha20-ietf", "123456"), 0)
if err != nil {
return err
}
client := &Client{
Connector: ShadowUDPConnector(url.UserPassword("chacha20-ietf", "123456")),
Transporter: UDPTransporter(),
}
server := &Server{
Handler: ShadowUDPdHandler(),
Listener: ln,
}
go server.Run()
defer server.Close()
return udpRoundtrip(client, server, host, data)
}
func TestShadowUDP(t *testing.T) {
udpSrv := newUDPTestServer(udpTestHandler)
udpSrv.Start()
defer udpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
err := shadowUDPRoundtrip(t, udpSrv.Addr(), sendData)
if err != nil {
t.Error(err)
}
}
// TODO: fix shadowsocks UDP relay benchmark.
func BenchmarkShadowUDP(b *testing.B) {
udpSrv := newUDPTestServer(udpTestHandler)
udpSrv.Start()
defer udpSrv.Close()
sendData := make([]byte, 128)
rand.Read(sendData)
ln, err := ShadowUDPListener("localhost:0", url.UserPassword("chacha20-ietf", "123456"), 1000*time.Millisecond)
if err != nil {
b.Error(err)
}
client := &Client{
Connector: ShadowUDPConnector(url.UserPassword("chacha20-ietf", "123456")),
Transporter: UDPTransporter(),
}
server := &Server{
Handler: ShadowUDPdHandler(),
Listener: ln,
}
go server.Run()
defer server.Close()
for i := 0; i < b.N; i++ {
if err := udpRoundtrip(client, server, udpSrv.Addr(), sendData); err != nil {
b.Error(err)
}
}
}

View File

@ -166,6 +166,7 @@ func sshRemoteForwardRoundtrip(t *testing.T, targetURL string, data []byte) (err
return httpRoundtrip(c, targetURL, data)
}
// TODO: fix this test
func _TestSSHRemoteForward(t *testing.T) {
httpSrv := httptest.NewServer(httpTestHandler)
defer httpSrv.Close()