initial commit

internal/util/ss/conn.go

@@ -0,0 +1,96 @@
+package ss
+
+import (
+	"bytes"
+	"net"
+
+	"github.com/go-gost/gosocks5"
+	"github.com/go-gost/gost/v3/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
+}

internal/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
+}