Fixing AVM origin tag issues

Author: Rumata888

barretenberg/cpp/src/barretenberg/vm2/avm_api.cpp

@@ -26,7 +26,8 @@ std::pair<AvmAPI::AvmProof, AvmAPI::AvmVerificationKey> AvmAPI::prove(const AvmA
     // Prove.
     info("Proving...");
     AvmProvingHelper proving_helper;
-    auto [proof, vk] = AVM_TRACK_TIME_V("proving/all", proving_helper.prove(std::move(trace)));
+    auto [proof, vk] =
+        AVM_TRACK_TIME_V("proving/all", proving_helper.prove(std::move(trace), inputs.publicInputs.to_columns()));
 
     info("Done!");
     return { std::move(proof), std::move(vk) };

barretenberg/cpp/src/barretenberg/vm2/constraining/prover.cpp

@@ -46,6 +46,20 @@ void AvmProver::execute_preamble_round()
     info("AVM vk hash in prover: ", vk_hash);
 }
 
+/**
+ * @brief Add public inputs to transcript
+ *
+ */
+void AvmProver::execute_public_inputs_round(const std::vector<std::vector<FF>>& public_inputs_cols)
+{
+    for (size_t i = 0; i < public_inputs_cols.size(); ++i) {
+        for (size_t j = 0; j < AVM_PUBLIC_INPUTS_COLUMNS_MAX_LENGTH; ++j) {
+            // The public inputs are added to the hash buffer, but do not increase the size of the proof
+            transcript->add_to_hash_buffer("public_input_" + std::to_string(i) + "_" + std::to_string(j),
+                                           j < public_inputs_cols[i].size() ? public_inputs_cols[i][j] : FF(0));
+        }
+    }
+}
 /**
  * @brief Compute commitments to all of the witness wires (apart from the logderivative inverse wires)
  *
@@ -141,11 +155,14 @@ HonkProof AvmProver::export_proof()
     return transcript->export_proof();
 }
 
-HonkProof AvmProver::construct_proof()
+HonkProof AvmProver::construct_proof(const std::vector<std::vector<FF>>& public_inputs_cols)
 {
     // Add circuit size public input size and public inputs to transcript.
     execute_preamble_round();
 
+    // Add public inputs to transcript.
+    AVM_TRACK_TIME("prove/public_inputs_round", execute_public_inputs_round(public_inputs_cols));
+
     // Compute wire commitments.
     AVM_TRACK_TIME("prove/wire_commitments_round", execute_wire_commitments_round());
 

barretenberg/cpp/src/barretenberg/vm2/constraining/prover.hpp

@@ -31,14 +31,15 @@ class AvmProver {
     // Note: all the following methods are virtual to allow Avm2 to tweak the behaviour.
     // We can remove this once the transition is done.
     virtual void execute_preamble_round();
+    virtual void execute_public_inputs_round(const std::vector<std::vector<FF>>& public_inputs_cols);
     virtual void execute_wire_commitments_round();
     virtual void execute_log_derivative_inverse_round();
     virtual void execute_log_derivative_inverse_commitments_round();
     virtual void execute_relation_check_rounds();
     virtual void execute_pcs_rounds();
 
     virtual HonkProof export_proof();
-    virtual HonkProof construct_proof();
+    virtual HonkProof construct_proof(const std::vector<std::vector<FF>>& public_inputs_cols);
 
     std::shared_ptr<Transcript> transcript = std::make_shared<Transcript>();
 

barretenberg/cpp/src/barretenberg/vm2/constraining/recursion/recursive_verifier.cpp

@@ -83,21 +83,43 @@ AvmRecursiveVerifier::PairingPoints AvmRecursiveVerifier::verify_proof(
     // TODO(#14234)[Unconditional PIs validation]: Remove the next 3 lines
     StdlibProof stdlib_proof = stdlib_proof_with_pi_flag;
     bool_t<Builder> pi_validation = !bool_t<Builder>(stdlib_proof.at(0));
+    // TODO(https://github.com/AztecProtocol/aztec-packages/issues/16716) Origin Tag security mechanism is screaming
+    // that there is a free witness affecting proof verificaton. Because it is and this bool allows completely disabling
+    // public input logic. So this has to be removed in the future.
+    pi_validation.unset_free_witness_tag();
     stdlib_proof.erase(stdlib_proof.begin());
 
     if (public_inputs.size() != AVM_NUM_PUBLIC_INPUT_COLUMNS) {
         throw_or_abort("AvmRecursiveVerifier::verify_proof: public inputs size mismatch");
     }
+    for (const auto& public_input : public_inputs) {
+        if (public_input.size() != AVM_PUBLIC_INPUTS_COLUMNS_MAX_LENGTH) {
+            throw_or_abort("AvmRecursiveVerifier::verify_proof: public input size mismatch");
+        }
+    }
 
     transcript->load_proof(stdlib_proof);
 
     // TODO(#15892): Fiat-Shamir the vk hash by uncommenting the add_to_hash_buffer.
     // transcript->add_to_hash_buffer("avm_vk_hash", vk_hash);
+    // TODO(https://github.com/AztecProtocol/aztec-packages/issues/16716) For now we are unsetting the free witness tags
+    // to stop triggering the Origin Tag security mechanism, but the problem is that the VK is not hashed.
+    for (auto& comm : key->get_all()) {
+        comm.unset_free_witness_tag();
+    }
+
     info("AVM vk hash in recursive verifier: ", vk_hash);
 
     RelationParams relation_parameters;
     VerifierCommitments commitments{ key };
 
+    // Add public inputs to transcript
+    for (size_t i = 0; i < AVM_NUM_PUBLIC_INPUT_COLUMNS; i++) {
+        for (size_t j = 0; j < public_inputs[i].size(); j++) {
+            transcript->add_to_hash_buffer("public_input_" + std::to_string(i) + "_" + std::to_string(j),
+                                           public_inputs[i][j]);
+        }
+    }
     // Get commitments to VM wires
     for (auto [comm, label] : zip_view(commitments.get_wires(), commitments.get_wires_labels())) {
         comm = transcript->template receive_from_prover<Commitment>(label);

barretenberg/cpp/src/barretenberg/vm2/constraining/recursion/recursive_verifier.test.cpp

@@ -41,12 +41,13 @@ class AvmRecursiveTests : public ::testing::Test {
     {
         auto [trace, public_inputs] = testing::get_minimal_trace_with_pi();
 
+        const auto public_inputs_cols = public_inputs.to_columns();
+
         InnerProver prover;
-        const auto [proof, vk_data] = prover.prove(std::move(trace));
+        const auto [proof, vk_data] = prover.prove(std::move(trace), public_inputs_cols);
         const auto verification_key = InnerProver::create_verification_key(vk_data);
         InnerVerifier verifier(verification_key);
 
-        const auto public_inputs_cols = public_inputs.to_columns();
         const bool verified = verifier.verify_proof(proof, public_inputs_cols);
 
         // Should be in principle ASSERT_TRUE, but compiler does not like it.

barretenberg/cpp/src/barretenberg/vm2/constraining/verifier.cpp

@@ -60,6 +60,22 @@ bool AvmVerifier::verify_proof(const HonkProof& proof, const std::vector<std::ve
     // transcript->add_to_hash_buffer("avm_vk_hash", vk_hash);
     info("AVM vk hash in verifier: ", vk_hash);
 
+    // Check public inputs size.
+    if (public_inputs.size() != AVM_NUM_PUBLIC_INPUT_COLUMNS) {
+        vinfo("Public inputs size mismatch");
+        return false;
+    }
+    // Public inputs from proof
+    for (size_t i = 0; i < AVM_NUM_PUBLIC_INPUT_COLUMNS; i++) {
+        if (public_inputs[i].size() != AVM_PUBLIC_INPUTS_COLUMNS_MAX_LENGTH) {
+            vinfo("Public input size mismatch");
+            return false;
+        }
+        for (size_t j = 0; j < public_inputs[i].size(); j++) {
+            transcript->add_to_hash_buffer("public_input_" + std::to_string(i) + "_" + std::to_string(j),
+                                           public_inputs[i][j]);
+        }
+    }
     VerifierCommitments commitments{ key };
     // Get commitments to VM wires
     for (auto [comm, label] : zip_view(commitments.get_wires(), commitments.get_wires_labels())) {
@@ -95,11 +111,6 @@ bool AvmVerifier::verify_proof(const HonkProof& proof, const std::vector<std::ve
         return false;
     }
 
-    if (public_inputs.size() != AVM_NUM_PUBLIC_INPUT_COLUMNS) {
-        vinfo("Public inputs size mismatch");
-        return false;
-    }
-
     std::array<FF, AVM_NUM_PUBLIC_INPUT_COLUMNS> claimed_evaluations = {
         output.claimed_evaluations.public_inputs_cols_0_,
         output.claimed_evaluations.public_inputs_cols_1_,

barretenberg/cpp/src/barretenberg/vm2/constraining/verifier.test.cpp

@@ -29,10 +29,10 @@ class AvmVerifierTests : public ::testing::Test {
         auto [trace, public_inputs] = testing::get_minimal_trace_with_pi();
 
         Prover prover;
-        const auto [proof, vk_data] = prover.prove(std::move(trace));
+        auto public_inputs_cols = public_inputs.to_columns();
+        const auto [proof, vk_data] = prover.prove(std::move(trace), public_inputs_cols);
         const auto verification_key = prover.create_verification_key(vk_data);
 
-        auto public_inputs_cols = public_inputs.to_columns();
         return { proof, verification_key, public_inputs_cols };
     }
 };

barretenberg/cpp/src/barretenberg/vm2/proving_helper.cpp

@@ -55,7 +55,8 @@ std::shared_ptr<AvmVerifier::VerificationKey> AvmProvingHelper::create_verificat
     return std::make_shared<VerificationKey>(precomputed_cmts);
 }
 
-std::pair<AvmProvingHelper::Proof, AvmProvingHelper::VkData> AvmProvingHelper::prove(tracegen::TraceContainer&& trace)
+std::pair<AvmProvingHelper::Proof, AvmProvingHelper::VkData> AvmProvingHelper::prove(
+    tracegen::TraceContainer&& trace, const std::vector<std::vector<FF>>& public_inputs_cols)
 {
     auto polynomials = AVM_TRACK_TIME_V("proving/prove:compute_polynomials", constraining::compute_polynomials(trace));
     auto proving_key = AVM_TRACK_TIME_V("proving/prove:proving_key", create_proving_key(polynomials));
@@ -65,7 +66,7 @@ std::pair<AvmProvingHelper::Proof, AvmProvingHelper::VkData> AvmProvingHelper::p
     auto prover = AVM_TRACK_TIME_V("proving/prove:construct_prover",
                                    AvmProver(proving_key, verification_key, proving_key->commitment_key));
 
-    auto proof = AVM_TRACK_TIME_V("proving/construct_proof", prover.construct_proof());
+    auto proof = AVM_TRACK_TIME_V("proving/construct_proof", prover.construct_proof(public_inputs_cols));
     auto serialized_vk = to_buffer(verification_key->to_field_elements());
 
     return { std::move(proof), std::move(serialized_vk) };

barretenberg/cpp/src/barretenberg/vm2/proving_helper.hpp

@@ -15,7 +15,8 @@ class AvmProvingHelper {
     using VkData = std::vector<uint8_t>;
 
     static std::shared_ptr<AvmVerifier::VerificationKey> create_verification_key(const VkData& vk_data);
-    std::pair<Proof, VkData> prove(tracegen::TraceContainer&& trace);
+    std::pair<Proof, VkData> prove(tracegen::TraceContainer&& trace,
+                                   const std::vector<std::vector<FF>>& public_inputs_cols);
     bool check_circuit(tracegen::TraceContainer&& trace);
     bool verify(const Proof& proof, const PublicInputs& pi, const VkData& vk_data);
 };