github.com/pure-x-eth/consensus_tm@v0.0.0-20230502163723-e3c2ff987250/p2p/conn_set.go (about)

     1  package p2p
     2  
     3  import (
     4  	"net"
     5  
     6  	tmsync "github.com/pure-x-eth/consensus_tm/libs/sync"
     7  )
     8  
     9  // ConnSet is a lookup table for connections and all their ips.
    10  type ConnSet interface {
    11  	Has(net.Conn) bool
    12  	HasIP(net.IP) bool
    13  	Set(net.Conn, []net.IP)
    14  	Remove(net.Conn)
    15  	RemoveAddr(net.Addr)
    16  }
    17  
    18  type connSetItem struct {
    19  	conn net.Conn
    20  	ips  []net.IP
    21  }
    22  
    23  type connSet struct {
    24  	tmsync.RWMutex
    25  
    26  	conns map[string]connSetItem
    27  }
    28  
    29  // NewConnSet returns a ConnSet implementation.
    30  func NewConnSet() ConnSet {
    31  	return &connSet{
    32  		conns: map[string]connSetItem{},
    33  	}
    34  }
    35  
    36  func (cs *connSet) Has(c net.Conn) bool {
    37  	cs.RLock()
    38  	defer cs.RUnlock()
    39  
    40  	_, ok := cs.conns[c.RemoteAddr().String()]
    41  
    42  	return ok
    43  }
    44  
    45  func (cs *connSet) HasIP(ip net.IP) bool {
    46  	cs.RLock()
    47  	defer cs.RUnlock()
    48  
    49  	for _, c := range cs.conns {
    50  		for _, known := range c.ips {
    51  			if known.Equal(ip) {
    52  				return true
    53  			}
    54  		}
    55  	}
    56  
    57  	return false
    58  }
    59  
    60  func (cs *connSet) Remove(c net.Conn) {
    61  	cs.Lock()
    62  	defer cs.Unlock()
    63  
    64  	delete(cs.conns, c.RemoteAddr().String())
    65  }
    66  
    67  func (cs *connSet) RemoveAddr(addr net.Addr) {
    68  	cs.Lock()
    69  	defer cs.Unlock()
    70  
    71  	delete(cs.conns, addr.String())
    72  }
    73  
    74  func (cs *connSet) Set(c net.Conn, ips []net.IP) {
    75  	cs.Lock()
    76  	defer cs.Unlock()
    77  
    78  	cs.conns[c.RemoteAddr().String()] = connSetItem{
    79  		conn: c,
    80  		ips:  ips,
    81  	}
    82  }