chains/zkvm/verifier.go

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

package zkvm

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"math/big"
	"sync"

	"github.com/luxfi/log"
	"github.com/luxfi/precompile/starkfri"

	"github.com/consensys/gnark-crypto/ecc/bn254"
	"github.com/consensys/gnark-crypto/ecc/bn254/fr"
	lru "github.com/hashicorp/golang-lru"
)

// errStrictPQClassicalForbidden is returned by the shielded-tx verifier
// when a strict-PQ chain is asked to verify a classical (bn254
// pairing-based) proof system. On a strict-PQ chain only the STARK/FRI
// system is accepted; groth16/plonk/bulletproofs are hard-disabled.
var errStrictPQClassicalForbidden = errors.New(
	"zkvm: classical proof system forbidden on strict-PQ chain — only STARK/FRI accepted")

// errStrictPQRealVKForbidden is returned at construction when a strict-PQ
// chain attempts to load a real (non-dummy) bn254 verifying key. Real
// classical VKs on a strict-PQ chain would re-enable the forgeable
// pairing path; refuse explicitly rather than rely on dummy-key detection.
var errStrictPQRealVKForbidden = errors.New(
	"zkvm: real bn254 verifying key forbidden on strict-PQ chain — shielded value uses STARK/FRI (P3Q) only")

// ProofVerifier verifies zero-knowledge proofs.
// When verifying keys are all zeros (dummy), proof verification is disabled
// and VerifyProof returns an error. This is fail-closed by design.
type ProofVerifier struct {
	config ZConfig
	log    log.Logger

	// bind is sha256(ChainID ‖ NetworkID). It is the first public input every
	// proof is checked against, so a proof made for another chain does not
	// verify here even when the notes it names are unspent on both.
	bind [32]byte

	// Proof verification cache
	proofCache *lru.Cache

	// Verifying keys
	verifyingKeys map[string][]byte // circuit type -> verifying key
	dummyKeys     bool              // true if all verifying keys are zero-filled

	// Statistics
	verifyCount uint64
	cacheHits   uint64
	cacheMisses uint64

	mu sync.RWMutex
}

// NewProofVerifier creates a new proof verifier
func NewProofVerifier(config ZConfig, bind [32]byte, log log.Logger) (*ProofVerifier, error) {
	// Create LRU cache for proof verification results
	cache, err := lru.New(int(config.ProofCacheSize))
	if err != nil {
		return nil, err
	}

	pv := &ProofVerifier{
		config:        config,
		bind:          bind,
		log:           log,
		proofCache:    cache,
		verifyingKeys: make(map[string][]byte),
	}

	// Load verifying keys
	if err := pv.loadVerifyingKeys(); err != nil {
		return nil, err
	}

	return pv, nil
}

// VerifyTransactionProof verifies a transaction's zero-knowledge proof.
// Returns an error if verifying keys are dummy (all zeros).
func (pv *ProofVerifier) VerifyTransactionProof(tx *Transaction) error {
	if tx.Proof == nil {
		return errors.New("transaction missing proof")
	}

	// Strict-PQ gate (Red H1). On a strict-PQ chain the classical
	// (bn254 pairing-based) systems are hard-disabled BEFORE any other
	// check — a CRQC that breaks bn254 must not be able to forge a
	// shield/unshield proof. Only STARK/FRI is accepted there.
	if err := pv.refuseClassicalUnderStrictPQ(tx.Proof.ProofType); err != nil {
		return err
	}

	// STARK/FRI path: delegate to the strict-PQ verifier (precompile/
	// starkfri). This is the ONLY accepted system on a strict-PQ chain
	// and is the quantum-safe shielded-proof path everywhere. It fails
	// closed (errVerifierUnavailable-equivalent) when no FRI binding is
	// registered — see verifySTARKProof.
	if tx.Proof.ProofType == "stark" {
		return pv.verifySTARKProof(tx)
	}

	// Classical path (non-strict chains only — the gate above already
	// rejected these under strict-PQ). Requires real (non-dummy) VKs.
	if pv.dummyKeys {
		return errors.New("zkvm: proof verification disabled — no real verifying keys loaded")
	}

	// The cache is keyed on the transaction's CONTENT — ComputeID covers the
	// nullifiers, the outputs and the proof — so a hit means this exact
	// transaction verified before.
	//
	// It used to be keyed on tx.ID as READ OFF THE WIRE, plus the proof and
	// its public inputs. All four were the peer's to choose: copy an accepted
	// transaction's id, proof type, proof and public inputs onto a transaction
	// spending entirely different notes, and the key matched. The hit returned
	// nil from here — before verifyPublicInputs, the only thing tying the
	// public inputs to what the transaction spends — and the shielded pool
	// gained value backed by a proof that attests to nothing about it.
	key := tx.ComputeID()

	pv.mu.Lock()
	pv.verifyCount++

	if cached, ok := pv.proofCache.Get(key); ok {
		pv.cacheHits++
		pv.mu.Unlock()

		if cached.(bool) {
			return nil
		}
		return errors.New("proof verification failed (cached)")
	}
	pv.cacheMisses++
	pv.mu.Unlock()

	// Verify proof based on type
	var err error
	switch tx.Proof.ProofType {
	case "groth16":
		err = pv.verifyGroth16Proof(tx)
	case "plonk":
		err = errPLONKIncomplete
	case "bulletproofs":
		err = errors.New("zkvm: Bulletproof verification not yet implemented, use groth16 or plonk")
	default:
		err = errors.New("unsupported proof type")
	}

	// Cache result
	pv.proofCache.Add(key, err == nil)

	return err
}

// refuseClassicalUnderStrictPQ is the single strict-PQ enforcement point
// for the shielded-tx verifier. On a strict-PQ chain it refuses every
// classical (quantum-breakable) proof system, leaving only "stark"
// (STARK/FRI). On a non-strict chain it is a no-op. Both VerifyTransactionProof
// and the GPU batch path call this so neither can verify a classical proof
// on a strict-PQ chain.
func (pv *ProofVerifier) refuseClassicalUnderStrictPQ(proofType string) error {
	if !pv.config.StrictPQ {
		return nil
	}
	switch proofType {
	case "stark":
		return nil
	default:
		// groth16, plonk, bulletproofs, anything else: forbidden.
		return errStrictPQClassicalForbidden
	}
}

// verifySTARKProof verifies a strict-PQ STARK/FRI shielded proof by
// delegating to precompile/starkfri (cSHAKE256 Merkle over Goldilocks,
// FRI low-degree test — no pairings, no bn254, no trusted setup). It
// fails closed when no FRI verifier binding is registered (CGO_ENABLED=0
// or the starkfri_p3q tag absent): a structurally well-formed proof is
// NEVER accepted without the real verifier, so there is no forgery oracle.
//
// The proof bytes are tx.Proof.ProofData (which must begin with the
// starkfri MagicHeader); the public inputs are the canonical
// concatenation of the tx's nullifiers and output commitments, binding
// the proof to this transaction's shielded value flow.
//
// PROVER STATUS (tracked, not faked). The FULL post-quantum shielded
// path also needs the prover side: a `p3q_prove` C ABI (Rust prover
// behind the starkfri_p3q cgo tag) and a shielded/ML-DSA AIR that
// arithmetises the Zcash-style spend/output circuit (note commitments,
// nullifier derivation, value balance, range proofs) over the Goldilocks
// field. Until that AIR + prover land, this verifier accepts no shielded
// proof on a strict-PQ chain (fail-closed), and shielded value transfer
// is effectively disabled there — which is the correct posture: no
// classical fallback, no forgeable path.
func (pv *ProofVerifier) verifySTARKProof(tx *Transaction) error {
	// Bind the proof to this tx's shielded value via the public inputs:
	// nullifiers (spends) ‖ output commitments (outputs).
	pub := make([]byte, 0, 96)
	pub = append(pub, pv.bind[:]...)
	for _, n := range tx.Nullifiers {
		pub = append(pub, n...)
	}
	for _, c := range tx.GetOutputCommitments() {
		pub = append(pub, c...)
	}

	ok, err := starkfri.Verify(tx.Proof.ProofData, pub)
	if err != nil {
		if errors.Is(err, starkfri.ErrVerifierNotRegistered) {
			// Binding pending: fail closed, distinguishable for operators.
			return fmt.Errorf("zkvm: strict-PQ STARK shielded verifier unbound (build with -tags starkfri_p3q): %w", err)
		}
		// Malformed / non-verifying proof.
		return fmt.Errorf("zkvm: STARK shielded proof verification failed: %w", err)
	}
	if !ok {
		return errors.New("zkvm: STARK shielded proof rejected")
	}
	return nil
}

// verifyGroth16Proof verifies a Groth16 proof using gnark
func (pv *ProofVerifier) verifyGroth16Proof(tx *Transaction) error {
	// Get verifying key for circuit type
	vkBytes, exists := pv.verifyingKeys[string(tx.Type)]
	if !exists {
		return fmt.Errorf("zkvm: no verifying key for circuit %q", tx.Type)
	}

	// Verify public inputs match transaction data
	if err := pv.verifyPublicInputs(tx); err != nil {
		return err
	}

	// Validate proof data length (Groth16: 2 G1 points + 1 G2 point)
	// BN254: G1 = 64 bytes (compressed), G2 = 128 bytes (compressed)
	// Total: 2*64 + 128 = 256 bytes minimum
	if len(tx.Proof.ProofData) < 256 {
		return errors.New("invalid proof data length for Groth16")
	}

	// Perform actual Groth16 verification using gnark-crypto
	if err := pv.verifyGroth16WithGnark(tx.Proof, vkBytes); err != nil {
		return fmt.Errorf("groth16 verification failed: %w", err)
	}

	pv.log.Debug("Groth16 proof verified",
		log.String("txID", tx.ID.String()),
		log.Int("vkLen", len(vkBytes)),
	)

	return nil
}

// verifyPublicInputs verifies that public inputs match transaction data.
//
// The first input is the chain binding, so a proof made for another chain does
// not verify here even when the notes it names happen to be unspent on both.
func (pv *ProofVerifier) verifyPublicInputs(tx *Transaction) error {
	if len(tx.Proof.PublicInputs) == 0 {
		return errors.New("no public inputs provided")
	}
	if !bytes.Equal(tx.Proof.PublicInputs[0], pv.bind[:]) {
		return errors.New("public input mismatch for chain binding")
	}

	// Verify nullifiers are included in public inputs (exact byte comparison)
	for i, nullifier := range tx.Nullifiers {
		if i+1 >= len(tx.Proof.PublicInputs) {
			return errors.New("missing public input for nullifier")
		}

		if !bytes.Equal(tx.Proof.PublicInputs[i+1], nullifier) {
			return errors.New("public input mismatch for nullifier")
		}
	}

	// Verify output commitments are included (exact byte comparison)
	outputCommitments := tx.GetOutputCommitments()
	offset := len(tx.Nullifiers) + 1

	for i, commitment := range outputCommitments {
		idx := offset + i
		if idx >= len(tx.Proof.PublicInputs) {
			return errors.New("missing public input for output commitment")
		}

		if !bytes.Equal(tx.Proof.PublicInputs[idx], commitment) {
			return errors.New("public input mismatch for output commitment")
		}
	}

	return nil
}

// loadVerifyingKeys takes the verifying key for each circuit the config keys and
// installs nothing for the others.
//
// A circuit the operator did not key has no key here, so the lookup in front of
// each verify refuses it by name. That is the one place the question is asked,
// and it is asked per circuit: a key is judged on its own, never on what the
// verifier holds for some other circuit. An absent key cannot be mistaken for a
// present one.
func (pv *ProofVerifier) loadVerifyingKeys() error {
	for _, ct := range []TransactionType{
		TransactionTypeTransfer,
		TransactionTypeShield,
		TransactionTypeUnshield,
	} {
		// Copied, so the verifier owns its bytes.
		if vk, ok := pv.config.VerifyingKeys[string(ct)]; ok && len(vk) > 0 {
			pv.verifyingKeys[string(ct)] = append([]byte(nil), vk...)
		}
	}

	// A verifier holding no key at all judges nothing, which is what the
	// classical path checks before it starts.
	pv.dummyKeys = len(pv.verifyingKeys) == 0

	// Strict-PQ hard gate (Red H1). Loading a REAL (non-dummy) bn254
	// verifying key on a strict-PQ chain is forbidden: such keys would
	// re-enable the forgeable classical pairing path for shielded value.
	// Refuse explicitly at construction rather than relying on the
	// implicit dummy-key detector — the shielded path on a strict-PQ
	// chain is STARK/FRI only.
	if pv.config.StrictPQ && !pv.dummyKeys {
		return errStrictPQRealVKForbidden
	}

	pv.log.Info("Loaded verifying keys",
		log.Int("count", len(pv.verifyingKeys)),
		log.Bool("strictPQ", pv.config.StrictPQ),
	)

	return nil
}

// VerifyingKeysLoaded returns true if real (non-dummy) verifying keys are loaded.
func (pv *ProofVerifier) VerifyingKeysLoaded() bool {
	return !pv.dummyKeys
}

// GetCacheSize returns the current size of the proof cache
func (pv *ProofVerifier) GetCacheSize() int {
	return pv.proofCache.Len()
}

// GetStats returns verifier statistics
func (pv *ProofVerifier) GetStats() (verifyCount, cacheHits, cacheMisses uint64) {
	pv.mu.RLock()
	defer pv.mu.RUnlock()

	return pv.verifyCount, pv.cacheHits, pv.cacheMisses
}

// Groth16Proof represents a Groth16 proof structure
type Groth16Proof struct {
	Ar  bn254.G1Affine // Proof component A
	Bs  bn254.G2Affine // Proof component B
	Krs bn254.G1Affine // Proof component C
}

// Groth16VerifyingKey represents a Groth16 verifying key
type Groth16VerifyingKey struct {
	Alpha bn254.G1Affine   // Alpha in G1
	Beta  bn254.G2Affine   // Beta in G2
	Gamma bn254.G2Affine   // Gamma in G2
	Delta bn254.G2Affine   // Delta in G2
	K     []bn254.G1Affine // K[i] for public inputs
}

// verifyGroth16WithGnark performs actual Groth16 verification using pairing operations
func (pv *ProofVerifier) verifyGroth16WithGnark(proof *ZKProof, vkBytes []byte) error {
	// Deserialize verifying key
	vk, err := deserializeVerifyingKey(vkBytes)
	if err != nil {
		return fmt.Errorf("failed to deserialize verifying key: %w", err)
	}

	// Validate verifying key with subgroup checks (CRITICAL for trusted setup validation)
	if err := validateVerifyingKey(vk); err != nil {
		return fmt.Errorf("verifying key validation failed: %w", err)
	}

	// Deserialize proof
	grothProof, err := deserializeGroth16Proof(proof.ProofData)
	if err != nil {
		return fmt.Errorf("failed to deserialize proof: %w", err)
	}

	// Deserialize public witness (public inputs)
	witness := make([]fr.Element, 0, len(proof.PublicInputs))
	for _, inputBytes := range proof.PublicInputs {
		var elem fr.Element
		elem.SetBytes(inputBytes)
		witness = append(witness, elem)
	}

	// Perform pairing-based verification
	if err := verifyGroth16Pairing(grothProof, vk, witness); err != nil {
		return fmt.Errorf("pairing verification failed: %w", err)
	}

	return nil
}

// verifyGroth16Pairing performs the Groth16 pairing check
// Verifies: e(A, B) = e(alpha, beta) * e(sum(pubInput_i * K_i), gamma) * e(C, delta)
//
// K carries one point per public input plus the constant term K[0], so a key of
// n points speaks about exactly n-1 public inputs — that count is a property of
// the circuit the key was made for, not something a transaction chooses.
//
// The witness arrives with the transaction, so a peer picks its length, and the
// sum below runs over the witness against K. A length either side of what the
// key states judges the proof against a statement the key does not describe —
// past the end of K in one direction, and short of the trailing K points in the
// other — so it has to be exactly what the key says.
func verifyGroth16Pairing(proof *Groth16Proof, vk *Groth16VerifyingKey, witness []fr.Element) error {
	if want := len(vk.K) - 1; len(witness) != want {
		return fmt.Errorf("public inputs: %d supplied, the verifying key's circuit takes %d",
			len(witness), want)
	}

	// Public input linear combination: K[0] + sum(witness_i * K[i+1]).
	// Every node has to reach the same point from the same key and witness,
	// so the sum stays on the CPU rather than being handed to an accelerator
	// whose answer is never compared against it.
	var publicInputLC bn254.G1Affine
	publicInputLC.Set(&vk.K[0])
	var scalar big.Int
	for i := range witness {
		var term bn254.G1Affine
		term.ScalarMultiplication(&vk.K[i+1], witness[i].BigInt(&scalar))
		publicInputLC.Add(&publicInputLC, &term)
	}

	// Pairing check: e(A, B) == e(alpha, beta) * e(publicInputLC, gamma) * e(C, delta)
	leftSide, err := bn254.Pair([]bn254.G1Affine{proof.Ar}, []bn254.G2Affine{proof.Bs})
	if err != nil {
		return fmt.Errorf("pairing A*B failed: %w", err)
	}

	alphaBeta, err := bn254.Pair([]bn254.G1Affine{vk.Alpha}, []bn254.G2Affine{vk.Beta})
	if err != nil {
		return fmt.Errorf("pairing alpha*beta failed: %w", err)
	}

	pubGamma, err := bn254.Pair([]bn254.G1Affine{publicInputLC}, []bn254.G2Affine{vk.Gamma})
	if err != nil {
		return fmt.Errorf("pairing pubInput*gamma failed: %w", err)
	}

	cDelta, err := bn254.Pair([]bn254.G1Affine{proof.Krs}, []bn254.G2Affine{vk.Delta})
	if err != nil {
		return fmt.Errorf("pairing C*delta failed: %w", err)
	}

	var rightSide bn254.GT
	rightSide.Set(&alphaBeta)
	rightSide.Mul(&rightSide, &pubGamma)
	rightSide.Mul(&rightSide, &cDelta)

	if !leftSide.Equal(&rightSide) {
		return errors.New("pairing check failed: proof is invalid")
	}

	return nil
}

// A point this verifier reads has to be in the prime-order subgroup and must
// not be the point at infinity. Both halves are load-bearing: gnark encodes
// infinity as all-zero bytes and reports it as in-subgroup, and a pairing drops
// any term whose argument is infinity, so an element at infinity would remove
// itself from the equation and leave a weaker check than the one written down.
// checkG1 and checkG2 are the one place that decides what a usable point is,
// and every decoder below routes through them.
var (
	errOffSubgroup = errors.New("point not in the prime-order subgroup")
	errAtInfinity  = errors.New("point at infinity")
)

func checkG1(p *bn254.G1Affine) error {
	if !p.IsInSubGroup() {
		return errOffSubgroup
	}
	if p.IsInfinity() {
		return errAtInfinity
	}
	return nil
}

func checkG2(p *bn254.G2Affine) error {
	if !p.IsInSubGroup() {
		return errOffSubgroup
	}
	if p.IsInfinity() {
		return errAtInfinity
	}
	return nil
}

// validateVerifyingKey checks every point of a Groth16 verifying key.
// A trusted setup never produces infinity for alpha, beta, gamma, delta or K,
// so a key that carries one is not a setup output and its pairing equation
// would collapse to something weaker.
func validateVerifyingKey(vk *Groth16VerifyingKey) error {
	if err := checkG1(&vk.Alpha); err != nil {
		return fmt.Errorf("Alpha: %w", err)
	}
	if err := checkG2(&vk.Beta); err != nil {
		return fmt.Errorf("Beta: %w", err)
	}
	if err := checkG2(&vk.Gamma); err != nil {
		return fmt.Errorf("Gamma: %w", err)
	}
	if err := checkG2(&vk.Delta); err != nil {
		return fmt.Errorf("Delta: %w", err)
	}
	for i := range vk.K {
		if err := checkG1(&vk.K[i]); err != nil {
			return fmt.Errorf("K[%d]: %w", i, err)
		}
	}

	return nil
}

// deserializeGroth16Proof deserializes a Groth16 proof from bytes.
// A proof it returns has already passed the point checks, so callers never
// repeat them.
func deserializeGroth16Proof(data []byte) (*Groth16Proof, error) {
	// Expected format: Ar (64 bytes) | Bs (128 bytes) | Krs (64 bytes) = 256 bytes
	if len(data) < 256 {
		return nil, errors.New("proof data too short")
	}

	proof := &Groth16Proof{}
	offset := 0

	// Deserialize Ar (G1 point, 64 bytes compressed)
	if err := proof.Ar.Unmarshal(data[offset : offset+64]); err != nil {
		return nil, fmt.Errorf("failed to unmarshal Ar: %w", err)
	}
	if err := checkG1(&proof.Ar); err != nil {
		return nil, fmt.Errorf("Ar: %w", err)
	}
	offset += 64

	// Deserialize Bs (G2 point, 128 bytes compressed)
	if err := proof.Bs.Unmarshal(data[offset : offset+128]); err != nil {
		return nil, fmt.Errorf("failed to unmarshal Bs: %w", err)
	}
	if err := checkG2(&proof.Bs); err != nil {
		return nil, fmt.Errorf("Bs: %w", err)
	}
	offset += 128

	// Deserialize Krs (G1 point, 64 bytes compressed)
	if err := proof.Krs.Unmarshal(data[offset : offset+64]); err != nil {
		return nil, fmt.Errorf("failed to unmarshal Krs: %w", err)
	}
	if err := checkG1(&proof.Krs); err != nil {
		return nil, fmt.Errorf("Krs: %w", err)
	}

	return proof, nil
}

// deserializeVerifyingKey deserializes a Groth16 verifying key from bytes
func deserializeVerifyingKey(data []byte) (*Groth16VerifyingKey, error) {
	// Format: Alpha (64) | Beta (128) | Gamma (128) | Delta (128) | numK (4) | K[...] (64*numK)
	minSize := 64 + 128 + 128 + 128 + 4
	if len(data) < minSize {
		return nil, errors.New("verifying key data too short")
	}

	vk := &Groth16VerifyingKey{}
	offset := 0

	// Alpha (G1)
	if err := vk.Alpha.Unmarshal(data[offset : offset+64]); err != nil {
		return nil, fmt.Errorf("failed to unmarshal Alpha: %w", err)
	}
	offset += 64

	// Beta (G2)
	if err := vk.Beta.Unmarshal(data[offset : offset+128]); err != nil {
		return nil, fmt.Errorf("failed to unmarshal Beta: %w", err)
	}
	offset += 128

	// Gamma (G2)
	if err := vk.Gamma.Unmarshal(data[offset : offset+128]); err != nil {
		return nil, fmt.Errorf("failed to unmarshal Gamma: %w", err)
	}
	offset += 128

	// Delta (G2)
	if err := vk.Delta.Unmarshal(data[offset : offset+128]); err != nil {
		return nil, fmt.Errorf("failed to unmarshal Delta: %w", err)
	}
	offset += 128

	// Number of K points
	numK := binary.BigEndian.Uint32(data[offset : offset+4])
	offset += 4

	if len(data) < offset+int(numK)*64 {
		return nil, errors.New("insufficient data for K points")
	}

	// K points (G1)
	vk.K = make([]bn254.G1Affine, numK)
	for i := uint32(0); i < numK; i++ {
		if err := vk.K[i].Unmarshal(data[offset : offset+64]); err != nil {
			return nil, fmt.Errorf("failed to unmarshal K[%d]: %w", i, err)
		}
		offset += 64
	}

	return vk, nil
}

// errPLONKIncomplete says why a PLONK proof is never accepted here.
//
// What stood in this place decoded the proof and the key — exact sizes,
// subgroup checks, points at infinity, some 200 lines of it — and then handed
// the result to a pairing that refused every proof unconditionally, because the
// verification equation
//
//	e(Wz + u·Wzw, [α]₂) = e(z·Wz + u·zω·Wzw, [1]₂)
//
// was never written. Decoding a proof carefully in order to refuse it is not a
// verifier; it is a shape that reads like one. The refusal is the whole of what
// this path does, so it is the whole of what this path says.
//
// A strict-PQ chain — the default here — refuses classical systems before
// reaching this at all. On a permissive chain, PLONK is refused here.
var errPLONKIncomplete = errors.New(
	"plonk: the verification equation is not implemented — failing closed rather than " +
		"accepting a proof on its shape")