chains/mpcvm/transport.go

// Copyright (C) 2019-2025, Lux Industries Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package mpcvm

// transport.go — the MessageRouter that lets an MPC ceremony (DKG, sign) run
// NATIVELY across the M-Chain's validators instead of an off-chain REST
// cluster. The executor drives a per-party threshold protocol.Handler; the
// router carries that handler's messages to the other parties and delivers
// theirs back.
//
// There is ONE router. It serialises a ceremony message into a session-tagged
// envelope and broadcasts it over the chain's app-gossip (or sends it to one
// peer); the receiving validator's Gossip handler demultiplexes by session and
// calls Deliver. Decoupled from node types via the small AppNetwork interface,
// so this file imports no node packages and the send side is testable against
// an in-memory fabric.
//
// threshold stays the shared library — the router only moves its opaque
// protocol.Message bytes; it never interprets the cryptography.

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"

	"github.com/luxfi/ids"
	"github.com/luxfi/math/set"
	"github.com/luxfi/threshold/pkg/party"
	"github.com/luxfi/threshold/pkg/protocol"
	"github.com/luxfi/warp"
)

// =============================================================================
// Cross-validator gossip transport
// =============================================================================

// AppNetwork is the minimal consensus-gossip send capability the gossipRouter
// needs. The VM implements it over the node AppSender it receives at
// Initialize. Kept as an interface so this file imports no node packages and
// stays unit-testable.
type AppNetwork interface {
	// Broadcast reliably gossips an app message to the chain's validators.
	Broadcast(ctx context.Context, msg []byte) error
	// SendTo sends an app message to a single validator.
	SendTo(ctx context.Context, nodeID ids.NodeID, msg []byte) error
}

// sessionEnvelope wraps one ceremony message with its session id so a receiving
// validator can demultiplex concurrent ceremonies to the right handler.
type sessionEnvelope struct {
	SessionID string `json:"sessionId"`
	Message   []byte `json:"message"` // protocol.Message, wire-encoded
}

// marshalEnvelope encodes a ceremony message for gossip.
func marshalEnvelope(sessionID string, msg *protocol.Message) ([]byte, error) {
	wire, err := msg.MarshalBinary()
	if err != nil {
		return nil, fmt.Errorf("mpcvm: marshal ceremony message: %w", err)
	}
	return json.Marshal(sessionEnvelope{SessionID: sessionID, Message: wire})
}

// unmarshalEnvelope decodes a gossiped ceremony message.
func unmarshalEnvelope(b []byte) (string, *protocol.Message, error) {
	var env sessionEnvelope
	if err := json.Unmarshal(b, &env); err != nil {
		return "", nil, fmt.Errorf("mpcvm: unmarshal envelope: %w", err)
	}
	msg := &protocol.Message{}
	if err := msg.UnmarshalBinary(env.Message); err != nil {
		return "", nil, fmt.Errorf("mpcvm: unmarshal ceremony message: %w", err)
	}
	return env.SessionID, msg, nil
}

// gossipRouter carries one party's ceremony messages over consensus gossip and
// feeds incoming ones (delivered by the VM's app-message handler) to the local
// handler. One gossipRouter per (session, self-party).
type gossipRouter struct {
	ctx       context.Context
	sessionID string
	self      party.ID
	net       AppNetwork
	// peers maps a threshold party.ID to the validator NodeID that runs it, so
	// a directed (To != "") ceremony message becomes an app-request to one peer.
	peers map[party.ID]ids.NodeID
	inbox chan *protocol.Message

	// mu guards closed so a late Deliver (from a peer still gossiping as the
	// session tears down) can never send on a closed inbox.
	mu     sync.Mutex
	closed bool
}

// newGossipRouter builds the cross-validator router for one session.
func newGossipRouter(
	ctx context.Context,
	sessionID string,
	self party.ID,
	net AppNetwork,
	peers map[party.ID]ids.NodeID,
) *gossipRouter {
	return &gossipRouter{
		ctx:       ctx,
		sessionID: sessionID,
		self:      self,
		net:       net,
		peers:     peers,
		inbox:     make(chan *protocol.Message, 1024),
	}
}

// Send implements MessageRouter: broadcast when To=="", else a directed send to
// the validator that runs the addressed party.
func (r *gossipRouter) Send(msg *protocol.Message) error {
	if msg == nil {
		return nil
	}
	env, err := marshalEnvelope(r.sessionID, msg)
	if err != nil {
		return err
	}
	if msg.To == "" {
		return r.net.Broadcast(r.ctx, env)
	}
	nodeID, ok := r.peers[msg.To]
	if !ok {
		return fmt.Errorf("mpcvm: no validator for party %s", msg.To)
	}
	return r.net.SendTo(r.ctx, nodeID, env)
}

// Receive implements MessageRouter.
func (r *gossipRouter) Receive() <-chan *protocol.Message { return r.inbox }

// Deliver hands an incoming ceremony message (already demultiplexed to this
// session) to the local handler. Called by the VM's app-message handler; drops
// on a full inbox rather than blocking consensus, and is a no-op once the
// session is closed (never sends on a closed channel).
func (r *gossipRouter) Deliver(msg *protocol.Message) {
	if msg == nil || msg.From == r.self {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		return
	}
	select {
	case r.inbox <- msg:
	default:
	}
}

// Close releases the receive channel when the session ends. Idempotent.
func (r *gossipRouter) Close() {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		return
	}
	r.closed = true
	close(r.inbox)
}

var _ MessageRouter = (*gossipRouter)(nil)

// =============================================================================
// warp.Sender-backed AppNetwork (the production send side)
// =============================================================================

// warpAppNetwork adapts the node's canonical warp.Sender (handed to the VM at
// Initialize as block.Init.Sender) to the AppNetwork the gossipRouter sends
// over. It is constructed PER CEREMONY with the exact set of committee node IDs
// that must receive the ceremony's messages: an MPC round is only sound if
// EVERY listed signer receives EVERY broadcast, so "Broadcast" here is an
// explicit multicast to the committee — never a validator SAMPLE (sampling a
// subset would silently starve a signer and stall the round). Both Broadcast
// and SendTo travel as app-gossip so they arrive at the receiver's single
// Gossip handler, which demultiplexes by session and calls the router's
// Deliver — one receive path, one demux point.
//
// Delivery is best-effort at the p2p layer: a dropped gossip datagram stalls a
// round until the threshold handler's RoundTimeout fires. Production hardening
// (reliable delivery via acks/retransmit, or SendAppRequest with retry) is
// tracked separately; the in-memory fabric used by the transport tests delivers
// synchronously so the ceremony is deterministic under test.
type warpAppNetwork struct {
	sender    warp.Sender
	self      ids.NodeID
	committee set.Set[ids.NodeID] // every committee node (self included; excluded on send)
}

// newWarpAppNetwork binds a ceremony's committee to the node's warp sender.
func newWarpAppNetwork(sender warp.Sender, self ids.NodeID, committee []ids.NodeID) *warpAppNetwork {
	s := set.NewSet[ids.NodeID](len(committee))
	s.Add(committee...)
	return &warpAppNetwork{sender: sender, self: self, committee: s}
}

// Broadcast multicasts msg to every committee node except self.
func (n *warpAppNetwork) Broadcast(ctx context.Context, msg []byte) error {
	if n.sender == nil {
		return fmt.Errorf("mpcvm: no warp sender configured")
	}
	targets := set.NewSet[ids.NodeID](n.committee.Len())
	targets.Add(n.committee.List()...)
	targets.Remove(n.self)
	if targets.Len() == 0 {
		return nil
	}
	return n.sender.SendGossip(ctx, warp.SendConfig{NodeIDs: targets}, msg)
}

// SendTo delivers msg to a single committee node.
func (n *warpAppNetwork) SendTo(ctx context.Context, nodeID ids.NodeID, msg []byte) error {
	if n.sender == nil {
		return fmt.Errorf("mpcvm: no warp sender configured")
	}
	if nodeID == n.self {
		return nil
	}
	return n.sender.SendGossip(ctx, warp.SendConfig{NodeIDs: set.Of(nodeID)}, msg)
}

// Compile-time check that warpAppNetwork satisfies AppNetwork.
var _ AppNetwork = (*warpAppNetwork)(nil)