vms/txs/mempool: add AdmissionVerifier hook for encrypted-payload txs (#115)

vms/txs/mempool/mempool.go

@@ -37,6 +37,11 @@ var (
 	ErrTxTooLarge           = errors.New("tx too large")
 	ErrMempoolFull          = errors.New("mempool is full")
 	ErrConflictsWithOtherTx = errors.New("tx conflicts with other tx")
+	// ErrAdmissionRejected wraps a rejection emitted by an AdmissionVerifier
+	// configured via NewWithAdmissionVerifier. The wrapped error carries the
+	// verifier's specific reason (e.g. NIZK proof invalid, fee bid below
+	// floor, budget meter exhausted).
+	ErrAdmissionRejected = errors.New("tx admission rejected")
 )
 
 type Tx interface {
@@ -45,6 +50,28 @@ type Tx interface {
 	Size() int
 }
 
+// AdmissionVerifier is an optional admission gate. When configured via
+// NewWithAdmissionVerifier the mempool calls VerifyAdmit on every Add()
+// after the cheap structural checks (duplicate / size / space / conflict)
+// pass and before the tx is inserted. A non-nil return rejects the tx; the
+// returned error is wrapped in ErrAdmissionRejected and recorded as the
+// drop reason.
+//
+// Intended consumers: validators that admit encrypted-payload transactions
+// (FHE ciphertext + NIZK proof of well-formedness) without decryption, per
+// LP-066 / luxfi/precompile/fhe. The verifier runs signature verify + fee
+// check + NIZK verify; the actual decryption never happens at admission
+// time.
+//
+// VerifyAdmit MUST be safe to call without holding any mempool lock — the
+// mempool holds its own write lock across Add() and the call is made from
+// within that critical section. Implementations that need to do expensive
+// work (NIZK verification) should not perform additional locking that
+// could re-enter the mempool.
+type AdmissionVerifier[T Tx] interface {
+	VerifyAdmit(tx T) error
+}
+
 type Metrics interface {
 	Update(numTxs, bytesAvailable int)
 }
@@ -83,17 +110,40 @@ type mempool[T Tx] struct {
 	droppedTxIDs   *lru.Cache[ids.ID, error] // TxID -> Verification error
 
 	metrics Metrics
+
+	// admissionVerifier is nil for the default New constructor (preserving
+	// the existing admission policy of signature/fee verification happening
+	// at the caller). When non-nil it is invoked from Add() after the cheap
+	// checks pass.
+	admissionVerifier AdmissionVerifier[T]
 }
 
 func New[T Tx](
 	metrics Metrics,
+) *mempool[T] {
+	return NewWithAdmissionVerifier[T](metrics, nil)
+}
+
+// NewWithAdmissionVerifier constructs a mempool that runs verifier.VerifyAdmit
+// on every Add() after the duplicate / size / space / conflict checks
+// succeed. A nil verifier is equivalent to New — no admission gate is
+// installed and behavior is byte-identical to the prior API.
+//
+// This is the entry point for encrypted-payload tx pools per luxfi/node#115:
+// validators construct a mempool partition with a verifier that runs
+// signature + fee + NIZK checks (sourced from luxfi/precompile/fhe) so
+// ciphertexts are admitted without decryption.
+func NewWithAdmissionVerifier[T Tx](
+	metrics Metrics,
+	verifier AdmissionVerifier[T],
 ) *mempool[T] {
 	m := &mempool[T]{
-		unissuedTxs:    linked.NewHashmap[ids.ID, T](),
-		consumedUTXOs:  setmap.New[ids.ID, ids.ID](),
-		bytesAvailable: maxMempoolSize,
-		droppedTxIDs:   lru.NewCache[ids.ID, error](droppedTxIDsCacheSize),
-		metrics:        metrics,
+		unissuedTxs:       linked.NewHashmap[ids.ID, T](),
+		consumedUTXOs:     setmap.New[ids.ID, ids.ID](),
+		bytesAvailable:    maxMempoolSize,
+		droppedTxIDs:      lru.NewCache[ids.ID, error](droppedTxIDsCacheSize),
+		metrics:           metrics,
+		admissionVerifier: verifier,
 	}
 	m.cond = lock.NewCond(&m.lock)
 	m.updateMetrics()
@@ -137,6 +187,20 @@ func (m *mempool[T]) Add(tx T) error {
 		return fmt.Errorf("%w: %s", ErrConflictsWithOtherTx, txID)
 	}
 
+	// Admission gate runs last among Add() checks: it is the most expensive
+	// (e.g. NIZK verify for encrypted-payload txs per LP-066). All cheap
+	// rejects have fired by this point so verification cost is only paid on
+	// txs that pass structural admission.
+	if m.admissionVerifier != nil {
+		if verr := m.admissionVerifier.VerifyAdmit(tx); verr != nil {
+			wrapped := fmt.Errorf("%w: %s: %w", ErrAdmissionRejected, txID, verr)
+			// Record the drop reason so downstream propagation / inspection
+			// surfaces match the existing dropped-tx tracking contract.
+			m.droppedTxIDs.Put(txID, wrapped)
+			return wrapped
+		}
+	}
+
 	m.bytesAvailable -= txSize
 	m.unissuedTxs.Put(txID, tx)
 	m.updateMetrics()

vms/txs/mempool/mempool_test.go

@@ -325,3 +325,119 @@ func TestWaitForEventWithTx(t *testing.T) {
 	require.Equal(vmcore.PendingTxs, msg.Type)
 	require.NoError(<-errs)
 }
+
+// -----------------------------------------------------------------------------
+// AdmissionVerifier tests (luxfi/node#115)
+// -----------------------------------------------------------------------------
+
+// fakeVerifier implements AdmissionVerifier[*dummyTx] for the tests below.
+// It tracks every call so tests can assert call counts and inspect the txs
+// that the gate observed.
+type fakeVerifier struct {
+	err  error
+	seen []ids.ID
+}
+
+func (v *fakeVerifier) VerifyAdmit(tx *dummyTx) error {
+	v.seen = append(v.seen, tx.ID())
+	return v.err
+}
+
+// Nil verifier path: NewWithAdmissionVerifier with nil must behave
+// byte-identically to New. We assert Add succeeds and Len reflects the tx.
+func TestAdmissionVerifier_nilVerifierMatchesNew(t *testing.T) {
+	require := require.New(t)
+
+	m := NewWithAdmissionVerifier[*dummyTx](&noMetrics{}, nil)
+	tx := newTx(0, 32)
+	require.NoError(m.Add(tx))
+	require.Equal(1, m.Len())
+}
+
+// Verifier returning nil: tx is admitted. The gate must run exactly once
+// per Add (not on Get, Peek, Iterate, or Remove).
+func TestAdmissionVerifier_acceptsWhenVerifierReturnsNil(t *testing.T) {
+	require := require.New(t)
+
+	v := &fakeVerifier{err: nil}
+	m := NewWithAdmissionVerifier[*dummyTx](&noMetrics{}, v)
+	tx := newTx(0, 32)
+	require.NoError(m.Add(tx))
+	require.Equal(1, m.Len())
+	require.Len(v.seen, 1)
+	require.Equal(tx.ID(), v.seen[0])
+
+	// Get / Peek / Iterate do not re-run the verifier.
+	_, _ = m.Get(tx.ID())
+	_, _ = m.Peek()
+	m.Iterate(func(*dummyTx) bool { return true })
+	require.Len(v.seen, 1, "verifier must run only on Add")
+
+	// Remove does not re-run the verifier.
+	m.Remove(tx)
+	require.Len(v.seen, 1, "Remove must not invoke the verifier")
+}
+
+// Verifier returning an error: tx is rejected, the returned error wraps
+// ErrAdmissionRejected and carries the verifier's reason, and the drop
+// reason is recorded.
+func TestAdmissionVerifier_rejectsWhenVerifierReturnsError(t *testing.T) {
+	require := require.New(t)
+
+	verifyErr := errors.New("nizk proof invalid")
+	v := &fakeVerifier{err: verifyErr}
+	m := NewWithAdmissionVerifier[*dummyTx](&noMetrics{}, v)
+	tx := newTx(0, 32)
+
+	addErr := m.Add(tx)
+	require.Error(addErr)
+	require.ErrorIs(addErr, ErrAdmissionRejected, "outer error must wrap ErrAdmissionRejected")
+	require.ErrorIs(addErr, verifyErr, "outer error must wrap the verifier's reason")
+
+	require.Equal(0, m.Len(), "rejected tx must not be inserted")
+
+	dropReason := m.GetDropReason(tx.ID())
+	require.Error(dropReason, "rejected tx must have a recorded drop reason")
+	require.ErrorIs(dropReason, ErrAdmissionRejected)
+	require.ErrorIs(dropReason, verifyErr)
+
+	require.Len(v.seen, 1, "verifier must be invoked exactly once")
+}
+
+// Cheap checks short-circuit before the verifier runs. We exercise the
+// three cheap reject paths (duplicate, oversize, conflict) and assert the
+// verifier never observed those txs — verification cost should not be paid
+// on a tx that would have been dropped anyway.
+func TestAdmissionVerifier_cheapChecksShortCircuit(t *testing.T) {
+	require := require.New(t)
+
+	v := &fakeVerifier{err: nil}
+	m := NewWithAdmissionVerifier[*dummyTx](&noMetrics{}, v)
+
+	// 1. Duplicate: first Add admits and runs the verifier once; second Add
+	// must reject with ErrDuplicateTx without re-running it.
+	tx := newTx(0, 32)
+	require.NoError(m.Add(tx))
+	require.Len(v.seen, 1)
+	dupErr := m.Add(tx)
+	require.ErrorIs(dupErr, ErrDuplicateTx)
+	require.Len(v.seen, 1, "verifier must not run on duplicate")
+
+	// 2. Oversize: tx larger than MaxTxSize is rejected before the
+	// verifier sees it.
+	bigTx := newTx(99, MaxTxSize+1)
+	bigErr := m.Add(bigTx)
+	require.ErrorIs(bigErr, ErrTxTooLarge)
+	require.Len(v.seen, 1, "verifier must not run on oversize tx")
+
+	// 3. Conflict: a second tx consuming the same UTXO as the admitted tx
+	// is rejected before the verifier sees it.
+	conflictTx := &dummyTx{
+		id:       ids.GenerateTestID(),
+		size:     32,
+		inputIDs: tx.inputIDs, // same inputs as the admitted tx
+	}
+	conflictErr := m.Add(conflictTx)
+	require.ErrorIs(conflictErr, ErrConflictsWithOtherTx)
+	require.Len(v.seen, 1, "verifier must not run on conflicting tx")
+}