code refactor for common package

pkg/common/bufpool/pool.go

@@ -0,0 +1,102 @@
+package bufpool
+
+import "sync"
+
+var (
+	pools = []struct {
+		size int
+		pool sync.Pool
+	}{
+		{
+			size: 128,
+			pool: sync.Pool{
+				New: func() interface{} {
+					return make([]byte, 128)
+				},
+			},
+		},
+		{
+			size: 512,
+			pool: sync.Pool{
+				New: func() interface{} {
+					return make([]byte, 512)
+				},
+			},
+		},
+		{
+			size: 1024,
+			pool: sync.Pool{
+				New: func() interface{} {
+					return make([]byte, 1024)
+				},
+			},
+		},
+		{
+			size: 4096,
+			pool: sync.Pool{
+				New: func() interface{} {
+					return make([]byte, 4096)
+				},
+			},
+		},
+		{
+			size: 8192,
+			pool: sync.Pool{
+				New: func() interface{} {
+					return make([]byte, 8192)
+				},
+			},
+		},
+		{
+			size: 16 * 1024,
+			pool: sync.Pool{
+				New: func() interface{} {
+					return make([]byte, 16*1024)
+				},
+			},
+		},
+		{
+			size: 32 * 1024,
+			pool: sync.Pool{
+				New: func() interface{} {
+					return make([]byte, 32*1024)
+				},
+			},
+		},
+		{
+			size: 64 * 1024,
+			pool: sync.Pool{
+				New: func() interface{} {
+					return make([]byte, 64*1024)
+				},
+			},
+		},
+		{
+			size: 65 * 1024,
+			pool: sync.Pool{
+				New: func() interface{} {
+					return make([]byte, 65*1024)
+				},
+			},
+		},
+	}
+)
+
+// Get returns a buffer size range from (0, 65]KB,
+// panic if size > 65KB.
+func Get(size int) []byte {
+	for i := range pools {
+		if size <= pools[i].size {
+			return pools[i].pool.Get().([]byte)
+		}
+	}
+	panic("size too large (max=65KB)")
+}
+
+func Put(b []byte) {
+	for i := range pools {
+		if cap(b) == pools[i].size {
+			pools[i].pool.Put(b[:cap(b)])
+		}
+	}
+}

pkg/common/util/kcp/kcp.go

@@ -0,0 +1,34 @@
+package kcp
+
+import (
+	"net"
+
+	"github.com/golang/snappy"
+)
+
+type kcpCompStreamConn struct {
+	net.Conn
+	w *snappy.Writer
+	r *snappy.Reader
+}
+
+func KCPCompStreamConn(conn net.Conn) net.Conn {
+	return &kcpCompStreamConn{
+		Conn: conn,
+		w:    snappy.NewBufferedWriter(conn),
+		r:    snappy.NewReader(conn),
+	}
+}
+
+func (c *kcpCompStreamConn) Read(b []byte) (n int, err error) {
+	return c.r.Read(b)
+}
+
+func (c *kcpCompStreamConn) Write(b []byte) (n int, err error) {
+	n, err = c.w.Write(b)
+	if err != nil {
+		return
+	}
+	err = c.w.Flush()
+	return n, err
+}

pkg/common/util/mux/mux.go

@@ -0,0 +1,85 @@
+package mux
+
+import (
+	"net"
+
+	smux "github.com/xtaci/smux"
+)
+
+type Session struct {
+	conn    net.Conn
+	session *smux.Session
+}
+
+func ClientSession(conn net.Conn) (*Session, error) {
+	s, err := smux.Client(conn, smux.DefaultConfig())
+	if err != nil {
+		return nil, err
+	}
+	return &Session{
+		conn:    conn,
+		session: s,
+	}, nil
+}
+
+func ServerSession(conn net.Conn) (*Session, error) {
+	s, err := smux.Server(conn, smux.DefaultConfig())
+	if err != nil {
+		return nil, err
+	}
+	return &Session{
+		conn:    conn,
+		session: s,
+	}, nil
+}
+
+func (session *Session) GetConn() (net.Conn, error) {
+	stream, err := session.session.OpenStream()
+	if err != nil {
+		return nil, err
+	}
+	return &StreamConn{Conn: session.conn, stream: stream}, nil
+}
+
+func (session *Session) Accept() (net.Conn, error) {
+	stream, err := session.session.AcceptStream()
+	if err != nil {
+		return nil, err
+	}
+	return &StreamConn{Conn: session.conn, stream: stream}, nil
+}
+
+func (session *Session) Close() error {
+	if session.session == nil {
+		return nil
+	}
+	return session.session.Close()
+}
+
+func (session *Session) IsClosed() bool {
+	if session.session == nil {
+		return true
+	}
+	return session.session.IsClosed()
+}
+
+func (session *Session) NumStreams() int {
+	return session.session.NumStreams()
+}
+
+type StreamConn struct {
+	net.Conn
+	stream *smux.Stream
+}
+
+func (c *StreamConn) Read(b []byte) (n int, err error) {
+	return c.stream.Read(b)
+}
+
+func (c *StreamConn) Write(b []byte) (n int, err error) {
+	return c.stream.Write(b)
+}
+
+func (c *StreamConn) Close() error {
+	return c.stream.Close()
+}

pkg/common/util/quic/quic.go

@@ -0,0 +1,104 @@
+package quic
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"errors"
+	"io"
+	"net"
+
+	"github.com/lucas-clemente/quic-go"
+)
+
+type quicConn struct {
+	quic.Session
+	quic.Stream
+}
+
+func QUICConn(session quic.Session, stream quic.Stream) net.Conn {
+	return &quicConn{
+		Session: session,
+		Stream:  stream,
+	}
+}
+
+type quicCipherConn struct {
+	net.PacketConn
+	key []byte
+}
+
+func QUICCipherConn(conn net.PacketConn, key []byte) net.PacketConn {
+	return &quicCipherConn{
+		PacketConn: conn,
+		key:        key,
+	}
+}
+
+func (conn *quicCipherConn) ReadFrom(data []byte) (n int, addr net.Addr, err error) {
+	n, addr, err = conn.PacketConn.ReadFrom(data)
+	if err != nil {
+		return
+	}
+	b, err := conn.decrypt(data[:n])
+	if err != nil {
+		return
+	}
+
+	copy(data, b)
+
+	return len(b), addr, nil
+}
+
+func (conn *quicCipherConn) WriteTo(data []byte, addr net.Addr) (n int, err error) {
+	b, err := conn.encrypt(data)
+	if err != nil {
+		return
+	}
+
+	_, err = conn.PacketConn.WriteTo(b, addr)
+	if err != nil {
+		return
+	}
+
+	return len(b), nil
+}
+
+func (conn *quicCipherConn) encrypt(data []byte) ([]byte, error) {
+	c, err := aes.NewCipher(conn.key)
+	if err != nil {
+		return nil, err
+	}
+
+	gcm, err := cipher.NewGCM(c)
+	if err != nil {
+		return nil, err
+	}
+
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, err
+	}
+
+	return gcm.Seal(nonce, nonce, data, nil), nil
+}
+
+func (conn *quicCipherConn) decrypt(data []byte) ([]byte, error) {
+	c, err := aes.NewCipher(conn.key)
+	if err != nil {
+		return nil, err
+	}
+
+	gcm, err := cipher.NewGCM(c)
+	if err != nil {
+		return nil, err
+	}
+
+	nonceSize := gcm.NonceSize()
+	if len(data) < nonceSize {
+		return nil, errors.New("ciphertext too short")
+	}
+
+	nonce, ciphertext := data[:nonceSize], data[nonceSize:]
+	return gcm.Open(nil, nonce, ciphertext, nil)
+}

pkg/common/util/socks/conn.go

@@ -0,0 +1,174 @@
+package socks
+
+import (
+	"bytes"
+	"net"
+
+	"github.com/go-gost/gosocks5"
+	"github.com/go-gost/gost/pkg/common/bufpool"
+)
+
+var (
+	_ net.PacketConn = (*UDPTunConn)(nil)
+	_ net.Conn       = (*UDPTunConn)(nil)
+
+	_ net.PacketConn = (*UDPConn)(nil)
+	_ net.Conn       = (*UDPConn)(nil)
+)
+
+type UDPTunConn struct {
+	net.Conn
+	taddr net.Addr
+}
+
+func UDPTunClientConn(c net.Conn, targetAddr net.Addr) *UDPTunConn {
+	return &UDPTunConn{
+		Conn:  c,
+		taddr: targetAddr,
+	}
+}
+
+func UDPTunServerConn(c net.Conn) *UDPTunConn {
+	return &UDPTunConn{
+		Conn: c,
+	}
+}
+
+func (c *UDPTunConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
+	socksAddr := gosocks5.Addr{}
+	header := gosocks5.UDPHeader{
+		Addr: &socksAddr,
+	}
+	dgram := gosocks5.UDPDatagram{
+		Header: &header,
+		Data:   b,
+	}
+	_, err = dgram.ReadFrom(c.Conn)
+	if err != nil {
+		return
+	}
+
+	n = len(dgram.Data)
+	if n > len(b) {
+		n = copy(b, dgram.Data)
+	}
+	addr, err = net.ResolveUDPAddr("udp", socksAddr.String())
+
+	return
+}
+
+func (c *UDPTunConn) Read(b []byte) (n int, err error) {
+	n, _, err = c.ReadFrom(b)
+	return
+}
+
+func (c *UDPTunConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
+	socksAddr := gosocks5.Addr{}
+	if err = socksAddr.ParseFrom(addr.String()); err != nil {
+		return
+	}
+
+	header := gosocks5.UDPHeader{
+		Addr: &socksAddr,
+	}
+	dgram := gosocks5.UDPDatagram{
+		Header: &header,
+		Data:   b,
+	}
+	dgram.Header.Rsv = uint16(len(dgram.Data))
+	dgram.Header.Frag = 0xff // UDP tun relay flag, used by shadowsocks
+	_, err = dgram.WriteTo(c.Conn)
+	n = len(b)
+
+	return
+}
+
+func (c *UDPTunConn) Write(b []byte) (n int, err error) {
+	return c.WriteTo(b, c.taddr)
+}
+
+var (
+	DefaultBufferSize = 4096
+)
+
+type UDPConn struct {
+	net.PacketConn
+	raddr      net.Addr
+	taddr      net.Addr
+	bufferSize int
+}
+
+func NewUDPConn(c net.PacketConn, bufferSize int) *UDPConn {
+	return &UDPConn{
+		PacketConn: c,
+		bufferSize: bufferSize,
+	}
+}
+
+// ReadFrom reads an UDP datagram.
+// NOTE: for server side,
+// the returned addr is the target address the client want to relay to.
+func (c *UDPConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
+	rbuf := bufpool.Get(c.bufferSize)
+	defer bufpool.Put(rbuf)
+
+	n, c.raddr, err = c.PacketConn.ReadFrom(rbuf)
+	if err != nil {
+		return
+	}
+
+	socksAddr := gosocks5.Addr{}
+	header := gosocks5.UDPHeader{
+		Addr: &socksAddr,
+	}
+	hlen, err := header.ReadFrom(bytes.NewReader(rbuf[:n]))
+	if err != nil {
+		return
+	}
+	n = copy(b, rbuf[hlen:n])
+
+	addr, err = net.ResolveUDPAddr("udp", socksAddr.String())
+	return
+}
+
+func (c *UDPConn) Read(b []byte) (n int, err error) {
+	n, _, err = c.ReadFrom(b)
+	return
+}
+
+func (c *UDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
+	wbuf := bufpool.Get(c.bufferSize)
+	defer bufpool.Put(wbuf)
+
+	socksAddr := gosocks5.Addr{}
+	if err = socksAddr.ParseFrom(addr.String()); err != nil {
+		return
+	}
+
+	header := gosocks5.UDPHeader{
+		Addr: &socksAddr,
+	}
+	dgram := gosocks5.UDPDatagram{
+		Header: &header,
+		Data:   b,
+	}
+
+	buf := bytes.NewBuffer(wbuf[:0])
+	_, err = dgram.WriteTo(buf)
+	if err != nil {
+		return
+	}
+
+	_, err = c.PacketConn.WriteTo(buf.Bytes(), c.raddr)
+	n = len(b)
+
+	return
+}
+
+func (c *UDPConn) Write(b []byte) (n int, err error) {
+	return c.WriteTo(b, c.taddr)
+}
+
+func (c *UDPConn) RemoteAddr() net.Addr {
+	return c.raddr
+}

pkg/common/util/socks/socks.go

@@ -0,0 +1,18 @@
+package socks
+
+const (
+	// MethodTLS is an extended SOCKS5 method with tls encryption support.
+	MethodTLS uint8 = 0x80
+	// MethodTLSAuth is an extended SOCKS5 method with tls encryption and authentication support.
+	MethodTLSAuth uint8 = 0x82
+	// MethodMux is an extended SOCKS5 method for stream multiplexing.
+	MethodMux = 0x88
+)
+
+const (
+	// CmdMuxBind is an extended SOCKS5 request CMD for
+	// multiplexing transport with the binding server.
+	CmdMuxBind uint8 = 0xF2
+	// CmdUDPTun is an extended SOCKS5 request CMD for UDP over TCP.
+	CmdUDPTun uint8 = 0xF3
+)

pkg/common/util/ss/conn.go

@@ -0,0 +1,96 @@
+package ss
+
+import (
+	"bytes"
+	"net"
+
+	"github.com/go-gost/gosocks5"
+	"github.com/go-gost/gost/pkg/common/bufpool"
+)
+
+var (
+	DefaultBufferSize = 4096
+)
+
+var (
+	_ net.PacketConn = (*UDPConn)(nil)
+	_ net.Conn       = (*UDPConn)(nil)
+)
+
+type UDPConn struct {
+	net.PacketConn
+	raddr      net.Addr
+	taddr      net.Addr
+	bufferSize int
+}
+
+func UDPClientConn(c net.PacketConn, remoteAddr, targetAddr net.Addr, bufferSize int) *UDPConn {
+	return &UDPConn{
+		PacketConn: c,
+		raddr:      remoteAddr,
+		taddr:      targetAddr,
+		bufferSize: bufferSize,
+	}
+}
+
+func UDPServerConn(c net.PacketConn, remoteAddr net.Addr, bufferSize int) *UDPConn {
+	return &UDPConn{
+		PacketConn: c,
+		raddr:      remoteAddr,
+		bufferSize: bufferSize,
+	}
+}
+
+func (c *UDPConn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
+	rbuf := bufpool.Get(c.bufferSize)
+	defer bufpool.Put(rbuf)
+
+	n, _, err = c.PacketConn.ReadFrom(rbuf)
+	if err != nil {
+		return
+	}
+
+	saddr := gosocks5.Addr{}
+	addrLen, err := saddr.ReadFrom(bytes.NewReader(rbuf[:n]))
+	if err != nil {
+		return
+	}
+
+	n = copy(b, rbuf[addrLen:n])
+	addr, err = net.ResolveUDPAddr("udp", saddr.String())
+
+	return
+}
+
+func (c *UDPConn) Read(b []byte) (n int, err error) {
+	n, _, err = c.ReadFrom(b)
+	return
+}
+
+func (c *UDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
+	wbuf := bufpool.Get(c.bufferSize)
+	defer bufpool.Put(wbuf)
+
+	socksAddr := gosocks5.Addr{}
+	if err = socksAddr.ParseFrom(addr.String()); err != nil {
+		return
+	}
+
+	addrLen, err := socksAddr.Encode(wbuf)
+	if err != nil {
+		return
+	}
+
+	n = copy(wbuf[addrLen:], b)
+	_, err = c.PacketConn.WriteTo(wbuf[:addrLen+n], c.raddr)
+
+	return
+}
+
+func (c *UDPConn) Write(b []byte) (n int, err error) {
+	return c.WriteTo(b, c.taddr)
+}
+
+func (c *UDPConn) RemoteAddr() net.Addr {
+	return c.raddr
+}

pkg/common/util/ss/ss.go

@@ -0,0 +1,60 @@
+package ss
+
+import (
+	"bytes"
+	"net"
+
+	"github.com/shadowsocks/go-shadowsocks2/core"
+	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
+)
+
+type shadowCipher struct {
+	cipher *ss.Cipher
+}
+
+func (c *shadowCipher) StreamConn(conn net.Conn) net.Conn {
+	return ss.NewConn(conn, c.cipher.Copy())
+}
+
+func (c *shadowCipher) PacketConn(conn net.PacketConn) net.PacketConn {
+	return ss.NewSecurePacketConn(conn, c.cipher.Copy())
+}
+
+func ShadowCipher(method, password string, key string) (core.Cipher, error) {
+	if method == "" && password == "" {
+		return nil, nil
+	}
+
+	c, _ := ss.NewCipher(method, password)
+	if c != nil {
+		return &shadowCipher{cipher: c}, nil
+	}
+
+	return core.PickCipher(method, []byte(key), password)
+}
+
+// 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
+	wbuf *bytes.Buffer
+}
+
+func ShadowConn(conn net.Conn, header []byte) net.Conn {
+	return &shadowConn{
+		Conn: conn,
+		wbuf: bytes.NewBuffer(header),
+	}
+}
+
+func (c *shadowConn) Write(b []byte) (n int, err error) {
+	n = len(b) // force byte length consistent
+	if c.wbuf.Len() > 0 {
+		c.wbuf.Write(b) // append the data to the cached header
+		_, err = c.Conn.Write(c.wbuf.Bytes())
+		c.wbuf.Reset()
+		return
+	}
+	_, err = c.Conn.Write(b)
+	return
+}

pkg/common/util/tcp.go

@@ -0,0 +1,32 @@
+package util
+
+import (
+	"net"
+	"time"
+)
+
+const (
+	defaultKeepAlivePeriod = 180 * time.Second
+)
+
+// TCPKeepAliveListener is a TCP listener with keep alive enabled.
+type TCPKeepAliveListener struct {
+	KeepAlivePeriod time.Duration
+	*net.TCPListener
+}
+
+func (l *TCPKeepAliveListener) Accept() (c net.Conn, err error) {
+	tc, err := l.AcceptTCP()
+	if err != nil {
+		return
+	}
+
+	tc.SetKeepAlive(true)
+	period := l.KeepAlivePeriod
+	if period <= 0 {
+		period = defaultKeepAlivePeriod
+	}
+	tc.SetKeepAlivePeriod(period)
+
+	return tc, nil
+}

pkg/common/util/tls/tls.go

@@ -0,0 +1,40 @@
+package tls
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"errors"
+	"io/ioutil"
+)
+
+// LoadTLSConfig loads the certificate from cert & key files and optional client CA file.
+func LoadTLSConfig(certFile, keyFile, caFile string) (*tls.Config, error) {
+	cert, err := tls.LoadX509KeyPair(certFile, keyFile)
+	if err != nil {
+		return nil, err
+	}
+
+	cfg := &tls.Config{Certificates: []tls.Certificate{cert}}
+
+	if pool, _ := loadCA(caFile); pool != nil {
+		cfg.ClientCAs = pool
+		cfg.ClientAuth = tls.RequireAndVerifyClientCert
+	}
+
+	return cfg, nil
+}
+
+func loadCA(caFile string) (cp *x509.CertPool, err error) {
+	if caFile == "" {
+		return
+	}
+	cp = x509.NewCertPool()
+	data, err := ioutil.ReadFile(caFile)
+	if err != nil {
+		return nil, err
+	}
+	if !cp.AppendCertsFromPEM(data) {
+		return nil, errors.New("AppendCertsFromPEM failed")
+	}
+	return
+}

pkg/common/util/ws/ws.go

@@ -0,0 +1,41 @@
+package ws
+
+import (
+	"net"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+type websocketConn struct {
+	*websocket.Conn
+	rb []byte
+}
+
+func WebsocketServerConn(conn *websocket.Conn) net.Conn {
+	return &websocketConn{
+		Conn: conn,
+	}
+}
+
+func (c *websocketConn) Read(b []byte) (n int, err error) {
+	if len(c.rb) == 0 {
+		_, c.rb, err = c.ReadMessage()
+	}
+	n = copy(b, c.rb)
+	c.rb = c.rb[n:]
+	return
+}
+
+func (c *websocketConn) Write(b []byte) (n int, err error) {
+	err = c.WriteMessage(websocket.BinaryMessage, b)
+	n = len(b)
+	return
+}
+
+func (c *websocketConn) SetDeadline(t time.Time) error {
+	if err := c.SetReadDeadline(t); err != nil {
+		return err
+	}
+	return c.SetWriteDeadline(t)
+}