github.com/noisysockets/noisysockets@v0.21.2-0.20240515114641-7f467e651c90/peer_list.go (about)

     1  // SPDX-License-Identifier: MPL-2.0
     2  /*
     3   * Copyright (C) 2024 The Noisy Sockets Authors.
     4   *
     5   * This Source Code Form is subject to the terms of the Mozilla Public
     6   * License, v. 2.0. If a copy of the MPL was not distributed with this
     7   * file, You can obtain one at http://mozilla.org/MPL/2.0/.
     8   */
     9  
    10  package noisysockets
    11  
    12  import (
    13  	"net/netip"
    14  	"sync"
    15  
    16  	"github.com/noisysockets/noisysockets/types"
    17  )
    18  
    19  type peerList struct {
    20  	mu            sync.RWMutex
    21  	byPublicKey   map[types.NoisePublicKey]*Peer
    22  	byName        map[string]*Peer
    23  	byAddr        map[netip.Addr]*Peer
    24  	byDestination map[netip.Prefix]*Peer
    25  }
    26  
    27  func newPeerList() *peerList {
    28  	return &peerList{
    29  		byPublicKey:   make(map[types.NoisePublicKey]*Peer),
    30  		byName:        make(map[string]*Peer),
    31  		byAddr:        make(map[netip.Addr]*Peer),
    32  		byDestination: make(map[netip.Prefix]*Peer),
    33  	}
    34  }
    35  
    36  func (pl *peerList) add(p *Peer) {
    37  	pl.mu.Lock()
    38  	defer pl.mu.Unlock()
    39  
    40  	pl.byPublicKey[p.PublicKey()] = p
    41  
    42  	if p.Name() != "" {
    43  		pl.byName[p.Name()] = p
    44  	}
    45  
    46  	for _, addr := range p.Addresses() {
    47  		pl.byAddr[addr] = p
    48  	}
    49  
    50  	for _, prefix := range p.DestinationForPrefixes() {
    51  		pl.byDestination[prefix] = p
    52  	}
    53  }
    54  
    55  func (pl *peerList) remove(publicKey types.NoisePublicKey) (*Peer, bool) {
    56  	pl.mu.Lock()
    57  	defer pl.mu.Unlock()
    58  
    59  	p, ok := pl.byPublicKey[publicKey]
    60  	if !ok {
    61  		return nil, false
    62  	}
    63  
    64  	delete(pl.byPublicKey, publicKey)
    65  
    66  	if p.Name() != "" {
    67  		delete(pl.byName, p.Name())
    68  	}
    69  
    70  	for _, addr := range p.Addresses() {
    71  		delete(pl.byAddr, addr)
    72  	}
    73  
    74  	for _, prefix := range p.DestinationForPrefixes() {
    75  		delete(pl.byDestination, prefix)
    76  	}
    77  
    78  	return p, true
    79  }
    80  
    81  func (pl *peerList) get(publicKey types.NoisePublicKey) (*Peer, bool) {
    82  	pl.mu.RLock()
    83  	defer pl.mu.RUnlock()
    84  
    85  	p, ok := pl.byPublicKey[publicKey]
    86  	return p, ok
    87  }
    88  
    89  func (pl *peerList) getByName(name string) (*Peer, bool) {
    90  	pl.mu.RLock()
    91  	defer pl.mu.RUnlock()
    92  
    93  	p, ok := pl.byName[name]
    94  	return p, ok
    95  }
    96  
    97  func (pl *peerList) getByAddress(addr netip.Addr) (*Peer, bool) {
    98  	pl.mu.RLock()
    99  	defer pl.mu.RUnlock()
   100  
   101  	p, ok := pl.byAddr[addr]
   102  	return p, ok
   103  }
   104  
   105  func (pl *peerList) getByDestination(prefix netip.Prefix) (*Peer, bool) {
   106  	pl.mu.RLock()
   107  	defer pl.mu.RUnlock()
   108  
   109  	p, ok := pl.byDestination[prefix]
   110  	return p, ok
   111  }
   112  
   113  func (pl *peerList) forEach(fn func(*Peer) error) error {
   114  	pl.mu.RLock()
   115  	defer pl.mu.RUnlock()
   116  
   117  	for _, p := range pl.byPublicKey {
   118  		if err := fn(p); err != nil {
   119  			return err
   120  		}
   121  	}
   122  
   123  	return nil
   124  }