From 4b856214f71a2d744d6ec43de9bc40a3827aa3a2 Mon Sep 17 00:00:00 2001 From: ginuerzh Date: Sat, 1 Feb 2020 06:49:51 +0000 Subject: [PATCH] update udp transparent proxy --- forward.go | 3 +- redirect.go | 160 +++++++++++++++++++++------------------------------- udp.go | 5 +- 3 files changed, 68 insertions(+), 100 deletions(-) diff --git a/forward.go b/forward.go index 326f87e..d7e93f8 100644 --- a/forward.go +++ b/forward.go @@ -610,7 +610,7 @@ func (l *tcpRemoteForwardListener) Close() error { type udpRemoteForwardListener struct { addr net.Addr chain *Chain - connMap udpConnMap + connMap *udpConnMap connChan chan net.Conn ln *net.UDPConn errChan chan error @@ -640,6 +640,7 @@ func UDPRemoteForwardListener(addr string, chain *Chain, cfg *UDPListenConfig) ( ln := &udpRemoteForwardListener{ addr: laddr, chain: chain, + connMap: new(udpConnMap), connChan: make(chan net.Conn, backlog), errChan: make(chan error, 1), closed: make(chan struct{}), diff --git a/redirect.go b/redirect.go index afac627..b2979b0 100644 --- a/redirect.go +++ b/redirect.go @@ -6,7 +6,9 @@ import ( "errors" "fmt" "net" + "sync" "syscall" + "time" "github.com/LiamHaworth/go-tproxy" "github.com/go-log/log" @@ -121,17 +123,15 @@ func (h *udpRedirectHandler) Init(options ...HandlerOption) { } } -func (h *udpRedirectHandler) Handle(c net.Conn) { - defer c.Close() +func (h *udpRedirectHandler) Handle(conn net.Conn) { + defer conn.Close() - conn, ok := c.(*udpRedirectServerConn) + raddr, ok := conn.LocalAddr().(*net.UDPAddr) if !ok { - log.Log("wrong connection type") + log.Log("[red-udp] wrong connection type") return } - raddr := conn.DstAddr() - var cc net.Conn var err error if h.options.Chain.IsEmpty() { @@ -167,11 +167,8 @@ func (h *udpRedirectHandler) Handle(c net.Conn) { } type udpRedirectListener struct { - ln *net.UDPConn - connChan chan net.Conn - errChan chan error - connMap udpConnMap - config *UDPListenConfig + *net.UDPConn + config *UDPListenConfig } // UDPRedirectListener creates a Listener for UDP transparent proxy server. @@ -189,103 +186,72 @@ func UDPRedirectListener(addr string, cfg *UDPListenConfig) (Listener, error) { if cfg == nil { cfg = &UDPListenConfig{} } - - backlog := cfg.Backlog - if backlog <= 0 { - backlog = defaultBacklog - } - - l := &udpRedirectListener{ - ln: ln, - connChan: make(chan net.Conn, backlog), - errChan: make(chan error, 1), - config: cfg, - } - go l.listenLoop() - return l, nil -} - -func (l *udpRedirectListener) listenLoop() { - for { - b := make([]byte, mediumBufferSize) - n, raddr, dstAddr, err := tproxy.ReadFromUDP(l.ln, b) - if err != nil { - log.Logf("[red-udp] peer -> %s : %s", l.Addr(), err) - l.Close() - l.errChan <- err - close(l.errChan) - return - } - - conn, ok := l.connMap.Get(raddr.String()) - if !ok { - conn = newUDPServerConn(l.ln, raddr, &udpServerConnConfig{ - ttl: l.config.TTL, - qsize: l.config.QueueSize, - onClose: func() { - l.connMap.Delete(raddr.String()) - log.Logf("[red-udp] %s closed (%d)", raddr, l.connMap.Size()) - }, - }) - - cc := udpRedirectServerConn{ - udpServerConn: conn, - dstAddr: dstAddr, - } - select { - case l.connChan <- cc: - l.connMap.Set(raddr.String(), conn) - log.Logf("[red-udp] %s -> %s (%d)", raddr, l.Addr(), l.connMap.Size()) - default: - conn.Close() - log.Logf("[red-udp] %s - %s: connection queue is full (%d)", - raddr, l.Addr(), cap(l.connChan)) - } - } - - select { - case conn.rChan <- b[:n]: - if Debug { - log.Logf("[red-udp] %s >>> %s : length %d", raddr, l.Addr(), n) - } - default: - log.Logf("[red-udp] %s -> %s : recv queue is full (%d)", - raddr, l.Addr(), cap(conn.rChan)) - } - } + return &udpRedirectListener{ + UDPConn: ln, + config: cfg, + }, nil } func (l *udpRedirectListener) Accept() (conn net.Conn, err error) { - var ok bool - select { - case conn = <-l.connChan: - case err, ok = <-l.errChan: - if !ok { - err = errors.New("accpet on closed listener") - } + b := make([]byte, mediumBufferSize) + + n, raddr, dstAddr, err := tproxy.ReadFromUDP(l.UDPConn, b) + if err != nil { + log.Logf("[red-udp] %s : %s", l.Addr(), err) + return + } + log.Logf("[red-udp] %s: %s -> %s", l.Addr(), raddr, dstAddr) + + c, err := tproxy.DialUDP("udp", dstAddr, raddr) + if err != nil { + log.Logf("[red-udp] %s -> %s : %s", raddr, dstAddr, err) + return + } + + ttl := l.config.TTL + if ttl <= 0 { + ttl = defaultTTL + } + + conn = &udpRedirectServerConn{ + Conn: c, + buf: b[:n], + ttl: ttl, } return } func (l *udpRedirectListener) Addr() net.Addr { - return l.ln.LocalAddr() -} - -func (l *udpRedirectListener) Close() error { - err := l.ln.Close() - l.connMap.Range(func(k interface{}, v *udpServerConn) bool { - v.Close() - return true - }) - - return err + return l.UDPConn.LocalAddr() } type udpRedirectServerConn struct { - *udpServerConn - dstAddr *net.UDPAddr + net.Conn + buf []byte + ttl time.Duration + once sync.Once } -func (c *udpRedirectServerConn) DstAddr() *net.UDPAddr { - return c.dstAddr +func (c *udpRedirectServerConn) Read(b []byte) (n int, err error) { + if c.ttl > 0 { + c.SetReadDeadline(time.Now().Add(c.ttl)) + defer c.SetReadDeadline(time.Time{}) + } + c.once.Do(func() { + n = copy(b, c.buf) + c.buf = nil + }) + + if n == 0 { + n, err = c.Conn.Read(b) + } + return +} + +func (c *udpRedirectServerConn) Write(b []byte) (n int, err error) { + if c.ttl > 0 { + c.SetWriteDeadline(time.Now().Add(c.ttl)) + defer c.SetWriteDeadline(time.Time{}) + } + return c.Conn.Write(b) } diff --git a/udp.go b/udp.go index b17e55a..34088ea 100644 --- a/udp.go +++ b/udp.go @@ -54,7 +54,7 @@ type udpListener struct { ln net.PacketConn connChan chan net.Conn errChan chan error - connMap udpConnMap + connMap *udpConnMap config *UDPListenConfig } @@ -82,6 +82,7 @@ func UDPListener(addr string, cfg *UDPListenConfig) (Listener, error) { ln: ln, connChan: make(chan net.Conn, backlog), errChan: make(chan error, 1), + connMap: new(udpConnMap), config: cfg, } go l.listenLoop() @@ -159,8 +160,8 @@ func (l *udpListener) Close() error { } type udpConnMap struct { - m sync.Map size int64 + m sync.Map } func (m *udpConnMap) Get(key interface{}) (conn *udpServerConn, ok bool) {