Limit the maximum number of udp connections

Author: mohanson

daze.go

@@ -22,7 +22,6 @@ import (
 	"net/url"
 	"os"
 	"path/filepath"
-	"slices"
 	"strconv"
 	"strings"
 	"sync"
@@ -50,12 +49,14 @@ import (
 var Conf = struct {
 	DialerTimeout time.Duration
 	RouterLruSize int
+	Socks5LruSize int
 }{
 	DialerTimeout: time.Second * 8,
 	// A single cache entry represents a single host or DNS name lookup. Make the cache as large as the maximum number
 	// of clients that access your web site concurrently. Note that setting the cache size too high is a waste of
 	// memory and degrades performance.
 	RouterLruSize: 64,
+	Socks5LruSize: 8,
 }
 
 // ResolverDns returns a DNS resolver.
@@ -409,6 +410,28 @@ func (l *Locale) ServeSocks5TCP(ctx *Context, cli io.ReadWriteCloser, dst string
 	return err
 }
 
+// ServeSocks5UDPRead handles the reading and forwarding of udp data.
+func (l *Locale) ServeSocks5UDPRead(srv io.Reader, bnd *net.UDPConn, app *net.UDPAddr, pre []byte) error {
+	var (
+		buf = make([]byte, 2048)
+		err error
+		m   = len(pre)
+		n   int
+	)
+	copy(buf[:m], pre)
+	for {
+		n, err = srv.Read(buf[m:])
+		if err != nil {
+			break
+		}
+		_, err = bnd.WriteToUDP(buf[:m+n], app)
+		if err != nil {
+			break
+		}
+	}
+	return err
+}
+
 // ServeSocks5UDP serves socks5 UDP protocol.
 func (l *Locale) ServeSocks5UDP(ctx *Context, cli io.ReadWriteCloser) error {
 	var (
@@ -420,7 +443,7 @@ func (l *Locale) ServeSocks5UDP(ctx *Context, cli io.ReadWriteCloser) error {
 		bndPort     uint16
 		bnd         *net.UDPConn
 		buf         = make([]byte, 2048)
-		cpl         = map[string]io.ReadWriteCloser{}
+		cpl         = lru.New[string, io.ReadWriteCloser](Conf.Socks5LruSize)
 		dstHost     string
 		dstPort     uint16
 		dst         string
@@ -437,6 +460,9 @@ func (l *Locale) ServeSocks5UDP(ctx *Context, cli io.ReadWriteCloser) error {
 	if err != nil {
 		return err
 	}
+	cpl.Drop = func(k string, v io.ReadWriteCloser) {
+		v.Close()
+	}
 
 	// https://datatracker.ietf.org/doc/html/rfc1928, Page 7, UDP ASSOCIATE:
 	// A UDP association terminates when the TCP connection that the UDP ASSOCIATE request arrived on terminates.
@@ -492,30 +518,26 @@ func (l *Locale) ServeSocks5UDP(ctx *Context, cli io.ReadWriteCloser) error {
 			dstPort = binary.BigEndian.Uint16(appHead[20:22])
 		}
 		dst = dstHost + ":" + strconv.Itoa(int(dstPort))
-		srv = cpl[dst]
-		if srv == nil {
+		if !cpl.Has(dst) {
 			log.Printf("conn: %08x  proto format=socks5", ctx.Cid)
 			srv, err = l.Dialer.Dial(ctx, "udp", dst)
 			if err != nil {
 				log.Printf("conn: %08x  error %s", ctx.Cid, err)
 				continue
 			}
-			cpl[dst] = srv
-			retHead := slices.Clone(appHead)
-			go ReadCall(srv, func(data []byte) error {
-				return doa.Err(bnd.WriteToUDP(append(retHead, data...), appAddr))
-			})
+			cpl.Set(dst, srv)
+			go l.ServeSocks5UDPRead(srv, bnd, appAddr, appHead)
 		}
+		srv = cpl.Get(dst)
 		_, err = srv.Write(buf[appHeadSize:appSize])
 		if err != nil {
 			log.Printf("conn: %08x  error %s", ctx.Cid, err)
-			srv.Close()
-			delete(cpl, dst)
+			cpl.Del(dst)
 			continue
 		}
 	}
-	for _, e := range cpl {
-		e.Close()
+	for k := range cpl.C {
+		cpl.Del(k)
 	}
 	return nil
 }
@@ -1030,25 +1052,6 @@ func (r *RandomReader) Read(p []byte) (int, error) {
 	return len(p), nil
 }
 
-// ReadCall reads data from the given io.Reader and passes it to the provided call function.
-func ReadCall(conn io.Reader, call func([]byte) error) error {
-	var (
-		buf = make([]byte, 2048)
-		err error
-		n   int
-	)
-	for {
-		n, err = conn.Read(buf)
-		if err != nil {
-			break
-		}
-		if err = call(buf[:n]); err != nil {
-			break
-		}
-	}
-	return err
-}
-
 // Salt converts the stupid password passed in by the user to 32-sized byte array.
 func Salt(s string) []byte {
 	h := sha256.Sum256([]byte(s))

lib/lru/lru.go

@@ -38,6 +38,15 @@ func (l *List[K, V]) Insert(e *Elem[K, V]) *Elem[K, V] {
 	return e
 }
 
+// Remove removes e from its list, decrements l's size.
+func (l *List[K, V]) Remove(e *Elem[K, V]) {
+	e.Prev.Next = e.Next
+	e.Next.Prev = e.Prev
+	e.Prev = nil // Avoid memory leaks
+	e.Next = nil // Avoid memory leaks
+	l.Size--
+}
+
 // Update e to next to root.
 func (l *List[K, V]) Update(e *Elem[K, V]) {
 	if l.Root.Next == e {
@@ -51,42 +60,17 @@ func (l *List[K, V]) Update(e *Elem[K, V]) {
 	e.Next.Prev = e
 }
 
-// Remove removes e from its list, decrements l's size.
-func (l *List[K, V]) Remove(e *Elem[K, V]) {
-	e.Prev.Next = e.Next
-	e.Next.Prev = e.Prev
-	e.Prev = nil // Avoid memory leaks
-	e.Next = nil // Avoid memory leaks
-	l.Size--
-}
-
 // Lru cache. It is safe for concurrent access.
 type Lru[K comparable, V any] struct {
-	// Size is the maximum number of cache entries before
-	// an item is evicted. Zero means no limit.
+	// Drop is called automatically when an elem is deleted.
+	Drop func(k K, v V)
+	// Size is the maximum number of cache entries before an item is evicted. Zero means no limit.
 	Size int
 	List *List[K, V]
 	C    map[K]*Elem[K, V]
 	M    *sync.Mutex
 }
 
-// Set adds a value to the cache.
-func (l *Lru[K, V]) Set(k K, v V) {
-	l.M.Lock()
-	defer l.M.Unlock()
-	if e, ok := l.C[k]; ok {
-		l.List.Update(e)
-		e.K = k
-		e.V = v
-		return
-	}
-	if l.List.Size == l.Size {
-		delete(l.C, l.List.Root.Prev.K)
-		l.List.Remove(l.List.Root.Prev)
-	}
-	l.C[k] = l.List.Insert(&Elem[K, V]{K: k, V: v})
-}
-
 // Get looks up a key's value from the cache.
 func (l *Lru[K, V]) GetExists(k K) (v V, ok bool) {
 	l.M.Lock()
@@ -111,21 +95,47 @@ func (l *Lru[K, V]) Del(k K) {
 	l.M.Lock()
 	defer l.M.Unlock()
 	if e, ok := l.C[k]; ok {
+		l.Drop(k, e.V)
 		delete(l.C, k)
 		l.List.Remove(e)
 	}
 }
 
+// Has returns true if a key exists.
+func (l *Lru[K, V]) Has(k K) bool {
+	_, b := l.C[k]
+	return b
+}
+
 // Len returns the number of items in the cache.
 func (l *Lru[K, V]) Len() int {
 	l.M.Lock()
 	defer l.M.Unlock()
 	return l.List.Size
 }
 
+// Set adds a value to the cache.
+func (l *Lru[K, V]) Set(k K, v V) {
+	l.M.Lock()
+	defer l.M.Unlock()
+	if e, ok := l.C[k]; ok {
+		l.List.Update(e)
+		e.K = k
+		e.V = v
+		return
+	}
+	if l.List.Size == l.Size {
+		l.Drop(l.List.Root.Prev.K, l.List.Root.Prev.V)
+		delete(l.C, l.List.Root.Prev.K)
+		l.List.Remove(l.List.Root.Prev)
+	}
+	l.C[k] = l.List.Insert(&Elem[K, V]{K: k, V: v})
+}
+
 // New returns a new LRU cache. If size is zero, the cache has no limit.
 func New[K comparable, V any](size int) *Lru[K, V] {
 	return &Lru[K, V]{
+		Drop: func(k K, v V) {},
 		Size: size,
 		List: new(List[K, V]).Init(),
 		C:    map[K]*Elem[K, V]{},