fix udp bind

pkg/listener/rtcp/listener.go

@@ -17,13 +17,13 @@ func init() {
 }
 
 type rtcpListener struct {
-	addr     string
-	laddr    net.Addr
-	chain    *chain.Chain
-	accepter connector.Accepter
-	md       metadata
-	logger   logger.Logger
-	closed   chan struct{}
+	addr   string
+	laddr  net.Addr
+	chain  *chain.Chain
+	ln     net.Listener
+	md     metadata
+	logger logger.Logger
+	closed chan struct{}
 }
 
 func NewListener(opts ...listener.Option) listener.Listener {
@@ -58,6 +58,34 @@ func (l *rtcpListener) Init(md md.Metadata) (err error) {
 	return
 }
 
+func (l *rtcpListener) Accept() (conn net.Conn, err error) {
+	select {
+	case <-l.closed:
+		return nil, net.ErrClosed
+	default:
+	}
+
+	if l.ln == nil {
+		r := (&chain.Router{}).
+			WithChain(l.chain).
+			WithRetry(l.md.retryCount).
+			WithLogger(l.logger)
+		l.ln, err = r.Bind(context.Background(), "tcp", l.laddr.String(),
+			connector.MuxBindOption(true),
+		)
+		if err != nil {
+			return nil, connector.NewAcceptError(err)
+		}
+	}
+	conn, err = l.ln.Accept()
+	if err != nil {
+		l.ln.Close()
+		l.ln = nil
+		return nil, connector.NewAcceptError(err)
+	}
+	return
+}
+
 func (l *rtcpListener) Addr() net.Addr {
 	return l.laddr
 }
@@ -67,27 +95,11 @@ func (l *rtcpListener) Close() error {
 	case <-l.closed:
 	default:
 		close(l.closed)
+		if l.ln != nil {
+			l.ln.Close()
+			l.ln = nil
+		}
 	}
 
 	return nil
 }
-
-func (l *rtcpListener) Accept() (conn net.Conn, err error) {
-	if l.accepter == nil {
-		r := (&chain.Router{}).
-			WithChain(l.chain).
-			WithRetry(l.md.retryCount).
-			WithLogger(l.logger)
-		l.accepter, err = r.Bind(context.Background(), "tcp", l.laddr.String())
-		if err != nil {
-			return nil, connector.NewAcceptError(err)
-		}
-	}
-	conn, err = l.accepter.Accept()
-	if err != nil {
-		l.accepter.Close()
-		l.accepter = nil
-		return nil, connector.NewAcceptError(err)
-	}
-	return
-}

pkg/listener/rudp/listener.go

@@ -2,15 +2,10 @@ 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/connector"
 	"github.com/go-gost/gost/pkg/listener"
 	"github.com/go-gost/gost/pkg/logger"
 	md "github.com/go-gost/gost/pkg/metadata"
@@ -22,14 +17,13 @@ func init() {
 }
 
 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
+	addr   string
+	laddr  *net.UDPAddr
+	chain  *chain.Chain
+	ln     net.Listener
+	md     metadata
+	logger logger.Logger
+	closed chan struct{}
 }
 
 func NewListener(opts ...listener.Option) listener.Listener {
@@ -60,21 +54,43 @@ func (l *rudpListener) Init(md md.Metadata) (err error) {
 	}
 
 	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
+		return nil, net.ErrClosed
+	default:
 	}
+
+	if l.ln == nil {
+		r := (&chain.Router{}).
+			WithChain(l.chain).
+			WithRetry(l.md.retryCount).
+			WithLogger(l.logger)
+		l.ln, err = r.Bind(context.Background(), "udp", l.laddr.String(),
+			connector.BacklogBindOption(l.md.backlog),
+			connector.UDPConnTTLBindOption(l.md.ttl),
+			connector.UDPDataBufferSizeBindOption(l.md.readBufferSize),
+			connector.UDPDataQueueSizeBindOption(l.md.readQueueSize),
+		)
+		if err != nil {
+			return nil, connector.NewAcceptError(err)
+		}
+	}
+	conn, err = l.ln.Accept()
+	if err != nil {
+		l.ln.Close()
+		l.ln = nil
+		return nil, connector.NewAcceptError(err)
+	}
+	return
+}
+
+func (l *rudpListener) Addr() net.Addr {
+	return l.laddr
 }
 
 func (l *rudpListener) Close() error {
@@ -82,136 +98,11 @@ func (l *rudpListener) Close() error {
 	case <-l.closed:
 	default:
 		close(l.closed)
-		l.connPool.Close()
+		if l.ln != nil {
+			l.ln.Close()
+			l.ln = nil
+		}
 	}
 
 	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.UDPTunClientPacketConn(conn), 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
-}

pkg/listener/rudp/metadata.go

@@ -7,8 +7,8 @@ import (
 )
 
 const (
-	defaultTTL            = 60 * time.Second
-	defaultReadBufferSize = 4096
+	defaultTTL            = 5 * time.Second
+	defaultReadBufferSize = 1024
 	defaultReadQueueSize  = 128
 	defaultBacklog        = 128
 )

pkg/listener/udp/listener.go

@@ -3,7 +3,6 @@ package udp
 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"
@@ -16,14 +15,10 @@ func init() {
 }
 
 type udpListener struct {
-	addr     string
-	md       metadata
-	conn     net.PacketConn
-	cqueue   chan net.Conn
-	errChan  chan error
-	closed   chan struct{}
-	connPool *udp.ConnPool
-	logger   logger.Logger
+	addr string
+	md   metadata
+	net.Listener
+	logger logger.Logger
 }
 
 func NewListener(opts ...listener.Option) listener.Listener {
@@ -32,10 +27,8 @@ func NewListener(opts ...listener.Option) listener.Listener {
 		opt(options)
 	}
 	return &udpListener{
-		addr:    options.Addr,
-		errChan: make(chan error, 1),
-		closed:  make(chan struct{}),
-		logger:  options.Logger,
+		addr:   options.Addr,
+		logger: options.Logger,
 	}
 }
 
@@ -49,82 +42,15 @@ func (l *udpListener) Init(md md.Metadata) (err error) {
 		return
 	}
 
-	l.conn, err = net.ListenUDP("udp", laddr)
+	conn, err := net.ListenUDP("udp", laddr)
 	if err != nil {
 		return
 	}
 
-	l.cqueue = make(chan net.Conn, l.md.backlog)
-	l.connPool = udp.NewConnPool(l.md.ttl).WithLogger(l.logger)
-
-	go l.listenLoop()
-
+	l.Listener = udp.NewListener(conn, laddr,
+		l.md.backlog,
+		l.md.readQueueSize, l.md.readBufferSize,
+		l.md.ttl,
+		l.logger)
 	return
 }
-
-func (l *udpListener) Accept() (conn net.Conn, err error) {
-	var ok bool
-	select {
-	case conn = <-l.cqueue:
-	case err, ok = <-l.errChan:
-		if !ok {
-			err = listener.ErrClosed
-		}
-	}
-	return
-}
-
-func (l *udpListener) Close() error {
-	select {
-	case <-l.closed:
-	default:
-		close(l.closed)
-		l.connPool.Close()
-		return l.conn.Close()
-	}
-
-	return nil
-}
-
-func (l *udpListener) Addr() net.Addr {
-	return l.conn.LocalAddr()
-}
-
-func (l *udpListener) listenLoop() {
-	for {
-		b := bufpool.Get(l.md.readBufferSize)
-
-		n, raddr, err := l.conn.ReadFrom(b)
-		if err != nil {
-			l.errChan <- err
-			close(l.errChan)
-			return
-		}
-
-		c := l.getConn(raddr)
-		if c == nil {
-			bufpool.Put(b)
-			continue
-		}
-
-		if err := c.WriteQueue(b[:n]); err != nil {
-			l.logger.Warn("data discarded: ", err)
-		}
-	}
-}
-
-func (l *udpListener) getConn(addr net.Addr) *udp.Conn {
-	c, ok := l.connPool.Get(addr.String())
-	if !ok {
-		c = udp.NewConn(l.conn, l.conn.LocalAddr(), addr, l.md.readQueueSize)
-		select {
-		case l.cqueue <- c:
-			l.connPool.Set(addr.String(), c)
-		default:
-			c.Close()
-			l.logger.Warnf("connection queue is full, client %s discarded", addr.String())
-			return nil
-		}
-	}
-	return c
-}

pkg/listener/udp/metadata.go

@@ -7,8 +7,8 @@ import (
 )
 
 const (
-	defaultTTL            = 60 * time.Second
-	defaultReadBufferSize = 4096
+	defaultTTL            = 5 * time.Second
+	defaultReadBufferSize = 1024
 	defaultReadQueueSize  = 128
 	defaultBacklog        = 128
 )