mirror of
https://github.com/Sagit-chu/flvx.git
synced 2026-10-12 20:26:37 +08:00
init
This commit is contained in:
qaq
@@ -0,0 +1,29 @@
|
||||
package udp
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net"
|
||||
|
||||
xnet "github.com/go-gost/x/internal/net"
|
||||
)
|
||||
|
||||
type Conn interface {
|
||||
net.PacketConn
|
||||
io.Reader
|
||||
io.Writer
|
||||
ReadUDP
|
||||
WriteUDP
|
||||
xnet.SetBuffer
|
||||
xnet.SyscallConn
|
||||
xnet.RemoteAddr
|
||||
}
|
||||
|
||||
type ReadUDP interface {
|
||||
ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error)
|
||||
ReadMsgUDP(b, oob []byte) (n, oobn, flags int, addr *net.UDPAddr, err error)
|
||||
}
|
||||
|
||||
type WriteUDP interface {
|
||||
WriteToUDP(b []byte, addr *net.UDPAddr) (int, error)
|
||||
WriteMsgUDP(b, oob []byte, addr *net.UDPAddr) (n, oobn int, err error)
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
package udp
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/go-gost/core/common/bufpool"
|
||||
"github.com/go-gost/core/logger"
|
||||
)
|
||||
|
||||
type ListenConfig struct {
|
||||
Addr net.Addr
|
||||
Backlog int
|
||||
ReadQueueSize int
|
||||
ReadBufferSize int
|
||||
TTL time.Duration
|
||||
Keepalive bool
|
||||
Logger logger.Logger
|
||||
}
|
||||
type listener struct {
|
||||
conn net.PacketConn
|
||||
cqueue chan net.Conn
|
||||
connPool *connPool
|
||||
closed chan struct{}
|
||||
errChan chan error
|
||||
config *ListenConfig
|
||||
}
|
||||
|
||||
func NewListener(conn net.PacketConn, cfg *ListenConfig) net.Listener {
|
||||
if cfg == nil {
|
||||
cfg = &ListenConfig{}
|
||||
}
|
||||
|
||||
ln := &listener{
|
||||
conn: conn,
|
||||
cqueue: make(chan net.Conn, cfg.Backlog),
|
||||
closed: make(chan struct{}),
|
||||
errChan: make(chan error, 1),
|
||||
config: cfg,
|
||||
}
|
||||
ln.connPool = newConnPool(cfg.TTL).WithLogger(cfg.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.config.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.config.Logger.Warn("data discarded: ", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (ln *listener) Addr() net.Addr {
|
||||
if ln.config.Addr != nil {
|
||||
return ln.config.Addr
|
||||
}
|
||||
return ln.conn.LocalAddr()
|
||||
}
|
||||
|
||||
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 {
|
||||
c, ok := ln.connPool.Get(raddr.String())
|
||||
if ok && !c.isClosed() {
|
||||
return c
|
||||
}
|
||||
|
||||
c = newConn(ln.conn, ln.Addr(), raddr, ln.config.ReadQueueSize, ln.config.Keepalive)
|
||||
select {
|
||||
case ln.cqueue <- c:
|
||||
ln.connPool.Set(raddr.String(), c)
|
||||
return c
|
||||
default:
|
||||
c.Close()
|
||||
ln.config.Logger.Warnf("connection queue is full, client %s discarded", raddr)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
keepalive bool
|
||||
}
|
||||
|
||||
func newConn(c net.PacketConn, laddr, remoteAddr net.Addr, queueSize int, keepalive bool) *conn {
|
||||
return &conn{
|
||||
PacketConn: c,
|
||||
localAddr: laddr,
|
||||
remoteAddr: remoteAddr,
|
||||
rc: make(chan []byte, queueSize),
|
||||
closed: make(chan struct{}),
|
||||
keepalive: keepalive,
|
||||
}
|
||||
}
|
||||
|
||||
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) WriteTo(b []byte, addr net.Addr) (n int, err error) {
|
||||
if !c.keepalive {
|
||||
defer c.Close()
|
||||
}
|
||||
return c.PacketConn.WriteTo(b, addr)
|
||||
}
|
||||
|
||||
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) isClosed() bool {
|
||||
select {
|
||||
case <-c.closed:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
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) {
|
||||
if p == nil {
|
||||
return
|
||||
}
|
||||
|
||||
v, ok := p.m.Load(key)
|
||||
if ok {
|
||||
c, ok = v.(*conn)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (p *connPool) Set(key any, c *conn) {
|
||||
if p == nil {
|
||||
return
|
||||
}
|
||||
|
||||
p.m.Store(key, c)
|
||||
}
|
||||
|
||||
func (p *connPool) Delete(key any) {
|
||||
if p == nil {
|
||||
return
|
||||
}
|
||||
p.m.Delete(key)
|
||||
}
|
||||
|
||||
func (p *connPool) Close() {
|
||||
if p == nil {
|
||||
return
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package udp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"net"
|
||||
|
||||
"github.com/go-gost/core/bypass"
|
||||
"github.com/go-gost/core/logger"
|
||||
)
|
||||
|
||||
const (
|
||||
MaxMessageSize = math.MaxUint16
|
||||
)
|
||||
|
||||
type Relay struct {
|
||||
pc1 net.PacketConn
|
||||
pc2 net.PacketConn
|
||||
|
||||
bypass bypass.Bypass
|
||||
logger logger.Logger
|
||||
}
|
||||
|
||||
func NewRelay(pc1, pc2 net.PacketConn) *Relay {
|
||||
return &Relay{
|
||||
pc1: pc1,
|
||||
pc2: pc2,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Relay) WithBypass(bp bypass.Bypass) *Relay {
|
||||
r.bypass = bp
|
||||
return r
|
||||
}
|
||||
|
||||
func (r *Relay) WithLogger(logger logger.Logger) *Relay {
|
||||
r.logger = logger
|
||||
return r
|
||||
}
|
||||
|
||||
func (r *Relay) Run(ctx context.Context) (err error) {
|
||||
errc := make(chan error, 2)
|
||||
|
||||
go func() {
|
||||
var b [MaxMessageSize]byte
|
||||
for {
|
||||
err := func() error {
|
||||
n, raddr, err := r.pc1.ReadFrom(b[:])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if r.bypass != nil && r.bypass.Contains(ctx, "udp", raddr.String()) {
|
||||
if r.logger != nil {
|
||||
r.logger.Warn("bypass: ", raddr)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
if _, err := r.pc2.WriteTo(b[:n], raddr); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if r.logger != nil {
|
||||
r.logger.Tracef("%s >>> %s data: %d",
|
||||
r.pc2.LocalAddr(), raddr, n)
|
||||
|
||||
}
|
||||
|
||||
return nil
|
||||
}()
|
||||
|
||||
if err != nil {
|
||||
errc <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
go func() {
|
||||
var b [MaxMessageSize]byte
|
||||
for {
|
||||
err := func() error {
|
||||
n, raddr, err := r.pc2.ReadFrom(b[:])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if r.bypass != nil && r.bypass.Contains(ctx, "udp", raddr.String()) {
|
||||
if r.logger != nil {
|
||||
r.logger.Warn("bypass: ", raddr)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
if _, err := r.pc1.WriteTo(b[:n], raddr); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if r.logger != nil {
|
||||
r.logger.Tracef("%s <<< %s data: %d",
|
||||
r.pc2.LocalAddr(), raddr, n)
|
||||
|
||||
}
|
||||
|
||||
return nil
|
||||
}()
|
||||
|
||||
if err != nil {
|
||||
errc <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return <-errc
|
||||
}
|
||||
Reference in New Issue
Block a user