fix: change the program ids to be a mapping program id =>verifier type

Author: maximopalopoli

contracts/src/core/AlignedProofAggregationService.sol

@@ -37,13 +37,16 @@ contract AlignedProofAggregationService is
     /// if the sp1 verifier address is set to this address, then we skip verification
     address public constant VERIFIER_MOCK_ADDRESS = address(0xFF);
 
-    /// The unique identifier (image ID) of the RISC Zero aggregator program.
-    /// This ensures that only proofs generated by a trusted Risc0 program can be verified.
-    bytes32 public risc0AggregatorProgramImageId;
-
-    /// The verification key hash for the SP1 aggregator program.
-    /// This ensures that only proofs generated by a trusted SP1 program can be verified.
-    bytes32 public sp1AggregatorProgramVKHash;
+    /// @notice A mapping to track valid program IDs (image IDs for RISC Zero or vk hashes for SP1)
+    /// to their corresponding verifier types. These program IDs are used to ensure that proofs are
+    /// verified against known and registered programs.
+    mapping(bytes32 => uint8) public programIds;
+
+    enum VerifierType {
+        INVALID, // This to prevent default value being considered valid (when mapping returns 0 for non-existing keys)
+        SP1,
+        RISC0
+    }
 
     constructor() {
         _disableInitializers();
@@ -54,45 +57,50 @@ contract AlignedProofAggregationService is
         address _alignedAggregatorAddress,
         address _sp1VerifierAddress,
         address _risc0VerifierAddress,
-        bytes32 _risc0AggregatorProgramImageId,
-        bytes32 _sp1AggregatorProgramVKHash
+        bytes32[] memory _programIds,
+        uint8[] memory _verifierTypes
     ) public initializer {
         __Ownable_init();
         __UUPSUpgradeable_init();
         _transferOwnership(newOwner);
         alignedAggregatorAddress = _alignedAggregatorAddress;
         sp1VerifierAddress = _sp1VerifierAddress;
         risc0VerifierAddress = _risc0VerifierAddress;
-        risc0AggregatorProgramImageId = _risc0AggregatorProgramImageId;
-        sp1AggregatorProgramVKHash = _sp1AggregatorProgramVKHash;
+        for (uint256 i = 0; i < _programIds.length; i++) {
+            programIds[_programIds[i]] = _verifierTypes[i];
+        }
     }
 
-    function verifySP1(bytes32 blobVersionedHash, bytes calldata sp1PublicValues, bytes calldata sp1ProofBytes)
+    function verifySP1(bytes32 blobVersionedHash, bytes calldata sp1PublicValues, bytes calldata sp1ProofBytes, bytes32 programId)
         public
         onlyAlignedAggregator
     {
         (bytes32 merkleRoot) = abi.decode(sp1PublicValues, (bytes32));
 
+        require(programIds[programId] == uint8(VerifierType.SP1), "Invalid program ID");
+
         // In dev mode, poofs are mocked, so we skip the verification part
         if (_isSP1VerificationEnabled()) {
-            ISP1Verifier(sp1VerifierAddress).verifyProof(sp1AggregatorProgramVKHash, sp1PublicValues, sp1ProofBytes);
+            ISP1Verifier(sp1VerifierAddress).verifyProof(programId, sp1PublicValues, sp1ProofBytes);
         }
 
         aggregatedProofs[merkleRoot] = true;
         emit AggregatedProofVerified(merkleRoot, blobVersionedHash);
     }
 
-    function verifyRisc0(bytes32 blobVersionedHash, bytes calldata risc0ReceiptSeal, bytes calldata risc0JournalBytes)
+    function verifyRisc0(bytes32 blobVersionedHash, bytes calldata risc0ReceiptSeal, bytes calldata risc0JournalBytes, bytes32 programId)
         public
         onlyAlignedAggregator
     {
         (bytes32 merkleRoot) = abi.decode(risc0JournalBytes, (bytes32));
 
+        require(programIds[programId] == uint8(VerifierType.RISC0), "Invalid program ID");
+
         // In dev mode, poofs are mocked, so we skip the verification part
         if (_isRisc0VerificationEnabled()) {
             bytes32 risc0JournalDigest = sha256(risc0JournalBytes);
             IRiscZeroVerifier(risc0VerifierAddress).verify(
-                risc0ReceiptSeal, risc0AggregatorProgramImageId, risc0JournalDigest
+                risc0ReceiptSeal, programId, risc0JournalDigest
             );
         }
 
@@ -113,13 +121,17 @@ contract AlignedProofAggregationService is
     /// @param merklePath The Merkle proof (sibling hashes) needed to reconstruct the Merkle root.
     /// @param programId The identifier for the ZK program (image_id in RISC0 or vk hash in SP1).
     /// @param publicInputs The public inputs bytes of the proof.
+    /// @param verifierType The type of verifier (SP1 or RISC0).
     ///
     /// @return bool Returns true if the computed Merkle root is a recognized valid aggregated proof.
-    function verifyProofInclusion(bytes32[] calldata merklePath, bytes32 programId, bytes calldata publicInputs)
+    function verifyProofInclusion(bytes32[] calldata merklePath, bytes32 programId, bytes calldata publicInputs, VerifierType verifierType)
         public
         view
         returns (bool)
     {
+        require(programIds[programId] == uint8(verifierType), "Program ID not registered for the received verifier type");
+
+        // TODO: consider adding the verifierType to the proof commitment to avoid potential collisions
         bytes32 proofCommitment = keccak256(abi.encodePacked(programId, publicInputs));
         bytes32 merkleRoot = MerkleProof.processProofCalldata(merklePath, proofCommitment);
         return aggregatedProofs[merkleRoot];
@@ -153,24 +165,19 @@ contract AlignedProofAggregationService is
         emit Risc0VerifierAddressUpdated(_risc0VerifierAddress);
     }
 
-    /// @notice Sets the image id of the Risc0 program
-    /// @param _risc0AggregatorProgramImageId The new imageid for the Risc0 aggregator program
-    function setRisc0AggregatorProgramImageId(bytes32 _risc0AggregatorProgramImageId) external onlyOwner {
-        risc0AggregatorProgramImageId = _risc0AggregatorProgramImageId;
-        emit Risc0AggregatorProgramImageIdUpdated(_risc0AggregatorProgramImageId);
-    }
-
     /// @notice Sets the address of the SP1 verifier contract
     /// @param _sp1VerifierAddress The new address for the SP1 verifier contract
     function setSP1VerifierAddress(address _sp1VerifierAddress) external onlyOwner {
         sp1VerifierAddress = _sp1VerifierAddress;
         emit SP1VerifierAddressUpdated(_sp1VerifierAddress);
     }
 
-    /// @notice Sets the vk hash of the sp1 program
-    /// @param _sp1AggregatorProgramVKHash The new vk hash for the sp1 aggregator program
-    function setSP1AggregatorProgramVKHash(bytes32 _sp1AggregatorProgramVKHash) external onlyOwner {
-        sp1AggregatorProgramVKHash = _sp1AggregatorProgramVKHash;
-        emit SP1AggregatorProgramVKHashUpdated(_sp1AggregatorProgramVKHash);
+    function addProgramId(bytes32 programId, VerifierType verifierType) external onlyOwner {
+        programIds[programId] = uint8(verifierType);
+    }
+
+    function deleteProgramId(bytes32 programId) external onlyOwner {
+        // Check if we should verify its existence before deleting
+        delete programIds[programId];
     }
 }