diff --git a/forward.go b/forward.go index 7f16144..5a9b14d 100644 --- a/forward.go +++ b/forward.go @@ -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 { diff --git a/ss.go b/ss.go index e492f51..a5eb563 100644 --- a/ss.go +++ b/ss.go @@ -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) +} diff --git a/ss_test.go b/ss_test.go index 0b12602..87ae6e8 100644 --- a/ss_test.go +++ b/ss_test.go @@ -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) + } + } +} diff --git a/ssh_test.go b/ssh_test.go index 6c11f2b..db3638e 100644 --- a/ssh_test.go +++ b/ssh_test.go @@ -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()