package node import ( "encoding/json" "fmt" "os" "path/filepath" "regexp" "sync" "time" "forge.lthn.ai/core/go-p2p/logging" "github.com/Snider/Poindexter" "github.com/adrg/xdg" ) // Peer represents a known remote node. type Peer struct { ID string `json:"id"` Name string `json:"name"` PublicKey string `json:"publicKey"` Address string `json:"address"` // host:port for WebSocket connection Role NodeRole `json:"role"` AddedAt time.Time `json:"addedAt"` LastSeen time.Time `json:"lastSeen"` // Poindexter metrics (updated dynamically) PingMS float64 `json:"pingMs"` // Latency in milliseconds Hops int `json:"hops"` // Network hop count GeoKM float64 `json:"geoKm"` // Geographic distance in kilometers Score float64 `json:"score"` // Reliability score 0-100 // Connection state (not persisted) Connected bool `json:"-"` } // saveDebounceInterval is the minimum time between disk writes. const saveDebounceInterval = 5 * time.Second // PeerAuthMode controls how unknown peers are handled type PeerAuthMode int const ( // PeerAuthOpen allows any peer to connect (original behavior) PeerAuthOpen PeerAuthMode = iota // PeerAuthAllowlist only allows pre-registered peers or those with allowed public keys PeerAuthAllowlist ) // Peer name validation constants const ( PeerNameMinLength = 1 PeerNameMaxLength = 64 ) // peerNameRegex validates peer names: alphanumeric, hyphens, underscores, and spaces var peerNameRegex = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9\-_ ]{0,62}[a-zA-Z0-9]$|^[a-zA-Z0-9]$`) // safeKeyPrefix returns a truncated key for logging, handling short keys safely func safeKeyPrefix(key string) string { if len(key) >= 16 { return key[:16] + "..." } if len(key) == 0 { return "(empty)" } return key } // validatePeerName checks if a peer name is valid. // Peer names must be 1-64 characters, start and end with alphanumeric, // and contain only alphanumeric, hyphens, underscores, and spaces. func validatePeerName(name string) error { if name == "" { return nil // Empty names are allowed (optional field) } if len(name) < PeerNameMinLength { return fmt.Errorf("peer name too short (min %d characters)", PeerNameMinLength) } if len(name) > PeerNameMaxLength { return fmt.Errorf("peer name too long (max %d characters)", PeerNameMaxLength) } if !peerNameRegex.MatchString(name) { return fmt.Errorf("peer name contains invalid characters (use alphanumeric, hyphens, underscores, spaces)") } return nil } // PeerRegistry manages known peers with KD-tree based selection. type PeerRegistry struct { peers map[string]*Peer kdTree *poindexter.KDTree[string] // KD-tree with peer ID as payload path string mu sync.RWMutex // Authentication settings authMode PeerAuthMode // How to handle unknown peers allowedPublicKeys map[string]bool // Allowlist of public keys (when authMode is Allowlist) allowedPublicKeyMu sync.RWMutex // Protects allowedPublicKeys // Debounce disk writes dirty bool // Whether there are unsaved changes saveTimer *time.Timer // Timer for debounced save saveMu sync.Mutex // Protects dirty and saveTimer stopChan chan struct{} // Signal to stop background save saveStopOnce sync.Once // Ensure stopChan is closed only once } // Dimension weights for peer selection // Lower ping, hops, geo are better; higher score is better var ( pingWeight = 1.0 hopsWeight = 0.7 geoWeight = 0.2 scoreWeight = 1.2 ) // NewPeerRegistry creates a new PeerRegistry, loading existing peers if available. func NewPeerRegistry() (*PeerRegistry, error) { peersPath, err := xdg.ConfigFile("lethean-desktop/peers.json") if err != nil { return nil, fmt.Errorf("failed to get peers path: %w", err) } return NewPeerRegistryWithPath(peersPath) } // NewPeerRegistryWithPath creates a new PeerRegistry with a custom path. // This is primarily useful for testing to avoid xdg path caching issues. func NewPeerRegistryWithPath(peersPath string) (*PeerRegistry, error) { pr := &PeerRegistry{ peers: make(map[string]*Peer), path: peersPath, stopChan: make(chan struct{}), authMode: PeerAuthOpen, // Default to open for backward compatibility allowedPublicKeys: make(map[string]bool), } // Try to load existing peers if err := pr.load(); err != nil { // No existing peers, that's ok pr.rebuildKDTree() return pr, nil } pr.rebuildKDTree() return pr, nil } // SetAuthMode sets the authentication mode for peer connections. func (r *PeerRegistry) SetAuthMode(mode PeerAuthMode) { r.allowedPublicKeyMu.Lock() defer r.allowedPublicKeyMu.Unlock() r.authMode = mode logging.Info("peer auth mode changed", logging.Fields{"mode": mode}) } // GetAuthMode returns the current authentication mode. func (r *PeerRegistry) GetAuthMode() PeerAuthMode { r.allowedPublicKeyMu.RLock() defer r.allowedPublicKeyMu.RUnlock() return r.authMode } // AllowPublicKey adds a public key to the allowlist. func (r *PeerRegistry) AllowPublicKey(publicKey string) { r.allowedPublicKeyMu.Lock() defer r.allowedPublicKeyMu.Unlock() r.allowedPublicKeys[publicKey] = true logging.Debug("public key added to allowlist", logging.Fields{"key": safeKeyPrefix(publicKey)}) } // RevokePublicKey removes a public key from the allowlist. func (r *PeerRegistry) RevokePublicKey(publicKey string) { r.allowedPublicKeyMu.Lock() defer r.allowedPublicKeyMu.Unlock() delete(r.allowedPublicKeys, publicKey) logging.Debug("public key removed from allowlist", logging.Fields{"key": safeKeyPrefix(publicKey)}) } // IsPublicKeyAllowed checks if a public key is in the allowlist. func (r *PeerRegistry) IsPublicKeyAllowed(publicKey string) bool { r.allowedPublicKeyMu.RLock() defer r.allowedPublicKeyMu.RUnlock() return r.allowedPublicKeys[publicKey] } // IsPeerAllowed checks if a peer is allowed to connect based on auth mode. // Returns true if: // - AuthMode is Open (allow all) // - AuthMode is Allowlist AND (peer is pre-registered OR public key is allowlisted) func (r *PeerRegistry) IsPeerAllowed(peerID string, publicKey string) bool { r.allowedPublicKeyMu.RLock() authMode := r.authMode keyAllowed := r.allowedPublicKeys[publicKey] r.allowedPublicKeyMu.RUnlock() // Open mode allows everyone if authMode == PeerAuthOpen { return true } // Allowlist mode: check if peer is pre-registered r.mu.RLock() _, isRegistered := r.peers[peerID] r.mu.RUnlock() if isRegistered { return true } // Check if public key is allowlisted return keyAllowed } // ListAllowedPublicKeys returns all allowlisted public keys. func (r *PeerRegistry) ListAllowedPublicKeys() []string { r.allowedPublicKeyMu.RLock() defer r.allowedPublicKeyMu.RUnlock() keys := make([]string, 0, len(r.allowedPublicKeys)) for key := range r.allowedPublicKeys { keys = append(keys, key) } return keys } // AddPeer adds a new peer to the registry. // Note: Persistence is debounced (writes batched every 5s). Call Close() to ensure // all changes are flushed to disk before shutdown. func (r *PeerRegistry) AddPeer(peer *Peer) error { r.mu.Lock() if peer.ID == "" { r.mu.Unlock() return fmt.Errorf("peer ID is required") } // Validate peer name (P2P-LOW-3) if err := validatePeerName(peer.Name); err != nil { r.mu.Unlock() return err } if _, exists := r.peers[peer.ID]; exists { r.mu.Unlock() return fmt.Errorf("peer %s already exists", peer.ID) } // Set defaults if peer.AddedAt.IsZero() { peer.AddedAt = time.Now() } if peer.Score == 0 { peer.Score = 50 // Default neutral score } r.peers[peer.ID] = peer r.rebuildKDTree() r.mu.Unlock() return r.save() } // UpdatePeer updates an existing peer's information. // Note: Persistence is debounced. Call Close() to flush before shutdown. func (r *PeerRegistry) UpdatePeer(peer *Peer) error { r.mu.Lock() if _, exists := r.peers[peer.ID]; !exists { r.mu.Unlock() return fmt.Errorf("peer %s not found", peer.ID) } r.peers[peer.ID] = peer r.rebuildKDTree() r.mu.Unlock() return r.save() } // RemovePeer removes a peer from the registry. // Note: Persistence is debounced. Call Close() to flush before shutdown. func (r *PeerRegistry) RemovePeer(id string) error { r.mu.Lock() if _, exists := r.peers[id]; !exists { r.mu.Unlock() return fmt.Errorf("peer %s not found", id) } delete(r.peers, id) r.rebuildKDTree() r.mu.Unlock() return r.save() } // GetPeer returns a peer by ID. func (r *PeerRegistry) GetPeer(id string) *Peer { r.mu.RLock() defer r.mu.RUnlock() peer, exists := r.peers[id] if !exists { return nil } // Return a copy peerCopy := *peer return &peerCopy } // ListPeers returns all registered peers. func (r *PeerRegistry) ListPeers() []*Peer { r.mu.RLock() defer r.mu.RUnlock() peers := make([]*Peer, 0, len(r.peers)) for _, peer := range r.peers { peerCopy := *peer peers = append(peers, &peerCopy) } return peers } // UpdateMetrics updates a peer's performance metrics. // Note: Persistence is debounced. Call Close() to flush before shutdown. func (r *PeerRegistry) UpdateMetrics(id string, pingMS, geoKM float64, hops int) error { r.mu.Lock() peer, exists := r.peers[id] if !exists { r.mu.Unlock() return fmt.Errorf("peer %s not found", id) } peer.PingMS = pingMS peer.GeoKM = geoKM peer.Hops = hops peer.LastSeen = time.Now() r.rebuildKDTree() r.mu.Unlock() return r.save() } // UpdateScore updates a peer's reliability score. // Note: Persistence is debounced. Call Close() to flush before shutdown. func (r *PeerRegistry) UpdateScore(id string, score float64) error { r.mu.Lock() peer, exists := r.peers[id] if !exists { r.mu.Unlock() return fmt.Errorf("peer %s not found", id) } // Clamp score to 0-100 if score < 0 { score = 0 } else if score > 100 { score = 100 } peer.Score = score r.rebuildKDTree() r.mu.Unlock() return r.save() } // SetConnected updates a peer's connection state. func (r *PeerRegistry) SetConnected(id string, connected bool) { r.mu.Lock() defer r.mu.Unlock() if peer, exists := r.peers[id]; exists { peer.Connected = connected if connected { peer.LastSeen = time.Now() } } } // Score adjustment constants const ( ScoreSuccessIncrement = 1.0 // Increment for successful interaction ScoreFailureDecrement = 5.0 // Decrement for failed interaction ScoreTimeoutDecrement = 3.0 // Decrement for timeout ScoreMinimum = 0.0 // Minimum score ScoreMaximum = 100.0 // Maximum score ScoreDefault = 50.0 // Default score for new peers ) // RecordSuccess records a successful interaction with a peer, improving their score. func (r *PeerRegistry) RecordSuccess(id string) { r.mu.Lock() peer, exists := r.peers[id] if !exists { r.mu.Unlock() return } peer.Score = min(peer.Score+ScoreSuccessIncrement, ScoreMaximum) peer.LastSeen = time.Now() r.mu.Unlock() r.save() } // RecordFailure records a failed interaction with a peer, reducing their score. func (r *PeerRegistry) RecordFailure(id string) { r.mu.Lock() peer, exists := r.peers[id] if !exists { r.mu.Unlock() return } peer.Score = max(peer.Score-ScoreFailureDecrement, ScoreMinimum) newScore := peer.Score r.mu.Unlock() r.save() logging.Debug("peer score decreased", logging.Fields{ "peer_id": id, "new_score": newScore, "reason": "failure", }) } // RecordTimeout records a timeout when communicating with a peer. func (r *PeerRegistry) RecordTimeout(id string) { r.mu.Lock() peer, exists := r.peers[id] if !exists { r.mu.Unlock() return } peer.Score = max(peer.Score-ScoreTimeoutDecrement, ScoreMinimum) newScore := peer.Score r.mu.Unlock() r.save() logging.Debug("peer score decreased", logging.Fields{ "peer_id": id, "new_score": newScore, "reason": "timeout", }) } // GetPeersByScore returns peers sorted by score (highest first). func (r *PeerRegistry) GetPeersByScore() []*Peer { r.mu.RLock() defer r.mu.RUnlock() peers := make([]*Peer, 0, len(r.peers)) for _, p := range r.peers { peers = append(peers, p) } // Sort by score descending for i := 0; i < len(peers)-1; i++ { for j := i + 1; j < len(peers); j++ { if peers[j].Score > peers[i].Score { peers[i], peers[j] = peers[j], peers[i] } } } return peers } // SelectOptimalPeer returns the best peer based on multi-factor optimization. // Uses Poindexter KD-tree to find the peer closest to ideal metrics. func (r *PeerRegistry) SelectOptimalPeer() *Peer { r.mu.RLock() defer r.mu.RUnlock() if r.kdTree == nil || len(r.peers) == 0 { return nil } // Target: ideal peer (0 ping, 0 hops, 0 geo, 100 score) // Score is inverted (100 - score) so lower is better in the tree target := []float64{0, 0, 0, 0} result, _, found := r.kdTree.Nearest(target) if !found { return nil } peer, exists := r.peers[result.Value] if !exists { return nil } peerCopy := *peer return &peerCopy } // SelectNearestPeers returns the n best peers based on multi-factor optimization. func (r *PeerRegistry) SelectNearestPeers(n int) []*Peer { r.mu.RLock() defer r.mu.RUnlock() if r.kdTree == nil || len(r.peers) == 0 { return nil } // Target: ideal peer target := []float64{0, 0, 0, 0} results, _ := r.kdTree.KNearest(target, n) peers := make([]*Peer, 0, len(results)) for _, result := range results { if peer, exists := r.peers[result.Value]; exists { peerCopy := *peer peers = append(peers, &peerCopy) } } return peers } // GetConnectedPeers returns all currently connected peers. func (r *PeerRegistry) GetConnectedPeers() []*Peer { r.mu.RLock() defer r.mu.RUnlock() peers := make([]*Peer, 0) for _, peer := range r.peers { if peer.Connected { peerCopy := *peer peers = append(peers, &peerCopy) } } return peers } // Count returns the number of registered peers. func (r *PeerRegistry) Count() int { r.mu.RLock() defer r.mu.RUnlock() return len(r.peers) } // rebuildKDTree rebuilds the KD-tree from current peers. // Must be called with lock held. func (r *PeerRegistry) rebuildKDTree() { if len(r.peers) == 0 { r.kdTree = nil return } points := make([]poindexter.KDPoint[string], 0, len(r.peers)) for _, peer := range r.peers { // Build 4D point with weighted, normalized values // Invert score so that higher score = lower value (better) point := poindexter.KDPoint[string]{ ID: peer.ID, Coords: []float64{ peer.PingMS * pingWeight, float64(peer.Hops) * hopsWeight, peer.GeoKM * geoWeight, (100 - peer.Score) * scoreWeight, // Invert score }, Value: peer.ID, } points = append(points, point) } // Build KD-tree with Euclidean distance tree, err := poindexter.NewKDTree(points, poindexter.WithMetric(poindexter.EuclideanDistance{})) if err != nil { // Log error but continue - worst case we don't have optimal selection return } r.kdTree = tree } // scheduleSave schedules a debounced save operation. // Multiple calls within saveDebounceInterval will be coalesced into a single save. // Must NOT be called with r.mu held. func (r *PeerRegistry) scheduleSave() { r.saveMu.Lock() defer r.saveMu.Unlock() r.dirty = true // If timer already running, let it handle the save if r.saveTimer != nil { return } // Start a new timer r.saveTimer = time.AfterFunc(saveDebounceInterval, func() { r.saveMu.Lock() r.saveTimer = nil shouldSave := r.dirty r.dirty = false r.saveMu.Unlock() if shouldSave { r.mu.RLock() err := r.saveNow() r.mu.RUnlock() if err != nil { // Log error but continue - best effort persistence logging.Warn("failed to save peer registry", logging.Fields{"error": err}) } } }) } // saveNow persists peers to disk immediately. // Must be called with r.mu held (at least RLock). func (r *PeerRegistry) saveNow() error { // Ensure directory exists dir := filepath.Dir(r.path) if err := os.MkdirAll(dir, 0755); err != nil { return fmt.Errorf("failed to create peers directory: %w", err) } // Convert to slice for JSON peers := make([]*Peer, 0, len(r.peers)) for _, peer := range r.peers { peers = append(peers, peer) } data, err := json.MarshalIndent(peers, "", " ") if err != nil { return fmt.Errorf("failed to marshal peers: %w", err) } // Use atomic write pattern: write to temp file, then rename tmpPath := r.path + ".tmp" if err := os.WriteFile(tmpPath, data, 0644); err != nil { return fmt.Errorf("failed to write peers temp file: %w", err) } if err := os.Rename(tmpPath, r.path); err != nil { os.Remove(tmpPath) // Clean up temp file return fmt.Errorf("failed to rename peers file: %w", err) } return nil } // Close flushes any pending changes and releases resources. func (r *PeerRegistry) Close() error { r.saveStopOnce.Do(func() { close(r.stopChan) }) // Cancel pending timer and save immediately if dirty r.saveMu.Lock() if r.saveTimer != nil { r.saveTimer.Stop() r.saveTimer = nil } shouldSave := r.dirty r.dirty = false r.saveMu.Unlock() if shouldSave { r.mu.RLock() err := r.saveNow() r.mu.RUnlock() return err } return nil } // save is a helper that schedules a debounced save. // Kept for backward compatibility but now debounces writes. // Must NOT be called with r.mu held. func (r *PeerRegistry) save() error { r.scheduleSave() return nil // Errors will be logged asynchronously } // load reads peers from disk. func (r *PeerRegistry) load() error { data, err := os.ReadFile(r.path) if err != nil { return fmt.Errorf("failed to read peers: %w", err) } var peers []*Peer if err := json.Unmarshal(data, &peers); err != nil { return fmt.Errorf("failed to unmarshal peers: %w", err) } r.peers = make(map[string]*Peer) for _, peer := range peers { r.peers[peer.ID] = peer } return nil } // Example usage inside a connection handler