github.com/okex/exchain@v1.8.0/libs/tendermint/p2p/conn_set.go (about)

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