update udp transparent proxy

This commit is contained in:
ginuerzh 2020-02-01 06:49:51 +00:00
parent bd9fc76466
commit 4b856214f7
3 changed files with 68 additions and 100 deletions

View File

@ -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{}),

View File

@ -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,10 +167,7 @@ func (h *udpRedirectHandler) Handle(c net.Conn) {
}
type udpRedirectListener struct {
ln *net.UDPConn
connChan chan net.Conn
errChan chan error
connMap udpConnMap
*net.UDPConn
config *UDPListenConfig
}
@ -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),
return &udpRedirectListener{
UDPConn: ln,
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))
}
}
}, 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)
}

5
udp.go
View File

@ -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) {