2025-03-03 12:12:30

Author: mohanson

daze.go

@@ -49,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.
@@ -441,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
@@ -458,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.
@@ -513,27 +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
+			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
 }

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]{},