Fix: some UDP issues (#265)

component/nat-table/nat.go

@@ -1,98 +0,0 @@
-package nat
-
-import (
-	"net"
-	"runtime"
-	"sync"
-	"time"
-)
-
-type Table struct {
-	*table
-}
-
-type table struct {
-	mapping sync.Map
-	janitor *janitor
-	timeout time.Duration
-}
-
-type element struct {
-	Expired    time.Time
-	RemoteAddr net.Addr
-	RemoteConn net.PacketConn
-}
-
-func (t *table) Set(key net.Addr, rConn net.PacketConn, rAddr net.Addr) {
-	// set conn read timeout
-	rConn.SetReadDeadline(time.Now().Add(t.timeout))
-	t.mapping.Store(key, &element{
-		RemoteAddr: rAddr,
-		RemoteConn: rConn,
-		Expired:    time.Now().Add(t.timeout),
-	})
-}
-
-func (t *table) Get(key net.Addr) (rConn net.PacketConn, rAddr net.Addr) {
-	item, exist := t.mapping.Load(key)
-	if !exist {
-		return
-	}
-	elm := item.(*element)
-	// expired
-	if time.Since(elm.Expired) > 0 {
-		t.mapping.Delete(key)
-		elm.RemoteConn.Close()
-		return
-	}
-	// reset expired time
-	elm.Expired = time.Now().Add(t.timeout)
-	return elm.RemoteConn, elm.RemoteAddr
-}
-
-func (t *table) cleanup() {
-	t.mapping.Range(func(k, v interface{}) bool {
-		key := k.(net.Addr)
-		elm := v.(*element)
-		if time.Since(elm.Expired) > 0 {
-			t.mapping.Delete(key)
-			elm.RemoteConn.Close()
-		}
-		return true
-	})
-}
-
-type janitor struct {
-	interval time.Duration
-	stop     chan struct{}
-}
-
-func (j *janitor) process(t *table) {
-	ticker := time.NewTicker(j.interval)
-	for {
-		select {
-		case <-ticker.C:
-			t.cleanup()
-		case <-j.stop:
-			ticker.Stop()
-			return
-		}
-	}
-}
-
-func stopJanitor(t *Table) {
-	t.janitor.stop <- struct{}{}
-}
-
-// New return *Cache
-func New(interval time.Duration) *Table {
-	j := &janitor{
-		interval: interval,
-		stop:     make(chan struct{}),
-	}
-	t := &table{janitor: j, timeout: interval}
-	go j.process(t)
-	T := &Table{t}
-	runtime.SetFinalizer(T, stopJanitor)
-	return T
-}

component/nat/table.go

@@ -0,0 +1,46 @@
+package nat
+
+import (
+	"net"
+	"sync"
+)
+
+type Table struct {
+	mapping sync.Map
+}
+
+type element struct {
+	RemoteAddr net.Addr
+	RemoteConn net.PacketConn
+}
+
+func (t *Table) Set(key string, pc net.PacketConn, addr net.Addr) {
+	// set conn read timeout
+	t.mapping.Store(key, &element{
+		RemoteConn: pc,
+		RemoteAddr: addr,
+	})
+}
+
+func (t *Table) Get(key string) (net.PacketConn, net.Addr) {
+	item, exist := t.mapping.Load(key)
+	if !exist {
+		return nil, nil
+	}
+	elm := item.(*element)
+	return elm.RemoteConn, elm.RemoteAddr
+}
+
+func (t *Table) GetOrCreateLock(key string) (*sync.WaitGroup, bool) {
+	item, loaded := t.mapping.LoadOrStore(key, &sync.WaitGroup{})
+	return item.(*sync.WaitGroup), loaded
+}
+
+func (t *Table) Delete(key string) {
+	t.mapping.Delete(key)
+}
+
+// New return *Cache
+func New() *Table {
+	return &Table{}
+}

component/socks5/socks5.go

@@ -338,6 +338,7 @@ func ParseAddr(s string) Addr {
 	return addr
 }
 
+// DecodeUDPPacket split `packet` to addr payload, and this function is mutable with `packet`
 func DecodeUDPPacket(packet []byte) (addr Addr, payload []byte, err error) {
 	if len(packet) < 5 {
 		err = errors.New("insufficient length of packet")
@@ -360,16 +361,15 @@ func DecodeUDPPacket(packet []byte) (addr Addr, payload []byte, err error) {
 		err = errors.New("failed to read UDP header")
 	}
 
-	payload = bytes.Join([][]byte{packet[3+len(addr):]}, []byte{})
+	payload = packet[3+len(addr):]
 	return
 }
 
-func EncodeUDPPacket(addr string, payload []byte) (packet []byte, err error) {
-	rAddr := ParseAddr(addr)
-	if rAddr == nil {
-		err = errors.New("cannot parse addr")
+func EncodeUDPPacket(addr Addr, payload []byte) (packet []byte, err error) {
+	if addr == nil {
+		err = errors.New("address is invalid")
 		return
 	}
-	packet = bytes.Join([][]byte{{0, 0, 0}, rAddr, payload}, []byte{})
+	packet = bytes.Join([][]byte{{0, 0, 0}, addr, payload}, []byte{})
 	return
 }