This commit is contained in:
ginuerzh
2021-11-16 16:12:16 +08:00
parent 2c0ce35b0b
commit 83dacf67d5
15 changed files with 543 additions and 266 deletions

View File

@ -26,7 +26,7 @@ type rtcpListener struct {
chain *chain.Chain
md metadata
ln net.Listener
connChan chan net.Conn
cqueue chan net.Conn
session *mux.Session
sessionMux sync.Mutex
logger logger.Logger
@ -57,7 +57,7 @@ func (l *rtcpListener) Init(md md.Metadata) (err error) {
}
l.laddr = laddr
l.connChan = make(chan net.Conn, l.md.connQueueSize)
l.cqueue = make(chan net.Conn, l.md.backlog)
if l.chain.IsEmpty() {
l.ln, err = net.ListenTCP("tcp", laddr)
@ -93,7 +93,7 @@ func (l *rtcpListener) Accept() (conn net.Conn, err error) {
}
select {
case conn = <-l.connChan:
case conn = <-l.cqueue:
case <-l.closed:
err = net.ErrClosed
}
@ -130,7 +130,7 @@ func (l *rtcpListener) listenLoop() {
tempDelay = 0
select {
case l.connChan <- conn:
case l.cqueue <- conn:
default:
conn.Close()
l.logger.Warnf("connection queue is full, client %s discarded", conn.RemoteAddr().String())
@ -169,36 +169,29 @@ func (l *rtcpListener) waitPeer(conn net.Conn) (net.Conn, error) {
addr.ParseFrom(l.addr)
req := gosocks5.NewRequest(gosocks5.CmdBind, &addr)
if err := req.Write(conn); err != nil {
l.logger.Error(err)
return nil, err
}
// first reply, bind status
rep, err := gosocks5.ReadReply(conn)
if err != nil {
l.logger.Error(err)
return nil, err
}
l.logger.Debug(rep)
if rep.Rep != gosocks5.Succeeded {
err = fmt.Errorf("bind on %s failed", l.addr)
l.logger.Error(err)
return nil, err
return nil, fmt.Errorf("bind on %s failed", l.addr)
}
l.logger.Debugf("bind on %s OK", rep.Addr)
// second reply, peer connected
rep, err = gosocks5.ReadReply(conn)
if err != nil {
l.logger.Error(err)
return nil, err
}
if rep.Rep != gosocks5.Succeeded {
err = fmt.Errorf("peer connect failed")
l.logger.Error(err)
return nil, err
return nil, fmt.Errorf("peer connect failed")
}
raddr, err := net.ResolveTCPAddr("tcp", rep.Addr.String())

View File

@ -8,28 +8,28 @@ import (
const (
defaultKeepAlivePeriod = 180 * time.Second
defaultConnQueueSize = 128
defaultBacklog = 128
)
type metadata struct {
enableMux bool
connQueueSize int
retryCount int
enableMux bool
backlog int
retryCount int
}
func (l *rtcpListener) parseMetadata(md md.Metadata) (err error) {
const (
enableMux = "mux"
connQueueSize = "connQueueSize"
retryCount = "retry"
enableMux = "mux"
backlog = "backlog"
retryCount = "retry"
)
l.md.enableMux = md.GetBool(enableMux)
l.md.retryCount = md.GetInt(retryCount)
l.md.connQueueSize = md.GetInt(connQueueSize)
if l.md.connQueueSize <= 0 {
l.md.connQueueSize = defaultConnQueueSize
l.md.backlog = md.GetInt(backlog)
if l.md.backlog <= 0 {
l.md.backlog = defaultBacklog
}
return
}

View File

@ -0,0 +1,213 @@
package rudp
import (
"context"
"fmt"
"net"
"time"
"github.com/go-gost/gosocks5"
"github.com/go-gost/gost/pkg/chain"
"github.com/go-gost/gost/pkg/common/bufpool"
"github.com/go-gost/gost/pkg/common/util/socks"
"github.com/go-gost/gost/pkg/common/util/udp"
"github.com/go-gost/gost/pkg/listener"
"github.com/go-gost/gost/pkg/logger"
md "github.com/go-gost/gost/pkg/metadata"
"github.com/go-gost/gost/pkg/registry"
)
func init() {
registry.RegisterListener("rudp", NewListener)
}
type rudpListener struct {
addr string
laddr *net.UDPAddr
chain *chain.Chain
md metadata
cqueue chan net.Conn
closed chan struct{}
connPool *udp.ConnPool
logger logger.Logger
}
func NewListener(opts ...listener.Option) listener.Listener {
options := &listener.Options{}
for _, opt := range opts {
opt(options)
}
return &rudpListener{
addr: options.Addr,
chain: options.Chain,
closed: make(chan struct{}),
logger: options.Logger,
}
}
func (l *rudpListener) Init(md md.Metadata) (err error) {
if err = l.parseMetadata(md); err != nil {
return
}
laddr, err := net.ResolveUDPAddr("udp", l.addr)
if err != nil {
return
}
l.laddr = laddr
l.cqueue = make(chan net.Conn, l.md.backlog)
l.connPool = udp.NewConnPool(l.md.ttl).WithLogger(l.logger)
go l.listenLoop()
return
}
func (l *rudpListener) Accept() (conn net.Conn, err error) {
select {
case conn = <-l.cqueue:
return
case <-l.closed:
return nil, listener.ErrClosed
}
}
func (l *rudpListener) Close() error {
select {
case <-l.closed:
default:
close(l.closed)
l.connPool.Close()
}
return nil
}
func (l *rudpListener) Addr() net.Addr {
return l.laddr
}
func (l *rudpListener) listenLoop() {
for {
conn, err := l.connect()
if err != nil {
l.logger.Error(err)
return
}
func() {
defer conn.Close()
for {
b := bufpool.Get(l.md.readBufferSize)
n, raddr, err := conn.ReadFrom(b)
if err != nil {
return
}
c := l.getConn(conn, raddr)
if c == nil {
bufpool.Put(b)
continue
}
if err := c.WriteQueue(b[:n]); err != nil {
l.logger.Warn("data discarded: ", err)
}
}
}()
}
}
func (l *rudpListener) connect() (conn net.PacketConn, err error) {
var tempDelay time.Duration
for {
select {
case <-l.closed:
return nil, net.ErrClosed
default:
}
conn, err = func() (net.PacketConn, error) {
if l.chain.IsEmpty() {
return net.ListenUDP("udp", l.laddr)
}
r := (&chain.Router{}).
WithChain(l.chain).
WithRetry(l.md.retryCount).
WithLogger(l.logger)
cc, err := r.Connect(context.Background())
if err != nil {
return nil, err
}
conn, err := l.initUDPTunnel(cc)
if err != nil {
cc.Close()
return nil, err
}
return conn, err
}()
if err == nil {
return
}
if tempDelay == 0 {
tempDelay = 1000 * time.Millisecond
} else {
tempDelay *= 2
}
if max := 6 * time.Second; tempDelay > max {
tempDelay = max
}
l.logger.Warnf("accept: %v, retrying in %v", err, tempDelay)
time.Sleep(tempDelay)
}
}
func (l *rudpListener) initUDPTunnel(conn net.Conn) (net.PacketConn, error) {
socksAddr := gosocks5.Addr{}
socksAddr.ParseFrom(l.laddr.String())
req := gosocks5.NewRequest(socks.CmdUDPTun, &socksAddr)
if err := req.Write(conn); err != nil {
return nil, err
}
l.logger.Debug(req)
reply, err := gosocks5.ReadReply(conn)
if err != nil {
return nil, err
}
l.logger.Debug(reply)
if reply.Rep != gosocks5.Succeeded {
return nil, fmt.Errorf("bind on %s failed", l.laddr)
}
baddr, err := net.ResolveUDPAddr("udp", reply.Addr.String())
if err != nil {
return nil, err
}
l.logger.Debugf("bind on %s OK", baddr)
return socks.UDPTunClientConn(conn, nil), nil
}
func (l *rudpListener) getConn(conn net.PacketConn, raddr net.Addr) *udp.Conn {
c, ok := l.connPool.Get(raddr.String())
if !ok {
c = udp.NewConn(conn, l.laddr, raddr, l.md.readQueueSize)
select {
case l.cqueue <- c:
l.connPool.Set(raddr.String(), c)
default:
c.Close()
l.logger.Warnf("connection queue is full, client %s discarded", raddr.String())
return nil
}
}
return c
}

View File

@ -0,0 +1,55 @@
package rudp
import (
"time"
md "github.com/go-gost/gost/pkg/metadata"
)
const (
defaultTTL = 60 * time.Second
defaultReadBufferSize = 4096
defaultReadQueueSize = 128
defaultBacklog = 128
)
type metadata struct {
ttl time.Duration
readBufferSize int
readQueueSize int
backlog int
retryCount int
}
func (l *rudpListener) parseMetadata(md md.Metadata) (err error) {
const (
ttl = "ttl"
readBufferSize = "readBufferSize"
readQueueSize = "readQueueSize"
backlog = "backlog"
retryCount = "retry"
)
l.md.ttl = md.GetDuration(ttl)
if l.md.ttl <= 0 {
l.md.ttl = defaultTTL
}
l.md.readBufferSize = md.GetInt(readBufferSize)
if l.md.readBufferSize <= 0 {
l.md.readBufferSize = defaultReadBufferSize
}
l.md.readQueueSize = md.GetInt(readQueueSize)
if l.md.readQueueSize <= 0 {
l.md.readQueueSize = defaultReadQueueSize
}
l.md.backlog = md.GetInt(backlog)
if l.md.backlog <= 0 {
l.md.backlog = defaultBacklog
}
l.md.retryCount = md.GetInt(retryCount)
return
}

View File

@ -1,190 +0,0 @@
package udp
import (
"errors"
"net"
"sync"
"sync/atomic"
"time"
"github.com/go-gost/gost/pkg/common/bufpool"
"github.com/go-gost/gost/pkg/logger"
)
// conn is a server side connection for UDP client peer, it implements net.Conn and net.PacketConn.
type conn struct {
net.PacketConn
remoteAddr net.Addr
rc chan []byte // data receive queue
idle int32
closed chan struct{}
closeMutex sync.Mutex
}
func newConn(c net.PacketConn, raddr net.Addr, queue int) *conn {
return &conn{
PacketConn: c,
remoteAddr: raddr,
rc: make(chan []byte, queue),
closed: make(chan struct{}),
}
}
func (c *conn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
select {
case bb := <-c.rc:
n = copy(b, bb)
c.SetIdle(false)
bufpool.Put(bb)
case <-c.closed:
err = net.ErrClosed
return
}
addr = c.remoteAddr
return
}
func (c *conn) Read(b []byte) (n int, err error) {
n, _, err = c.ReadFrom(b)
return
}
func (c *conn) Write(b []byte) (n int, err error) {
return c.WriteTo(b, c.remoteAddr)
}
func (c *conn) Close() error {
c.closeMutex.Lock()
defer c.closeMutex.Unlock()
select {
case <-c.closed:
default:
close(c.closed)
}
return nil
}
func (c *conn) RemoteAddr() net.Addr {
return c.remoteAddr
}
func (c *conn) IsIdle() bool {
return atomic.LoadInt32(&c.idle) > 0
}
func (c *conn) SetIdle(idle bool) {
v := int32(0)
if idle {
v = 1
}
atomic.StoreInt32(&c.idle, v)
}
func (c *conn) Queue(b []byte) error {
select {
case c.rc <- b:
return nil
case <-c.closed:
return net.ErrClosed
default:
return errors.New("recv queue is full")
}
}
type connPool struct {
m sync.Map
ttl time.Duration
closed chan struct{}
logger logger.Logger
}
func newConnPool(ttl time.Duration) *connPool {
p := &connPool{
ttl: ttl,
closed: make(chan struct{}),
}
go p.idleCheck()
return p
}
func (p *connPool) WithLogger(logger logger.Logger) *connPool {
p.logger = logger
return p
}
func (p *connPool) Get(key interface{}) (c *conn, ok bool) {
v, ok := p.m.Load(key)
if ok {
c, ok = v.(*conn)
}
return
}
func (p *connPool) Set(key interface{}, c *conn) {
p.m.Store(key, c)
}
func (p *connPool) Delete(key interface{}) {
p.m.Delete(key)
}
func (p *connPool) Close() {
select {
case <-p.closed:
return
default:
}
close(p.closed)
p.m.Range(func(k, v interface{}) bool {
if c, ok := v.(*conn); ok && c != nil {
c.Close()
}
return true
})
}
func (p *connPool) idleCheck() {
ticker := time.NewTicker(p.ttl)
defer ticker.Stop()
for {
select {
case <-ticker.C:
size := 0
idles := 0
p.m.Range(func(key, value interface{}) bool {
c, ok := value.(*conn)
if !ok || c == nil {
p.Delete(key)
return true
}
size++
if c.IsIdle() {
idles++
p.Delete(key)
c.Close()
return true
}
c.SetIdle(true)
return true
})
if idles > 0 {
p.logger.Debugf("connection pool: size=%d, idle=%d", size, idles)
}
case <-p.closed:
return
}
}
}

View File

@ -4,6 +4,7 @@ import (
"net"
"github.com/go-gost/gost/pkg/common/bufpool"
"github.com/go-gost/gost/pkg/common/util/udp"
"github.com/go-gost/gost/pkg/listener"
"github.com/go-gost/gost/pkg/logger"
md "github.com/go-gost/gost/pkg/metadata"
@ -18,10 +19,10 @@ type udpListener struct {
addr string
md metadata
conn net.PacketConn
connChan chan net.Conn
cqueue chan net.Conn
errChan chan error
closed chan struct{}
connPool *connPool
connPool *udp.ConnPool
logger logger.Logger
}
@ -53,8 +54,8 @@ func (l *udpListener) Init(md md.Metadata) (err error) {
return
}
l.connChan = make(chan net.Conn, l.md.connQueueSize)
l.connPool = newConnPool(l.md.ttl).WithLogger(l.logger)
l.cqueue = make(chan net.Conn, l.md.backlog)
l.connPool = udp.NewConnPool(l.md.ttl).WithLogger(l.logger)
go l.listenLoop()
@ -64,7 +65,7 @@ func (l *udpListener) Init(md md.Metadata) (err error) {
func (l *udpListener) Accept() (conn net.Conn, err error) {
var ok bool
select {
case conn = <-l.connChan:
case conn = <-l.cqueue:
case err, ok = <-l.errChan:
if !ok {
err = listener.ErrClosed
@ -106,18 +107,18 @@ func (l *udpListener) listenLoop() {
continue
}
if err := c.Queue(b[:n]); err != nil {
if err := c.WriteQueue(b[:n]); err != nil {
l.logger.Warn("data discarded: ", err)
}
}
}
func (l *udpListener) getConn(addr net.Addr) *conn {
func (l *udpListener) getConn(addr net.Addr) *udp.Conn {
c, ok := l.connPool.Get(addr.String())
if !ok {
c = newConn(l.conn, addr, l.md.readQueueSize)
c = udp.NewConn(l.conn, l.conn.LocalAddr(), addr, l.md.readQueueSize)
select {
case l.connChan <- c:
case l.cqueue <- c:
l.connPool.Set(addr.String(), c)
default:
c.Close()

View File

@ -10,7 +10,7 @@ const (
defaultTTL = 60 * time.Second
defaultReadBufferSize = 4096
defaultReadQueueSize = 128
defaultConnQueueSize = 128
defaultBacklog = 128
)
type metadata struct {
@ -18,7 +18,7 @@ type metadata struct {
readBufferSize int
readQueueSize int
connQueueSize int
backlog int
}
func (l *udpListener) parseMetadata(md md.Metadata) (err error) {
@ -26,7 +26,7 @@ func (l *udpListener) parseMetadata(md md.Metadata) (err error) {
ttl = "ttl"
readBufferSize = "readBufferSize"
readQueueSize = "readQueueSize"
connQueueSize = "connQueueSize"
backlog = "backlog"
)
l.md.ttl = md.GetDuration(ttl)
@ -43,9 +43,9 @@ func (l *udpListener) parseMetadata(md md.Metadata) (err error) {
l.md.readQueueSize = defaultReadQueueSize
}
l.md.connQueueSize = md.GetInt(connQueueSize)
if l.md.connQueueSize <= 0 {
l.md.connQueueSize = defaultConnQueueSize
l.md.backlog = md.GetInt(backlog)
if l.md.backlog <= 0 {
l.md.backlog = defaultBacklog
}
return