pg/disco/types.go

package disco

import (
	"errors"
	"log/slog"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/rkonfj/peerguard/peer"
	"github.com/rkonfj/peerguard/secure"
)

var (
	ErrUseOfClosedConnection error = errors.New("use of closed network connection")
)

const (
	OP_PEER_DISCO       = 1
	OP_PEER_CONFIRM     = 2
	OP_PEER_DELETE      = 3
	OP_PEER_HEALTHCHECK = 10
)

type PeerStore interface {
	FindPeer(peer.PeerID) (*PeerContext, bool)
}

type PeerContext struct {
	PeerID     peer.PeerID
	States     map[string]*PeerState
	CreateTime time.Time

	exitSig           chan struct{}
	ping              func(addr *net.UDPAddr)
	keepaliveInterval time.Duration

	statesMutex sync.RWMutex
}

func (peer *PeerContext) AddAddr(addr *net.UDPAddr) {
	peer.statesMutex.Lock()
	defer peer.statesMutex.Unlock()
	if _, ok := peer.States[addr.String()]; ok {
		return
	}
	peer.States[addr.String()] = &PeerState{Addr: addr}
}

func (peer *PeerContext) Heartbeat(addr *net.UDPAddr) {
	peer.statesMutex.Lock()
	defer peer.statesMutex.Unlock()
	for _, state := range peer.States {
		if state.Addr.IP.Equal(addr.IP) && state.Addr.Port == addr.Port {
			updated := time.Since(state.LastActiveTime) > 2*peer.keepaliveInterval+2*time.Second
			if updated {
				slog.Info("[UDP] AddPeer", "peer", peer.PeerID, "addr", addr)
			}
			state.LastActiveTime = time.Now()
			return
		}
	}
	slog.Info("[UDP][0RTT] AddPeer", "peer", peer.PeerID, "addr", addr)
	peer.States[addr.String()] = &PeerState{Addr: addr, LastActiveTime: time.Now()}
}

func (peer *PeerContext) Healthcheck() {
	if time.Since(peer.CreateTime) > 3*peer.keepaliveInterval {
		for addr, state := range peer.States {
			if time.Since(state.LastActiveTime) > 2*peer.keepaliveInterval {
				if state.LastActiveTime.IsZero() {
					slog.Debug("[UDP] RemovePeer", "peer", peer.PeerID, "addr", state.Addr)
				} else {
					slog.Info("[UDP] RemovePeer", "peer", peer.PeerID, "addr", state.Addr)
				}
				peer.statesMutex.Lock()
				delete(peer.States, addr)
				peer.statesMutex.Unlock()
			}
		}
	}
}

func (peer *PeerContext) IPv4Ready() bool {
	peer.statesMutex.RLock()
	defer peer.statesMutex.RUnlock()
	for _, state := range peer.States {
		if state.Addr.IP.To4() != nil && time.Since(state.LastActiveTime) <= peer.keepaliveInterval+2*time.Second {
			return true
		}
	}
	return false
}

func (peer *PeerContext) Ready() bool {
	peer.statesMutex.RLock()
	defer peer.statesMutex.RUnlock()
	for _, state := range peer.States {
		if time.Since(state.LastActiveTime) <= peer.keepaliveInterval+2*time.Second {
			return true
		}
	}
	return false
}

func (peer *PeerContext) Select() *net.UDPAddr {
	peer.statesMutex.RLock()
	defer peer.statesMutex.RUnlock()
	addrs := make([]*net.UDPAddr, 0, len(peer.States))
	for _, state := range peer.States {
		if time.Since(state.LastActiveTime) <= peer.keepaliveInterval+2*time.Second {
			addrs = append(addrs, state.Addr)
		}
	}
	return addrs[rand.Intn(len(addrs))]
}

func (peer *PeerContext) Keepalive() {
	ticker := time.NewTicker(peer.keepaliveInterval)
	ping := func() {
		addrs := make([]*net.UDPAddr, 0, len(peer.States))
		peer.statesMutex.RLock()
		for _, v := range peer.States {
			if time.Since(v.LastActiveTime) > 2*peer.keepaliveInterval+2*time.Second {
				continue
			}
			addrs = append(addrs, v.Addr)
		}
		peer.statesMutex.RUnlock()
		for _, addr := range addrs {
			peer.ping(addr)
		}
	}
	ping()
	for {
		select {
		case <-peer.exitSig:
			ticker.Stop()
			slog.Debug("[UDP] Ping exit", "peer", peer.PeerID)
			return
		case <-ticker.C:
			ping()
		}
	}
}

func (peer *PeerContext) Close() error {
	close(peer.exitSig)
	return nil
}

type PeerState struct {
	Addr           *net.UDPAddr
	LastActiveTime time.Time
}

type PeerOP struct {
	Op     int
	Addr   *net.UDPAddr
	PeerID peer.PeerID
}

type STUNSession struct {
	PeerID peer.PeerID
	CTime  time.Time
}

type Datagram struct {
	PeerID peer.PeerID
	Data   []byte
}

func (d *Datagram) TryDecrypt(symmAlgo secure.SymmAlgo) []byte {
	b, err := symmAlgo.Decrypt(d.Data, d.PeerID.String())
	if err != nil {
		slog.Debug("Datagram decrypt error", "err", err)
		return d.Data
	}
	return b
}

func (d *Datagram) TryEncrypt(symmAlgo secure.SymmAlgo) []byte {
	b, err := symmAlgo.Encrypt(d.Data, d.PeerID.String())
	if err != nil {
		slog.Debug("Datagram encrypt error", "err", err)
		return d.Data
	}
	return b
}

type PeerFindEvent struct {
	PeerID   peer.PeerID
	Metadata peer.Metadata
}

type PeerUDPAddrEvent struct {
	PeerID peer.PeerID
	Addr   *net.UDPAddr
}