initial commit

listener/ftcp/conn.go

@@ -0,0 +1,124 @@
+package ftcp
+
+import (
+	"errors"
+	"net"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// serverConn is a server side connection for UDP client peer, it implements net.Conn and net.PacketConn.
+type serverConn struct {
+	pc         net.PacketConn
+	raddr      net.Addr
+	rc         chan []byte // data receive queue
+	fresh      int32
+	closed     chan struct{}
+	closeMutex sync.Mutex
+	config     *serverConnConfig
+}
+
+type serverConnConfig struct {
+	ttl     time.Duration
+	qsize   int
+	onClose func()
+}
+
+func newServerConn(conn net.PacketConn, raddr net.Addr, cfg *serverConnConfig) *serverConn {
+	if conn == nil || raddr == nil {
+		return nil
+	}
+
+	if cfg == nil {
+		cfg = &serverConnConfig{}
+	}
+	c := &serverConn{
+		pc:     conn,
+		raddr:  raddr,
+		rc:     make(chan []byte, cfg.qsize),
+		closed: make(chan struct{}),
+		config: cfg,
+	}
+	go c.ttlWait()
+	return c
+}
+
+func (c *serverConn) send(b []byte) error {
+	select {
+	case c.rc <- b:
+		return nil
+	default:
+		return errors.New("queue is full")
+	}
+}
+
+func (c *serverConn) Read(b []byte) (n int, err error) {
+	select {
+	case bb := <-c.rc:
+		n = copy(b, bb)
+		atomic.StoreInt32(&c.fresh, 1)
+	case <-c.closed:
+		err = errors.New("read from closed connection")
+		return
+	}
+
+	return
+}
+
+func (c *serverConn) Write(b []byte) (n int, err error) {
+	return c.pc.WriteTo(b, c.raddr)
+}
+
+func (c *serverConn) Close() error {
+	c.closeMutex.Lock()
+	defer c.closeMutex.Unlock()
+
+	select {
+	case <-c.closed:
+		return errors.New("connection is closed")
+	default:
+		if c.config.onClose != nil {
+			c.config.onClose()
+		}
+		close(c.closed)
+	}
+	return nil
+}
+
+func (c *serverConn) LocalAddr() net.Addr {
+	return c.pc.LocalAddr()
+}
+
+func (c *serverConn) RemoteAddr() net.Addr {
+	return c.raddr
+}
+
+func (c *serverConn) SetDeadline(t time.Time) error {
+	return c.pc.SetDeadline(t)
+}
+
+func (c *serverConn) SetReadDeadline(t time.Time) error {
+	return c.pc.SetReadDeadline(t)
+}
+
+func (c *serverConn) SetWriteDeadline(t time.Time) error {
+	return c.pc.SetWriteDeadline(t)
+}
+
+func (c *serverConn) ttlWait() {
+	ticker := time.NewTicker(c.config.ttl)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			if !atomic.CompareAndSwapInt32(&c.fresh, 1, 0) {
+				c.Close()
+				return
+			}
+		case <-c.closed:
+			return
+		}
+	}
+}

listener/ftcp/listener.go

@@ -0,0 +1,159 @@
+package ftcp
+
+import (
+	"net"
+	"sync"
+	"sync/atomic"
+
+	"github.com/go-gost/gost/v3/pkg/common/metrics"
+	"github.com/go-gost/gost/v3/pkg/listener"
+	"github.com/go-gost/gost/v3/pkg/logger"
+	md "github.com/go-gost/gost/v3/pkg/metadata"
+	"github.com/go-gost/gost/v3/pkg/registry"
+	"github.com/xtaci/tcpraw"
+)
+
+func init() {
+	registry.ListenerRegistry().Register("ftcp", NewListener)
+}
+
+type ftcpListener struct {
+	conn     net.PacketConn
+	connChan chan net.Conn
+	errChan  chan error
+	connPool connPool
+	logger   logger.Logger
+	md       metadata
+	options  listener.Options
+}
+
+func NewListener(opts ...listener.Option) listener.Listener {
+	options := listener.Options{}
+	for _, opt := range opts {
+		opt(&options)
+	}
+	return &ftcpListener{
+		logger:  options.Logger,
+		options: options,
+	}
+}
+
+func (l *ftcpListener) Init(md md.Metadata) (err error) {
+	if err = l.parseMetadata(md); err != nil {
+		return
+	}
+
+	l.conn, err = tcpraw.Listen("tcp", l.options.Addr)
+	if err != nil {
+		return
+	}
+
+	l.connChan = make(chan net.Conn, l.md.connQueueSize)
+	l.errChan = make(chan error, 1)
+
+	go l.listenLoop()
+
+	return
+}
+
+func (l *ftcpListener) Accept() (conn net.Conn, err error) {
+	var ok bool
+	select {
+	case conn = <-l.connChan:
+		conn = metrics.WrapConn(l.options.Service, conn)
+	case err, ok = <-l.errChan:
+		if !ok {
+			err = listener.ErrClosed
+		}
+	}
+	return
+}
+
+func (l *ftcpListener) Close() error {
+	err := l.conn.Close()
+	l.connPool.Range(func(k any, v *serverConn) bool {
+		v.Close()
+		return true
+	})
+	return err
+}
+
+func (l *ftcpListener) Addr() net.Addr {
+	return l.conn.LocalAddr()
+}
+
+func (l *ftcpListener) listenLoop() {
+	for {
+		b := make([]byte, l.md.readBufferSize)
+
+		n, raddr, err := l.conn.ReadFrom(b)
+		if err != nil {
+			l.logger.Error("accept:", err)
+			l.errChan <- err
+			close(l.errChan)
+			return
+		}
+
+		conn, ok := l.connPool.Get(raddr.String())
+		if !ok {
+			conn = newServerConn(l.conn, raddr,
+				&serverConnConfig{
+					ttl:   l.md.ttl,
+					qsize: l.md.readQueueSize,
+					onClose: func() {
+						l.connPool.Delete(raddr.String())
+					},
+				})
+
+			select {
+			case l.connChan <- conn:
+				l.connPool.Set(raddr.String(), conn)
+			default:
+				conn.Close()
+				l.logger.Error("connection queue is full")
+			}
+		}
+
+		if err := conn.send(b[:n]); err != nil {
+			l.logger.Warn("data discarded:", err)
+		}
+		l.logger.Debug("recv", n)
+	}
+}
+
+func (l *ftcpListener) parseMetadata(md md.Metadata) (err error) {
+	return
+}
+
+type connPool struct {
+	size int64
+	m    sync.Map
+}
+
+func (p *connPool) Get(key any) (conn *serverConn, ok bool) {
+	v, ok := p.m.Load(key)
+	if ok {
+		conn, ok = v.(*serverConn)
+	}
+	return
+}
+
+func (p *connPool) Set(key any, conn *serverConn) {
+	p.m.Store(key, conn)
+	atomic.AddInt64(&p.size, 1)
+}
+
+func (p *connPool) Delete(key any) {
+	p.m.Delete(key)
+	atomic.AddInt64(&p.size, -1)
+}
+
+func (p *connPool) Range(f func(key any, value *serverConn) bool) {
+	p.m.Range(func(k, v any) bool {
+		return f(k, v.(*serverConn))
+	})
+}
+
+func (p *connPool) Size() int64 {
+	return atomic.LoadInt64(&p.size)
+}

listener/ftcp/metadata.go

@@ -0,0 +1,22 @@
+package ftcp
+
+import "time"
+
+const (
+	defaultTTL            = 60 * time.Second
+	defaultReadBufferSize = 1024
+	defaultReadQueueSize  = 128
+	defaultConnQueueSize  = 128
+)
+
+const (
+	addr = "addr"
+)
+
+type metadata struct {
+	ttl time.Duration
+
+	readBufferSize int
+	readQueueSize  int
+	connQueueSize  int
+}