github.com/metacubex/mihomo@v1.18.5/component/nat/table.go (about)

     1  package nat
     2  
     3  import (
     4  	"net"
     5  	"sync"
     6  
     7  	C "github.com/metacubex/mihomo/constant"
     8  
     9  	"github.com/puzpuzpuz/xsync/v3"
    10  )
    11  
    12  type Table struct {
    13  	mapping *xsync.MapOf[string, *Entry]
    14  	lockMap *xsync.MapOf[string, *sync.Cond]
    15  }
    16  
    17  type Entry struct {
    18  	PacketConn      C.PacketConn
    19  	WriteBackProxy  C.WriteBackProxy
    20  	LocalUDPConnMap *xsync.MapOf[string, *net.UDPConn]
    21  	LocalLockMap    *xsync.MapOf[string, *sync.Cond]
    22  }
    23  
    24  func (t *Table) Set(key string, e C.PacketConn, w C.WriteBackProxy) {
    25  	t.mapping.Store(key, &Entry{
    26  		PacketConn:      e,
    27  		WriteBackProxy:  w,
    28  		LocalUDPConnMap: xsync.NewMapOf[string, *net.UDPConn](),
    29  		LocalLockMap:    xsync.NewMapOf[string, *sync.Cond](),
    30  	})
    31  }
    32  
    33  func (t *Table) Get(key string) (C.PacketConn, C.WriteBackProxy) {
    34  	entry, exist := t.getEntry(key)
    35  	if !exist {
    36  		return nil, nil
    37  	}
    38  	return entry.PacketConn, entry.WriteBackProxy
    39  }
    40  
    41  func (t *Table) GetOrCreateLock(key string) (*sync.Cond, bool) {
    42  	item, loaded := t.lockMap.LoadOrCompute(key, makeLock)
    43  	return item, loaded
    44  }
    45  
    46  func (t *Table) Delete(key string) {
    47  	t.mapping.Delete(key)
    48  }
    49  
    50  func (t *Table) DeleteLock(lockKey string) {
    51  	t.lockMap.Delete(lockKey)
    52  }
    53  
    54  func (t *Table) GetForLocalConn(lAddr, rAddr string) *net.UDPConn {
    55  	entry, exist := t.getEntry(lAddr)
    56  	if !exist {
    57  		return nil
    58  	}
    59  	item, exist := entry.LocalUDPConnMap.Load(rAddr)
    60  	if !exist {
    61  		return nil
    62  	}
    63  	return item
    64  }
    65  
    66  func (t *Table) AddForLocalConn(lAddr, rAddr string, conn *net.UDPConn) bool {
    67  	entry, exist := t.getEntry(lAddr)
    68  	if !exist {
    69  		return false
    70  	}
    71  	entry.LocalUDPConnMap.Store(rAddr, conn)
    72  	return true
    73  }
    74  
    75  func (t *Table) RangeForLocalConn(lAddr string, f func(key string, value *net.UDPConn) bool) {
    76  	entry, exist := t.getEntry(lAddr)
    77  	if !exist {
    78  		return
    79  	}
    80  	entry.LocalUDPConnMap.Range(f)
    81  }
    82  
    83  func (t *Table) GetOrCreateLockForLocalConn(lAddr, key string) (*sync.Cond, bool) {
    84  	entry, loaded := t.getEntry(lAddr)
    85  	if !loaded {
    86  		return nil, false
    87  	}
    88  	item, loaded := entry.LocalLockMap.LoadOrCompute(key, makeLock)
    89  	return item, loaded
    90  }
    91  
    92  func (t *Table) DeleteForLocalConn(lAddr, key string) {
    93  	entry, loaded := t.getEntry(lAddr)
    94  	if !loaded {
    95  		return
    96  	}
    97  	entry.LocalUDPConnMap.Delete(key)
    98  }
    99  
   100  func (t *Table) DeleteLockForLocalConn(lAddr, key string) {
   101  	entry, loaded := t.getEntry(lAddr)
   102  	if !loaded {
   103  		return
   104  	}
   105  	entry.LocalLockMap.Delete(key)
   106  }
   107  
   108  func (t *Table) getEntry(key string) (*Entry, bool) {
   109  	return t.mapping.Load(key)
   110  }
   111  
   112  func makeLock() *sync.Cond {
   113  	return sync.NewCond(&sync.Mutex{})
   114  }
   115  
   116  // New return *Cache
   117  func New() *Table {
   118  	return &Table{
   119  		mapping: xsync.NewMapOf[string, *Entry](),
   120  		lockMap: xsync.NewMapOf[string, *sync.Cond](),
   121  	}
   122  }