permission.go 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. package client
  2. import (
  3. "net"
  4. "sync"
  5. "sync/atomic"
  6. )
  7. type permState int32
  8. const (
  9. permStateIdle permState = iota
  10. permStatePermitted
  11. )
  12. type permission struct {
  13. st permState // thread-safe (atomic op)
  14. mutex sync.RWMutex // thread-safe
  15. }
  16. func (p *permission) setState(state permState) {
  17. atomic.StoreInt32((*int32)(&p.st), int32(state))
  18. }
  19. func (p *permission) state() permState {
  20. return permState(atomic.LoadInt32((*int32)(&p.st)))
  21. }
  22. // Thread-safe permission map
  23. type permissionMap struct {
  24. permMap map[string]*permission
  25. mutex sync.RWMutex
  26. }
  27. func (m *permissionMap) insert(addr net.Addr, p *permission) bool {
  28. m.mutex.Lock()
  29. defer m.mutex.Unlock()
  30. udpAddr, ok := addr.(*net.UDPAddr)
  31. if !ok {
  32. return false
  33. }
  34. m.permMap[udpAddr.IP.String()] = p
  35. return true
  36. }
  37. func (m *permissionMap) find(addr net.Addr) (*permission, bool) {
  38. m.mutex.RLock()
  39. defer m.mutex.RUnlock()
  40. udpAddr, ok := addr.(*net.UDPAddr)
  41. if !ok {
  42. return nil, false
  43. }
  44. p, ok := m.permMap[udpAddr.IP.String()]
  45. return p, ok
  46. }
  47. func (m *permissionMap) delete(addr net.Addr) {
  48. m.mutex.Lock()
  49. defer m.mutex.Unlock()
  50. udpAddr, ok := addr.(*net.UDPAddr)
  51. if !ok {
  52. return
  53. }
  54. delete(m.permMap, udpAddr.IP.String())
  55. }
  56. func (m *permissionMap) addrs() []net.Addr {
  57. m.mutex.RLock()
  58. defer m.mutex.RUnlock()
  59. addrs := []net.Addr{}
  60. for k := range m.permMap {
  61. addrs = append(addrs, &net.UDPAddr{
  62. IP: net.ParseIP(k),
  63. })
  64. }
  65. return addrs
  66. }
  67. func newPermissionMap() *permissionMap {
  68. return &permissionMap{
  69. permMap: map[string]*permission{},
  70. }
  71. }