Merge pull request #40 from flashbots/N-03/custom-errors-with-require

N-03 use custom errors everywhere

Author: Melvillian

src/FlashtestationRegistry.sol

@@ -141,34 +141,32 @@ contract FlashtestationRegistry is
         limitBytesSize(extendedRegistrationData)
     {
         (bool success, bytes memory output) = attestationContract.verifyAndAttestOnChain{value: msg.value}(rawQuote);
-        if (!success) {
-            revert InvalidQuote(output);
-        }
+        require(success, InvalidQuote(output));
 
         // now we know the quote is valid, we can safely parse the output into the TDX report body,
         // from which we'll extract the data we need to register the TEE
         TD10ReportBody memory td10ReportBody = QuoteParser.parseV4VerifierOutput(output);
 
         // Binding the tee address and extended report data to the quote
-        if (td10ReportBody.reportData.length < TD_REPORTDATA_LENGTH) {
-            revert InvalidReportDataLength(td10ReportBody.reportData.length);
-        }
+        require(
+            td10ReportBody.reportData.length >= TD_REPORTDATA_LENGTH,
+            InvalidReportDataLength(td10ReportBody.reportData.length)
+        );
 
         (address teeAddress, bytes32 extendedDataReportHash) = QuoteParser.parseReportData(td10ReportBody.reportData);
 
         // Ensure that the caller is the TEE-controlled address, otherwise we have no guarantees that
         // the TEE-controlled address is the one that is registering the TEE
-        if (signer != teeAddress) {
-            revert SignerMustMatchTEEAddress(signer, teeAddress);
-        }
+        require(signer == teeAddress, SignerMustMatchTEEAddress(signer, teeAddress));
 
         // Verify that the extended registration data matches the hash in the TDX report data
         // This is to ensure that the values in extendedRegistrationData are the same as the values
         // in the TDX report data, which cannot be forged by the TEE-controlled address
         bytes32 extendedRegistrationDataHash = keccak256(extendedRegistrationData);
-        if (extendedRegistrationDataHash != extendedDataReportHash) {
-            revert InvalidRegistrationDataHash(extendedDataReportHash, extendedRegistrationDataHash);
-        }
+        require(
+            extendedRegistrationDataHash == extendedDataReportHash,
+            InvalidRegistrationDataHash(extendedDataReportHash, extendedRegistrationDataHash)
+        );
 
         bytes32 newQuoteHash = keccak256(rawQuote);
         bool previouslyRegistered = checkPreviousRegistration(teeAddress, newQuoteHash);
@@ -203,9 +201,7 @@ contract FlashtestationRegistry is
      */
     function checkPreviousRegistration(address teeAddress, bytes32 newQuoteHash) internal view returns (bool) {
         bytes32 existingQuoteHash = registeredTEEs[teeAddress].quoteHash;
-        if (newQuoteHash == existingQuoteHash) {
-            revert TEEServiceAlreadyRegistered(teeAddress);
-        }
+        require(newQuoteHash != existingQuoteHash, TEEServiceAlreadyRegistered(teeAddress));
 
         // if the TEE is already registered, but we're using a different quote,
         // return true to signal that the TEE is already registered but is updating its quote
@@ -234,28 +230,17 @@ contract FlashtestationRegistry is
         // if the TEE-controlled address is not registered with the FlashtestationRegistry,
         // it doesn't make sense to invalidate the attestation
         RegisteredTEE memory registeredTEE = registeredTEEs[teeAddress];
-        if (registeredTEE.rawQuote.length == 0) {
-            revert TEEServiceNotRegistered(teeAddress);
-        }
-
-        if (!registeredTEE.isValid) {
-            revert TEEServiceAlreadyInvalid(teeAddress);
-        }
+        require(registeredTEE.rawQuote.length > 0, TEEServiceNotRegistered(teeAddress));
+        require(registeredTEE.isValid, TEEServiceAlreadyInvalid(teeAddress));
 
         // now we check the attestation, and invalidate the TEE if it's no longer valid.
         // This will only happen if the DCAP Endorsements associated with the TEE's quote
         // have been updated
         (bool success,) = attestationContract.verifyAndAttestOnChain{value: msg.value}(registeredTEE.rawQuote);
-        if (success) {
-            // if the attestation is still valid, then this function call is a no-op except for
-            // wasting the caller's gas. So we revert here to signal that the TEE is still valid.
-            // Offchain users who want to monitor for potential invalid TEEs can do so by calling
-            // this function and checking for the `TEEIsStillValid` error
-            revert TEEIsStillValid(teeAddress);
-        } else {
-            registeredTEEs[teeAddress].isValid = false;
-            emit TEEServiceInvalidated(teeAddress);
-        }
+        require(!success, TEEIsStillValid(teeAddress));
+
+        registeredTEEs[teeAddress].isValid = false;
+        emit TEEServiceInvalidated(teeAddress);
     }
 
     /// @inheritdoc IFlashtestationRegistry

src/utils/QuoteParser.sol

@@ -107,9 +107,7 @@ library QuoteParser {
      */
     function checkTEEVersion(bytes memory rawReportBody) internal pure {
         uint16 version = uint16(bytes2((rawReportBody.substring(0, 2)))); // uint16 is 2 bytes
-        if (version != ACCEPTED_TDX_VERSION) {
-            revert InvalidTEEVersion(version);
-        }
+        require(version == ACCEPTED_TDX_VERSION, InvalidTEEVersion(version));
     }
 
     /**
@@ -119,8 +117,6 @@ library QuoteParser {
      */
     function checkTEEType(bytes memory rawReportBody) internal pure {
         bytes4 teeType = bytes4(rawReportBody.substring(2, 4)); // 4 bytes
-        if (teeType != TDX_TEE) {
-            revert InvalidTEEType(teeType);
-        }
+        require(teeType == TDX_TEE, InvalidTEEType(teeType));
     }
 }

test/mocks/MockAutomataDcapAttestationFee.sol

@@ -25,9 +25,7 @@ contract MockAutomataDcapAttestationFee {
     error InsufficientFee(uint256 required, uint256 provided);
 
     function verifyAndAttestOnChain(bytes calldata rawQuote) external payable returns (bool, bytes memory) {
-        if (msg.value < baseFee) {
-            revert InsufficientFee(baseFee, msg.value);
-        }
+        require(msg.value >= baseFee, InsufficientFee(baseFee, msg.value));
 
         QuoteResult memory result = quoteResults[rawQuote];
         return (result.success, result.output);