memory.go 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. package fakeip
  2. import (
  3. "net/netip"
  4. "github.com/metacubex/mihomo/common/lru"
  5. )
  6. type memoryStore struct {
  7. cacheIP *lru.LruCache[string, netip.Addr]
  8. cacheHost *lru.LruCache[netip.Addr, string]
  9. }
  10. // GetByHost implements store.GetByHost
  11. func (m *memoryStore) GetByHost(host string) (netip.Addr, bool) {
  12. if ip, exist := m.cacheIP.Get(host); exist {
  13. // ensure ip --> host on head of linked list
  14. m.cacheHost.Get(ip)
  15. return ip, true
  16. }
  17. return netip.Addr{}, false
  18. }
  19. // PutByHost implements store.PutByHost
  20. func (m *memoryStore) PutByHost(host string, ip netip.Addr) {
  21. m.cacheIP.Set(host, ip)
  22. }
  23. // GetByIP implements store.GetByIP
  24. func (m *memoryStore) GetByIP(ip netip.Addr) (string, bool) {
  25. if host, exist := m.cacheHost.Get(ip); exist {
  26. // ensure host --> ip on head of linked list
  27. m.cacheIP.Get(host)
  28. return host, true
  29. }
  30. return "", false
  31. }
  32. // PutByIP implements store.PutByIP
  33. func (m *memoryStore) PutByIP(ip netip.Addr, host string) {
  34. m.cacheHost.Set(ip, host)
  35. }
  36. // DelByIP implements store.DelByIP
  37. func (m *memoryStore) DelByIP(ip netip.Addr) {
  38. if host, exist := m.cacheHost.Get(ip); exist {
  39. m.cacheIP.Delete(host)
  40. }
  41. m.cacheHost.Delete(ip)
  42. }
  43. // Exist implements store.Exist
  44. func (m *memoryStore) Exist(ip netip.Addr) bool {
  45. return m.cacheHost.Exist(ip)
  46. }
  47. // CloneTo implements store.CloneTo
  48. // only for memoryStore to memoryStore
  49. func (m *memoryStore) CloneTo(store store) {
  50. if ms, ok := store.(*memoryStore); ok {
  51. m.cacheIP.CloneTo(ms.cacheIP)
  52. m.cacheHost.CloneTo(ms.cacheHost)
  53. }
  54. }
  55. // FlushFakeIP implements store.FlushFakeIP
  56. func (m *memoryStore) FlushFakeIP() error {
  57. _ = m.cacheIP.Clear()
  58. return m.cacheHost.Clear()
  59. }
  60. func newMemoryStore(size int) *memoryStore {
  61. return &memoryStore{
  62. cacheIP: lru.New[string, netip.Addr](lru.WithSize[string, netip.Addr](size)),
  63. cacheHost: lru.New[netip.Addr, string](lru.WithSize[netip.Addr, string](size)),
  64. }
  65. }