initial commit

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

internal/util/socks/conn.go

@@ -0,0 +1,172 @@
+package socks
+
+import (
+	"bytes"
+	"net"
+
+	"github.com/go-gost/core/common/bufpool"
+	"github.com/go-gost/gosocks5"
+)
+
+type udpTunConn struct {
+	net.Conn
+	taddr net.Addr
+}
+
+func UDPTunClientConn(c net.Conn, targetAddr net.Addr) net.Conn {
+	return &udpTunConn{
+		Conn:  c,
+		taddr: targetAddr,
+	}
+}
+
+func UDPTunClientPacketConn(c net.Conn) net.PacketConn {
+	return &udpTunConn{
+		Conn: c,
+	}
+}
+
+func UDPTunServerConn(c net.Conn) net.PacketConn {
+	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 UDPConn(c net.PacketConn, bufferSize int) net.PacketConn {
+	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
+}

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

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