github.com/noisysockets/noisysockets@v0.21.2-0.20240515114641-7f467e651c90/peer.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  	"fmt"
    14  	"net/netip"
    15  	"slices"
    16  	"sync"
    17  
    18  	"github.com/noisysockets/noisysockets/internal/conn"
    19  	"github.com/noisysockets/noisysockets/internal/transport"
    20  	"github.com/noisysockets/noisysockets/types"
    21  )
    22  
    23  var (
    24  	ErrNoEndpoint = fmt.Errorf("no known endpoint for peer")
    25  )
    26  
    27  // Peer represents a wireguard peer in the network.
    28  type Peer struct {
    29  	*transport.Peer
    30  	mu                     sync.RWMutex
    31  	name                   string
    32  	publicKey              types.NoisePublicKey
    33  	addrs                  []netip.Addr
    34  	destinationForPrefixes []netip.Prefix
    35  }
    36  
    37  func newPeer(transportPeer *transport.Peer, name string, publicKey types.NoisePublicKey) *Peer {
    38  	return &Peer{
    39  		Peer:      transportPeer,
    40  		name:      name,
    41  		publicKey: publicKey,
    42  	}
    43  }
    44  
    45  // Name returns the human friendly name of the peer.
    46  func (p *Peer) Name() string {
    47  	return p.name
    48  }
    49  
    50  // PublicKey returns the public key of the peer.
    51  func (p *Peer) PublicKey() types.NoisePublicKey {
    52  	return p.publicKey
    53  }
    54  
    55  // Addresses returns the list of addresses of the peer.
    56  func (p *Peer) Addresses() []netip.Addr {
    57  	p.mu.RLock()
    58  	defer p.mu.RUnlock()
    59  
    60  	return p.addrs
    61  }
    62  
    63  // AddAddress adds one or more addresses to the peer.
    64  func (p *Peer) AddAddresses(addrs ...netip.Addr) {
    65  	p.mu.Lock()
    66  	defer p.mu.Unlock()
    67  
    68  	p.addrs = slices.CompactFunc(append(p.addrs, addrs...), func(a, b netip.Addr) bool { return a == b })
    69  }
    70  
    71  // RemoveAddress removes one or more addresses from the peer.
    72  func (p *Peer) RemoveAddresses(addrs ...netip.Addr) {
    73  	p.mu.Lock()
    74  	defer p.mu.Unlock()
    75  
    76  	toRemove := make(map[netip.Addr]bool, len(addrs))
    77  	for _, a := range addrs {
    78  		toRemove[a] = true
    79  	}
    80  
    81  	for i, a := range p.addrs {
    82  		if toRemove[a] {
    83  			p.addrs = append(p.addrs[:i], p.addrs[i+1:]...)
    84  		}
    85  	}
    86  }
    87  
    88  // DestinationForPrefixes returns the list of prefixes the peer is the destination for.
    89  func (p *Peer) DestinationForPrefixes() []netip.Prefix {
    90  	p.mu.RLock()
    91  	defer p.mu.RUnlock()
    92  
    93  	return p.destinationForPrefixes
    94  }
    95  
    96  // AddDestinationForPrefix adds one or more prefixes the peer is the destination for.
    97  func (p *Peer) AddDestinationPrefixes(prefixes ...netip.Prefix) {
    98  	p.mu.Lock()
    99  	defer p.mu.Unlock()
   100  
   101  	p.destinationForPrefixes = slices.CompactFunc(append(p.destinationForPrefixes, prefixes...),
   102  		func(a, b netip.Prefix) bool { return a == b })
   103  }
   104  
   105  // RemoveDestinationForPrefix removes one or more prefixes the peer is the destination for.
   106  func (p *Peer) RemoveDestinationPrefixes(prefixes ...netip.Prefix) {
   107  	p.mu.Lock()
   108  	defer p.mu.Unlock()
   109  
   110  	toRemove := make(map[netip.Prefix]bool, len(prefixes))
   111  	for _, prefix := range prefixes {
   112  		toRemove[prefix] = true
   113  	}
   114  
   115  	for i, prefix := range p.destinationForPrefixes {
   116  		if toRemove[prefix] {
   117  			p.destinationForPrefixes = append(p.destinationForPrefixes[:i], p.destinationForPrefixes[i+1:]...)
   118  		}
   119  	}
   120  }
   121  
   122  // GetEndpoint returns the endpoint (public address) of the peer.
   123  func (p *Peer) GetEndpoint() (netip.AddrPort, error) {
   124  	p.mu.RLock()
   125  	defer p.mu.RUnlock()
   126  
   127  	endpoint := p.Peer.GetEndpoint()
   128  	if endpoint == nil {
   129  		return netip.AddrPort{}, ErrNoEndpoint
   130  	}
   131  
   132  	return netip.ParseAddrPort(endpoint.DstToString())
   133  }
   134  
   135  // SetEndpoint sets the endpoint (public address) of the peer.
   136  func (p *Peer) SetEndpoint(endpoint netip.AddrPort) {
   137  	p.mu.Lock()
   138  	defer p.mu.Unlock()
   139  
   140  	p.Peer.SetEndpoint(&conn.StdNetEndpoint{AddrPort: endpoint})
   141  }