zeromesh/internal/vl1/peer.go

package vl1

import (
	"log/slog"
	"net"
	"sync"
	"time"

	"zeromesh/internal/identity"
)

type Peer struct {
	Address    identity.Address
	PublicKey  [32]byte
	Endpoint   *net.UDPAddr
	Cipher     *NoiseCipher
	LastSeen   time.Time
	LastSend   time.Time
	Connected  bool
	mu         sync.RWMutex
	log        *slog.Logger
}

func (p *Peer) Touch() {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.LastSeen = time.Now()
}

func (p *Peer) SetCipher(c *NoiseCipher) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.Cipher = c
	p.Connected = true
}

func (p *Peer) IsConnected() bool {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.Connected
}

func (p *Peer) IsAlive() bool {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return time.Since(p.LastSeen) < 90*time.Second
}

func (p *Peer) Encrypt(plaintext []byte) ([]byte, error) {
	p.mu.RLock()
	defer p.mu.RUnlock()
	if p.Cipher == nil {
		return nil, ErrNoCipher
	}
	return p.Cipher.Encrypt(plaintext)
}

func (p *Peer) Decrypt(ciphertext []byte) ([]byte, error) {
	p.mu.RLock()
	defer p.mu.RUnlock()
	if p.Cipher == nil {
		return nil, ErrNoCipher
	}
	return p.Cipher.Decrypt(ciphertext)
}

func (p *Peer) EncryptTo(buf []byte, plaintext []byte) (int, error) {
	p.mu.RLock()
	defer p.mu.RUnlock()
	if p.Cipher == nil {
		return 0, ErrNoCipher
	}
	return p.Cipher.EncryptTo(buf, plaintext)
}

func (p *Peer) DecryptTo(buf []byte, ciphertext []byte) ([]byte, error) {
	p.mu.RLock()
	defer p.mu.RUnlock()
	if p.Cipher == nil {
		return nil, ErrNoCipher
	}
	return p.Cipher.DecryptTo(buf, ciphertext)
}

func (p *Peer) NeedsKeepalive() bool {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.Connected && time.Since(p.LastSend) > 25*time.Second
}

var ErrNoCipher = errNoCipher()

func errNoCipher() error {
	return &noCipherError{}
}

type noCipherError struct{}

func (e *noCipherError) Error() string {
	return "no cipher established"
}

type PeerManager struct {
	peers map[identity.Address]*Peer
	mu    sync.RWMutex
	log   *slog.Logger
}

func NewPeerManager(log *slog.Logger) *PeerManager {
	return &PeerManager{
		peers: make(map[identity.Address]*Peer),
		log:   log,
	}
}

func (pm *PeerManager) AddPeer(addr identity.Address, pubKey [32]byte, endpoint *net.UDPAddr) *Peer {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	peer := &Peer{
		Address:   addr,
		PublicKey: pubKey,
		Endpoint:  endpoint,
		LastSeen:  time.Now(),
		LastSend:  time.Now(),
		log:       pm.log.With("peer", addr.String()),
	}
	pm.peers[addr] = peer
	return peer
}

func (pm *PeerManager) GetPeer(addr identity.Address) *Peer {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	return pm.peers[addr]
}

func (pm *PeerManager) GetPeerByEndpoint(endpoint *net.UDPAddr) *Peer {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	for _, p := range pm.peers {
		if p.Endpoint != nil && p.Endpoint.IP.Equal(endpoint.IP) && p.Endpoint.Port == endpoint.Port {
			return p
		}
	}
	return nil
}

func (pm *PeerManager) RemovePeer(addr identity.Address) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	delete(pm.peers, addr)
}

func (pm *PeerManager) UpdatePeerEndpoint(addr identity.Address, endpoint *net.UDPAddr) {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	if p, ok := pm.peers[addr]; ok {
		p.Endpoint = endpoint
	}
}

func (pm *PeerManager) ConnectedPeers() []*Peer {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	var result []*Peer
	for _, p := range pm.peers {
		if p.IsConnected() {
			result = append(result, p)
		}
	}
	return result
}

func (pm *PeerManager) AllPeers() []*Peer {
	pm.mu.RLock()
	defer pm.mu.RUnlock()
	result := make([]*Peer, 0, len(pm.peers))
	for _, p := range pm.peers {
		result = append(result, p)
	}
	return result
}

func (pm *PeerManager) CleanDead() {
	pm.mu.Lock()
	defer pm.mu.Unlock()
	for addr, p := range pm.peers {
		if p.Connected && !p.IsAlive() {
			p.log.Warn("peer timed out, removing")
			delete(pm.peers, addr)
		}
	}
}