initial commit

common/util/udp/conn.go

@@ -0,0 +1,102 @@
+package udp
+
+import (
+	"errors"
+	"net"
+	"sync"
+	"sync/atomic"
+
+	"github.com/go-gost/core/common/bufpool"
+)
+
+// Conn is a server side connection for UDP client peer, it implements net.Conn and net.PacketConn.
+type Conn struct {
+	net.PacketConn
+	localAddr  net.Addr
+	remoteAddr net.Addr
+	rc         chan []byte // data receive queue
+	idle       int32       // indicate the connection is idle
+	closed     chan struct{}
+	closeMutex sync.Mutex
+}
+
+func NewConn(c net.PacketConn, localAddr, remoteAddr net.Addr, queueSize int) *Conn {
+	return &Conn{
+		PacketConn: c,
+		localAddr:  localAddr,
+		remoteAddr: remoteAddr,
+		rc:         make(chan []byte, queueSize),
+		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) LocalAddr() net.Addr {
+	return c.localAddr
+}
+
+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) WriteQueue(b []byte) error {
+	select {
+	case c.rc <- b:
+		return nil
+
+	case <-c.closed:
+		return net.ErrClosed
+
+	default:
+		return errors.New("recv queue is full")
+	}
+}

common/util/udp/listener.go

@@ -0,0 +1,120 @@
+package udp
+
+import (
+	"net"
+	"sync"
+	"time"
+
+	"github.com/go-gost/core/common/bufpool"
+	"github.com/go-gost/core/logger"
+)
+
+type listener struct {
+	addr           net.Addr
+	conn           net.PacketConn
+	cqueue         chan net.Conn
+	readQueueSize  int
+	readBufferSize int
+	connPool       *ConnPool
+	mux            sync.Mutex
+	closed         chan struct{}
+	errChan        chan error
+	logger         logger.Logger
+}
+
+func NewListener(conn net.PacketConn, addr net.Addr, backlog, dataQueueSize, dataBufferSize int, ttl time.Duration, logger logger.Logger) net.Listener {
+	ln := &listener{
+		conn:           conn,
+		addr:           addr,
+		cqueue:         make(chan net.Conn, backlog),
+		connPool:       NewConnPool(ttl).WithLogger(logger),
+		readQueueSize:  dataQueueSize,
+		readBufferSize: dataBufferSize,
+		closed:         make(chan struct{}),
+		errChan:        make(chan error, 1),
+		logger:         logger,
+	}
+	go ln.listenLoop()
+
+	return ln
+}
+
+func (ln *listener) Accept() (conn net.Conn, err error) {
+	select {
+	case conn = <-ln.cqueue:
+		return
+	case <-ln.closed:
+		return nil, net.ErrClosed
+	case err = <-ln.errChan:
+		if err == nil {
+			err = net.ErrClosed
+		}
+		return
+	}
+}
+
+func (ln *listener) listenLoop() {
+	for {
+		select {
+		case <-ln.closed:
+			return
+		default:
+		}
+
+		b := bufpool.Get(ln.readBufferSize)
+
+		n, raddr, err := ln.conn.ReadFrom(*b)
+		if err != nil {
+			ln.errChan <- err
+			close(ln.errChan)
+			return
+		}
+
+		c := ln.getConn(raddr)
+		if c == nil {
+			bufpool.Put(b)
+			continue
+		}
+
+		if err := c.WriteQueue((*b)[:n]); err != nil {
+			ln.logger.Warn("data discarded: ", err)
+		}
+	}
+}
+
+func (ln *listener) Addr() net.Addr {
+	return ln.addr
+}
+
+func (ln *listener) Close() error {
+	select {
+	case <-ln.closed:
+	default:
+		close(ln.closed)
+		ln.conn.Close()
+		ln.connPool.Close()
+	}
+
+	return nil
+}
+
+func (ln *listener) getConn(raddr net.Addr) *Conn {
+	ln.mux.Lock()
+	defer ln.mux.Unlock()
+
+	c, ok := ln.connPool.Get(raddr.String())
+	if ok {
+		return c
+	}
+
+	c = NewConn(ln.conn, ln.addr, raddr, ln.readQueueSize)
+	select {
+	case ln.cqueue <- c:
+		ln.connPool.Set(raddr.String(), c)
+		return c
+	default:
+		c.Close()
+		ln.logger.Warnf("connection queue is full, client %s discarded", raddr)
+		return nil
+	}
+}

common/util/udp/pool.go

@@ -0,0 +1,100 @@
+package udp
+
+import (
+	"sync"
+	"time"
+
+	"github.com/go-gost/core/logger"
+)
+
+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 any) (c *Conn, ok bool) {
+	v, ok := p.m.Load(key)
+	if ok {
+		c, ok = v.(*Conn)
+	}
+	return
+}
+
+func (p *ConnPool) Set(key any, c *Conn) {
+	p.m.Store(key, c)
+}
+
+func (p *ConnPool) Delete(key any) {
+	p.m.Delete(key)
+}
+
+func (p *ConnPool) Close() {
+	select {
+	case <-p.closed:
+		return
+	default:
+	}
+
+	close(p.closed)
+
+	p.m.Range(func(k, v any) 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 any) 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
+		}
+	}
+}