table.go 2.6 KB

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