From 49b7c05eea2890417890fd4d545534bd92c826c3 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 12 Nov 2026 04:45:24 -0800 Subject: [PATCH] refactor(kms-auth)!: slim DstackKms bytecode DstackKms was 24,041 bytes against the 24,576-byte EIP-170 limit, leaving little room for new features. This takes it to 22,621 without splitting the contract: - Revert with custom errors instead of revert strings. Authorization rejection reasons, which the auth backends forward, are unchanged. - Drop the single-epoch addKmsAggregatedMrAtEpoch, addKmsDeviceAtEpoch, addOsImageHashAtEpoch and their remove counterparts; *AtEpochs(e, e, v) writes one root epoch. Manage.s.sol already used the range forms. - Pass _emitPolicy the policy hash, folded at compile time, instead of hashing a string at run time. Event values are unchanged. BREAKING CHANGE: DstackKms setters revert with custom errors. The revert strings of registerApp, setAppImplementation and deployAndRegisterApp shipped in v0.6.0; the others and the single-epoch setters are unreleased. Signed-off-by: Kevin Wang --- CHANGELOG.md | 19 +- docs/app-key-epochs.md | 7 +- docs/formal/key-rotation/EpochPolicy.tla | 4 +- docs/tutorials/gateway-build-configuration.md | 2 +- .../troubleshooting-gateway-deployment.md | 4 +- dstack/kms/auth-eth/README.md | 5 +- dstack/kms/auth-eth/contracts/DstackKms.sol | 126 +++++------ dstack/kms/auth-eth/docs/specification.md | 8 +- dstack/kms/auth-eth/index.test.ts | 2 +- .../auth-eth/test/DstackKms.symbolic.t.sol | 10 +- dstack/kms/auth-eth/test/DstackKms.t.sol | 210 +++++++++--------- .../test/DstackKmsEpoch.symbolic.t.sol | 30 +-- dstack/kms/auth-eth/test/EventAudit.t.sol | 10 +- .../kms/auth-eth/test/LegacyAppCompat.t.sol | 4 +- .../auth-eth/test/UpgradesWithPlugin.t.sol | 14 +- 15 files changed, 230 insertions(+), 225 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 71b996488..2ae84ca37 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -82,8 +82,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - DstackKms, DstackApp: material approvals are per epoch. A request is decided on the row of the epoch it selects, so a value is revoked only once every row allowing it is cleared. Material setters name their rows: - `*AtEpoch(epoch, value)` writes one epoch, `*AtEpochs(from, to, value)` a - range; `remove*AtEpochs(0, latest root epoch, value)` in `DstackKms` and + `*AtEpochs(from, to, value)` writes a range, and in `DstackApp` + `*AtEpoch(epoch, value)` one epoch (in `DstackKms`, pass the epoch twice); + `remove*AtEpochs(0, latest root epoch, value)` in `DstackKms` and `remove*AtEpochs(0, currentAppEpoch, value)` in `DstackApp` revoke a value everywhere. `DstackApp` reports its setters as the new ERC-165 interface `IAppEpochManagement` (`0x8c0524c4`). The `Manage.s.sol` scripts require @@ -124,12 +125,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `ghcr.io/dstack-tee/dstack-kms-auth-eth-bun` remain available. Update image references and source paths in deployments. +- **Breaking:** DstackKms: reverts with custom errors such as + `UnknownRootEpoch()` instead of revert strings, to stay under the EIP-170 + size limit. This includes `registerApp`, `setAppImplementation` and + `deployAndRegisterApp`, whose v0.6.0 strings (`Invalid app ID`, `Invalid + implementation address`, `DstackApp implementation not set`, `Invalid owner + address`) become `InvalidAppId()`, `InvalidImplementationAddress()`, + `AppImplementationNotSet()` and `InvalidOwnerAddress()`; tooling that + matches the strings must decode the errors. Authorization rejection reasons + are unchanged. + ### Removed - DstackKms: the epoch-less `addOsImageHash`, `removeOsImageHash`, `addKmsAggregatedMr`, `removeKmsAggregatedMr`, `addKmsDevice` and - `removeKmsDevice`. Use the `*AtEpoch` or `*AtEpochs` forms; before any - rotation, `addOsImageHashAtEpoch(0, hash)` is the old call. + `removeKmsDevice`. Use the `*AtEpochs` forms; before any rotation, + `addOsImageHashAtEpochs(0, 0, hash)` is the old call. - **Breaking:** DstackApp: the epoch-less `addComposeHash`, `removeComposeHash`, `addDevice` and `removeDevice`, and with them `IAppAuthBasicManagement`: an upgraded app no longer reports interface ID diff --git a/docs/app-key-epochs.md b/docs/app-key-epochs.md index bbb759ecc..f0d4c7a9b 100644 --- a/docs/app-key-epochs.md +++ b/docs/app-key-epochs.md @@ -64,9 +64,10 @@ the old row after the handover. The old image never needs access to the new row. A request may select any epoch up to the current one and is decided on that epoch's row, so a value is revoked only once every row allowing it is cleared. -Neither contract has epoch-less material setters: `*AtEpoch(epoch, material)` -changes one row and `*AtEpochs(fromEpoch, toEpoch, material)` every row in the -range, so one transaction revokes a value from all rows, with +Neither contract has epoch-less material setters: +`*AtEpochs(fromEpoch, toEpoch, material)` changes every row in the range, and in +`DstackApp` `*AtEpoch(epoch, material)` one row (in `DstackKms`, pass the epoch +twice). One transaction revokes a value from all rows, with `remove*AtEpochs(0, latest, …)` in `DstackKms` and `remove*AtEpochs(0, currentAppEpoch, …)` in `DstackApp`. `DstackKms` writes registered root epochs (before any is registered, only epoch zero); `DstackApp` diff --git a/docs/formal/key-rotation/EpochPolicy.tla b/docs/formal/key-rotation/EpochPolicy.tla index e47a305a0..0a33a423e 100644 --- a/docs/formal/key-rotation/EpochPolicy.tla +++ b/docs/formal/key-rotation/EpochPolicy.tla @@ -7,7 +7,7 @@ (* on the selected rows. *) (* *) (* Root side (DstackKms): one row of allowed values (OS images; KMS MRs *) -(* and devices alike) per root epoch, written by *AtEpoch(s) over a range *) +(* and devices alike) per root epoch, written by *AtEpochs over a range *) (* of registered rows. A pending epoch can be popped and registered again, *) (* with the same or other keys. Rows are keyed by registration *) (* (RootEpoch.registration); KeyByChainHash = TRUE models the earlier *) @@ -77,7 +77,7 @@ Pop == /\ regs' = SubSeq(regs, 1, Latest) /\ UNCHANGED <> -\* add*AtEpochs / remove*AtEpochs(from, to, v); *AtEpoch is from = to. +\* add*AtEpochs / remove*AtEpochs(from, to, v); one row is from = to. Write(from, to, v, allowed) == /\ from <= to /\ to <= Latest /\ LET written == {<>: e \in from..to} diff --git a/docs/tutorials/gateway-build-configuration.md b/docs/tutorials/gateway-build-configuration.md index dd29fcc32..10dfc20c9 100644 --- a/docs/tutorials/gateway-build-configuration.md +++ b/docs/tutorials/gateway-build-configuration.md @@ -278,7 +278,7 @@ echo "OS image hash: 0x$OS_IMAGE_HASH" ``` ```bash -cast send "$KMS_CONTRACT_ADDR" "addOsImageHashAtEpoch(uint256,bytes32)" 0 "0x$OS_IMAGE_HASH" --rpc-url "$ETH_RPC_URL" --private-key "$PRIVATE_KEY" +cast send "$KMS_CONTRACT_ADDR" "addOsImageHashAtEpochs(uint256,uint256,bytes32)" 0 0 "0x$OS_IMAGE_HASH" --rpc-url "$ETH_RPC_URL" --private-key "$PRIVATE_KEY" ``` Verify it was added: diff --git a/docs/tutorials/troubleshooting-gateway-deployment.md b/docs/tutorials/troubleshooting-gateway-deployment.md index 053c42c57..73eb7f238 100644 --- a/docs/tutorials/troubleshooting-gateway-deployment.md +++ b/docs/tutorials/troubleshooting-gateway-deployment.md @@ -141,8 +141,8 @@ echo "OS image hash: 0x$OS_IMAGE_HASH" # Whitelist it on the KMS contract export KMS_CONTRACT_ADDR=$(cat ~/.dstack/secrets/kms-contract-address) cast send "$KMS_CONTRACT_ADDR" \ - "addOsImageHashAtEpoch(uint256,bytes32)" \ - 0 "0x$OS_IMAGE_HASH" \ + "addOsImageHashAtEpochs(uint256,uint256,bytes32)" \ + 0 0 "0x$OS_IMAGE_HASH" \ --rpc-url "https://ethereum-sepolia-rpc.publicnode.com" \ --private-key "$(cat ~/.dstack/secrets/sepolia-private-key)" ``` diff --git a/dstack/kms/auth-eth/README.md b/dstack/kms/auth-eth/README.md index 0e2536d21..517f49e47 100644 --- a/dstack/kms/auth-eth/README.md +++ b/dstack/kms/auth-eth/README.md @@ -332,8 +332,9 @@ See `script/README.md` for complete documentation of all available scripts. #### Epoch-scoped material approvals Material approvals use exact root/app epoch rows; neither older nor newer rows -are implicitly authorized. Material setters name their rows: `*AtEpoch` one -row, `*AtEpochs` a range; `DstackApp` reports them as `IAppEpochManagement`. +are implicitly authorized. Material setters name their rows: `*AtEpochs` a +range, and in `DstackApp` `*AtEpoch` one row; `DstackApp` reports them as +`IAppEpochManagement`. The management scripts take `ROOT_EPOCH`/`APP_EPOCH`, or `FROM_ROOT_EPOCH`/`TO_ROOT_EPOCH` and `FROM_APP_EPOCH`/`TO_APP_EPOCH`, and use the epoch-less setters of app policies predating app epochs. Original bool mappings and events remain epoch-zero views after diff --git a/dstack/kms/auth-eth/contracts/DstackKms.sol b/dstack/kms/auth-eth/contracts/DstackKms.sol index d35b95350..2eb798ddc 100644 --- a/dstack/kms/auth-eth/contracts/DstackKms.sol +++ b/dstack/kms/auth-eth/contracts/DstackKms.sol @@ -95,6 +95,23 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E event EpochPolicyChanged(bytes32 indexed policy, bytes32 indexed value, uint256 indexed epoch, bool enabled); + error InvalidEpochCaPubkey(); + error InvalidEpochK256Pubkey(); + error EmptyEpochEvidenceUrl(); + error Epoch0KeysDifferFromKmsInfo(); + error InvalidRootEpochActivation(); + error RootEpochChainHashMismatch(); + error NotPendingRootEpoch(); + error NoRootEpochRegistered(); + error UnknownRootEpoch(); + error InvalidAppId(); + error InvalidImplementationAddress(); + error AppImplementationNotSet(); + error InvalidOwnerAddress(); + error InvalidRootEpochRange(); + error AppNotRegistered(); + error AppNotDeployed(); + // Events // appId is unindexed for backward compatibility with deployed log indexers. // slither-disable-next-line unindexed-event-address @@ -138,7 +155,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E if (_appImplementation != address(0)) { appImplementation = _appImplementation; emit AppImplementationSet(_appImplementation); - _emitPolicy("app-implementation", bytes32(uint256(uint160(_appImplementation))), true); + _emitPolicy(keccak256("app-implementation"), bytes32(uint256(uint160(_appImplementation))), true); } } @@ -161,11 +178,11 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // Function to authorize upgrades (required by UUPSUpgradeable) function _authorizeUpgrade(address newImplementation) internal override onlyOwner { - _emitPolicy("implementation-upgrade", bytes32(uint256(uint160(newImplementation))), true); + _emitPolicy(keccak256("implementation-upgrade"), bytes32(uint256(uint160(newImplementation))), true); } - function _emitPolicy(string memory policy, bytes32 value, bool enabled) internal { - emit PolicyChanged(msg.sender, keccak256(bytes(policy)), value, enabled); + function _emitPolicy(bytes32 policy, bytes32 value, bool enabled) internal { + emit PolicyChanged(msg.sender, policy, value, enabled); } // Register the root public keys of the next KMS root epoch, with where to @@ -183,24 +200,23 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E onlyOwner returns (uint256 epoch) { - require(caPubkey.length == 91 && bytes27(caPubkey[:27]) == P256_SPKI_PREFIX, "Invalid epoch CA public key"); - require( - k256Pubkey.length == 33 && (k256Pubkey[0] == 0x02 || k256Pubkey[0] == 0x03), "Invalid epoch k256 public key" - ); - require(bytes(evidenceUrl).length > 0, "Empty epoch evidence URL"); + if (caPubkey.length != 91 || bytes27(caPubkey[:27]) != P256_SPKI_PREFIX) revert InvalidEpochCaPubkey(); + if (k256Pubkey.length != 33 || (k256Pubkey[0] != 0x02 && k256Pubkey[0] != 0x03)) { + revert InvalidEpochK256Pubkey(); + } + if (bytes(evidenceUrl).length == 0) revert EmptyEpochEvidenceUrl(); epoch = _rootEpochs.length; if (epoch == 0 && kmsInfo.k256Pubkey.length > 0) { - require( - keccak256(caPubkey) == keccak256(kmsInfo.caPubkey) - && keccak256(k256Pubkey) == keccak256(kmsInfo.k256Pubkey), - "Epoch 0 keys differ from kmsInfo" - ); + if ( + keccak256(caPubkey) != keccak256(kmsInfo.caPubkey) + || keccak256(k256Pubkey) != keccak256(kmsInfo.k256Pubkey) + ) revert Epoch0KeysDifferFromKmsInfo(); } bytes32 previous = epoch == 0 ? bytes32(0) : _rootEpochs[epoch - 1].chainHash; bytes32 chainHash = keccak256(abi.encodePacked(previous, keccak256(caPubkey), keccak256(k256Pubkey))); _rootEpochs.push(RootEpoch(caPubkey, k256Pubkey, evidenceUrl, chainHash, _rootEpochRegistrations++)); emit RootEpochRegistered(epoch, chainHash, caPubkey, k256Pubkey, evidenceUrl); - _emitPolicy("root-epoch", chainHash, true); + _emitPolicy(keccak256("root-epoch"), chainHash, true); if (epoch == 0) { _useRootEpoch(0); } @@ -209,11 +225,11 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // Make the next registered root epoch the one KMS instances serve. Bind // activation to its chain hash in case a pending epoch has been replaced. function activateRootEpoch(uint256 epoch, bytes32 chainHash) external onlyOwner { - require(epoch == _currentRootEpoch + 1 && epoch < _rootEpochs.length, "Invalid root epoch activation"); - require(_rootEpochs[epoch].chainHash == chainHash, "root epoch chain hash mismatch"); + if (epoch != _currentRootEpoch + 1 || epoch >= _rootEpochs.length) revert InvalidRootEpochActivation(); + if (_rootEpochs[epoch].chainHash != chainHash) revert RootEpochChainHashMismatch(); _currentRootEpoch = epoch; _useRootEpoch(epoch); - _emitPolicy("root-epoch-activated", bytes32(epoch), true); + _emitPolicy(keccak256("root-epoch-activated"), bytes32(epoch), true); } // Point kmsInfo at the keys of the now current epoch, for readers of the @@ -232,8 +248,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // Point an epoch at a new copy of its evidence, for when the old location // goes away. The evidence proves itself, so only its location changes. function setRootEpochEvidenceUrl(uint256 epoch, string calldata evidenceUrl) external onlyOwner { - require(epoch < _rootEpochs.length, "Unknown root epoch"); - require(bytes(evidenceUrl).length > 0, "Empty epoch evidence URL"); + if (epoch >= _rootEpochs.length) revert UnknownRootEpoch(); + if (bytes(evidenceUrl).length == 0) revert EmptyEpochEvidenceUrl(); _rootEpochs[epoch].evidenceUrl = evidenceUrl; emit RootEpochEvidenceUrlSet(epoch, evidenceUrl); } @@ -245,11 +261,11 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // next registration of the number starts with empty rows, whatever its keys, // which must be prepared again before activation. function popRootEpoch(uint256 epoch, bytes32 chainHash) external onlyOwner { - require(epoch + 1 == _rootEpochs.length && epoch > _currentRootEpoch, "Not the pending root epoch"); - require(_rootEpochs[epoch].chainHash == chainHash, "Root epoch chain hash mismatch"); + if (epoch + 1 != _rootEpochs.length || epoch <= _currentRootEpoch) revert NotPendingRootEpoch(); + if (_rootEpochs[epoch].chainHash != chainHash) revert RootEpochChainHashMismatch(); _rootEpochs.pop(); emit RootEpochPopped(epoch, chainHash); - _emitPolicy("root-epoch", chainHash, false); + _emitPolicy(keccak256("root-epoch"), chainHash, false); } function rootEpochCount() external view returns (uint256) { @@ -258,7 +274,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E // The root epoch KMS instances serve and its chain hash function currentRootEpoch() external view returns (uint256 epoch, bytes32 chainHash) { - require(_rootEpochs.length > 0, "No root epoch registered"); + if (_rootEpochs.length == 0) revert NoRootEpochRegistered(); epoch = _currentRootEpoch; chainHash = _rootEpochs[epoch].chainHash; } @@ -268,7 +284,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E view returns (bytes memory caPubkey, bytes memory k256Pubkey, string memory evidenceUrl, bytes32 chainHash) { - require(epoch < _rootEpochs.length, "Unknown root epoch"); + if (epoch >= _rootEpochs.length) revert UnknownRootEpoch(); RootEpoch storage keys = _rootEpochs[epoch]; return (keys.caPubkey, keys.k256Pubkey, keys.evidenceUrl, keys.chainHash); } @@ -288,20 +304,20 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E function setAllowAnyTeeVariant(bool allowed) external onlyOwner { _teeVariants.restricted = !allowed; emit AllowAnyTeeVariantSet(allowed); - _emitPolicy("allow-any-tee-variant", bytes32(0), allowed); + _emitPolicy(keccak256("allow-any-tee-variant"), bytes32(0), allowed); } function setAllowedTeeVariants(string[] calldata teeVariants) external onlyOwner { uint64 mask = TeeVariantPolicy.setAllowed(_teeVariants, teeVariants); emit AllowedTeeVariantsSet(teeVariants); - _emitPolicy("allowed-tee-variants", bytes32(uint256(mask)), true); + _emitPolicy(keccak256("allowed-tee-variants"), bytes32(uint256(mask)), true); } // Function to set trusted Gateway App ID function setGatewayAppId(string memory appId) external onlyOwner { gatewayAppId = appId; emit GatewayAppIdSet(appId); - _emitPolicy("gateway-app-id", keccak256(bytes(appId)), true); + _emitPolicy(keccak256("gateway-app-id"), keccak256(bytes(appId)), true); } /// @notice Register an app address as known to this KMS. @@ -314,18 +330,18 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E /// and by the registered app's own {IAppAuth-isAppAllowed}, so /// registration alone confers no privilege. function registerApp(address appId) public { - require(appId != address(0), "Invalid app ID"); + if (appId == address(0)) revert InvalidAppId(); registeredApps[appId] = true; emit AppRegistered(appId); - _emitPolicy("registered-app", bytes32(uint256(uint160(appId))), true); + _emitPolicy(keccak256("registered-app"), bytes32(uint256(uint160(appId))), true); } // Function to set DstackApp implementation contract address function setAppImplementation(address _implementation) external onlyOwner { - require(_implementation != address(0), "Invalid implementation address"); + if (_implementation == address(0)) revert InvalidImplementationAddress(); appImplementation = _implementation; emit AppImplementationSet(_implementation); - _emitPolicy("app-implementation", bytes32(uint256(uint160(_implementation))), true); + _emitPolicy(keccak256("app-implementation"), bytes32(uint256(uint160(_implementation))), true); } // Factory method: Deploy and register DstackApp in single transaction @@ -340,8 +356,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E public returns (address appId) { - require(appImplementation != address(0), "DstackApp implementation not set"); - require(initialOwner != address(0), "Invalid owner address"); + if (appImplementation == address(0)) revert AppImplementationNotSet(); + if (initialOwner == address(0)) revert InvalidOwnerAddress(); bytes memory initData = abi.encodeWithSelector( bytes4(keccak256("initialize(address,bool,bool,bool,bytes32,bytes32)")), @@ -387,18 +403,10 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E return epoch < _rootEpochs.length && _kmsAggregatedMrsByRoot[_rootEpochs[epoch].registration][mrAggregated]; } - function addKmsAggregatedMrAtEpoch(uint256 epoch, bytes32 mrAggregated) public onlyOwner { - _setKmsAggregatedMr(epoch, mrAggregated, true); - } - function addKmsAggregatedMrAtEpochs(uint256 fromEpoch, uint256 toEpoch, bytes32 mrAggregated) external onlyOwner { _setAtEpochs(_setKmsAggregatedMr, fromEpoch, toEpoch, mrAggregated, true); } - function removeKmsAggregatedMrAtEpoch(uint256 epoch, bytes32 mrAggregated) public onlyOwner { - _setKmsAggregatedMr(epoch, mrAggregated, false); - } - function removeKmsAggregatedMrAtEpochs( uint256 fromEpoch, uint256 toEpoch, @@ -411,12 +419,12 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E } function _setKmsAggregatedMr(uint256 epoch, bytes32 mrAggregated, bool allowed) private { - require(epoch <= _latestRootEpoch(), "Unknown root epoch"); + if (epoch > _latestRootEpoch()) revert UnknownRootEpoch(); if (epoch == 0) { kmsAllowedAggregatedMrs[mrAggregated] = allowed; if (allowed) emit KmsAggregatedMrAdded(mrAggregated); else emit KmsAggregatedMrRemoved(mrAggregated); - _emitPolicy("kms-aggregated-mr", mrAggregated, allowed); + _emitPolicy(keccak256("kms-aggregated-mr"), mrAggregated, allowed); } else { _kmsAggregatedMrsByRoot[_rootEpochs[epoch].registration][mrAggregated] = allowed; } @@ -428,29 +436,21 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E return epoch < _rootEpochs.length && _kmsDeviceIdsByRoot[_rootEpochs[epoch].registration][deviceId]; } - function addKmsDeviceAtEpoch(uint256 epoch, bytes32 deviceId) public onlyOwner { - _setKmsDevice(epoch, deviceId, true); - } - function addKmsDeviceAtEpochs(uint256 fromEpoch, uint256 toEpoch, bytes32 deviceId) external onlyOwner { _setAtEpochs(_setKmsDevice, fromEpoch, toEpoch, deviceId, true); } - function removeKmsDeviceAtEpoch(uint256 epoch, bytes32 deviceId) public onlyOwner { - _setKmsDevice(epoch, deviceId, false); - } - function removeKmsDeviceAtEpochs(uint256 fromEpoch, uint256 toEpoch, bytes32 deviceId) external onlyOwner { _setAtEpochs(_setKmsDevice, fromEpoch, toEpoch, deviceId, false); } function _setKmsDevice(uint256 epoch, bytes32 deviceId, bool allowed) private { - require(epoch <= _latestRootEpoch(), "Unknown root epoch"); + if (epoch > _latestRootEpoch()) revert UnknownRootEpoch(); if (epoch == 0) { kmsAllowedDeviceIds[deviceId] = allowed; if (allowed) emit KmsDeviceAdded(deviceId); else emit KmsDeviceRemoved(deviceId); - _emitPolicy("kms-device", deviceId, allowed); + _emitPolicy(keccak256("kms-device"), deviceId, allowed); } else { _kmsDeviceIdsByRoot[_rootEpochs[epoch].registration][deviceId] = allowed; } @@ -462,29 +462,21 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E return epoch < _rootEpochs.length && _osImagesByRoot[_rootEpochs[epoch].registration][osImageHash]; } - function addOsImageHashAtEpoch(uint256 epoch, bytes32 osImageHash) public onlyOwner { - _setOsImageHash(epoch, osImageHash, true); - } - function addOsImageHashAtEpochs(uint256 fromEpoch, uint256 toEpoch, bytes32 osImageHash) external onlyOwner { _setAtEpochs(_setOsImageHash, fromEpoch, toEpoch, osImageHash, true); } - function removeOsImageHashAtEpoch(uint256 epoch, bytes32 osImageHash) public onlyOwner { - _setOsImageHash(epoch, osImageHash, false); - } - function removeOsImageHashAtEpochs(uint256 fromEpoch, uint256 toEpoch, bytes32 osImageHash) external onlyOwner { _setAtEpochs(_setOsImageHash, fromEpoch, toEpoch, osImageHash, false); } function _setOsImageHash(uint256 epoch, bytes32 osImageHash, bool allowed) private { - require(epoch <= _latestRootEpoch(), "Unknown root epoch"); + if (epoch > _latestRootEpoch()) revert UnknownRootEpoch(); if (epoch == 0) { allowedOsImages[osImageHash] = allowed; if (allowed) emit OsImageHashAdded(osImageHash); else emit OsImageHashRemoved(osImageHash); - _emitPolicy("os-image", osImageHash, allowed); + _emitPolicy(keccak256("os-image"), osImageHash, allowed); } else { _osImagesByRoot[_rootEpochs[epoch].registration][osImageHash] = allowed; } @@ -503,7 +495,7 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E ) private { - require(fromEpoch <= toEpoch, "Invalid root epoch range"); + if (fromEpoch > toEpoch) revert InvalidRootEpochRange(); for (uint256 epoch = fromEpoch; epoch <= toEpoch; ++epoch) { setter(epoch, value, allowed); } @@ -516,8 +508,8 @@ contract DstackKms is Initializable, Ownable2StepUpgradeable, UUPSUpgradeable, E /// Reverts for an unregistered or undeployed app rather than reporting epoch /// zero for it. Legacy policies have epoch zero; an advertised getter must succeed. function currentAppEpoch(address appId) external view returns (uint256) { - require(registeredApps[appId], "App not registered"); - require(appId.code.length != 0, "App not deployed or invalid address"); + if (!registeredApps[appId]) revert AppNotRegistered(); + if (appId.code.length == 0) revert AppNotDeployed(); return _appEpoch(appId); } diff --git a/dstack/kms/auth-eth/docs/specification.md b/dstack/kms/auth-eth/docs/specification.md index 5bbb1d7f7..0d3263675 100644 --- a/dstack/kms/auth-eth/docs/specification.md +++ b/dstack/kms/auth-eth/docs/specification.md @@ -55,9 +55,9 @@ return-value spoofing are still possible — see §6. | 0–3 | `kmsInfo` | `KmsInfo` (struct, 4 fields) | `registerRootEpoch` (epoch 0), `activateRootEpoch` (owner) | Preserved for proxy upgrades and legacy readers; quote and event log are cleared on activation | | 4 | `gatewayAppId` | `string` | `setGatewayAppId` (owner) | Renamed from `tproxyAppId`; carries `@custom:oz-renamed-from` | | 5 | `registeredApps` | `mapping(address => bool)` | `registerApp` (anyone), `deployAndRegisterApp` (anyone) | Permissionless — see §6.1 | -| 6 | `kmsAllowedAggregatedMrs` | `mapping(bytes32 => bool)` | `addKmsAggregatedMrAtEpoch(s)`, `removeKmsAggregatedMrAtEpoch(s)` at epoch 0 (owner) | | -| 7 | `kmsAllowedDeviceIds` | `mapping(bytes32 => bool)` | `addKmsDeviceAtEpoch(s)`, `removeKmsDeviceAtEpoch(s)` at epoch 0 (owner) | | -| 8 | `allowedOsImages` | `mapping(bytes32 => bool)` | `addOsImageHashAtEpoch(s)`, `removeOsImageHashAtEpoch(s)` at epoch 0 (owner) | | +| 6 | `kmsAllowedAggregatedMrs` | `mapping(bytes32 => bool)` | `addKmsAggregatedMrAtEpochs`, `removeKmsAggregatedMrAtEpochs` at epoch 0 (owner) | | +| 7 | `kmsAllowedDeviceIds` | `mapping(bytes32 => bool)` | `addKmsDeviceAtEpochs`, `removeKmsDeviceAtEpochs` at epoch 0 (owner) | | +| 8 | `allowedOsImages` | `mapping(bytes32 => bool)` | `addOsImageHashAtEpochs`, `removeOsImageHashAtEpochs` at epoch 0 (owner) | | | 9 | `appImplementation` | `address` | `setAppImplementation` (owner) | Used by `deployAndRegisterApp` | | 10 | `__gap[50]` | `uint256[50]` | (none — must always read zero) | Reserved | @@ -173,7 +173,7 @@ setup uses `registerRootEpoch`. On an upgraded proxy with existing legacy roots, epoch 0 must match the stored `kmsInfo` keys; the upgrade preserves that storage. For each of `setGatewayAppId`, -`setAppImplementation`, and the `*AtEpoch` and `*AtEpochs` forms of +`setAppImplementation`, and the `*AtEpochs` forms of `addKmsAggregatedMr`, `removeKmsAggregatedMr`, `addKmsDevice`, `removeKmsDevice`, `addOsImageHash`, `removeOsImageHash`: diff --git a/dstack/kms/auth-eth/index.test.ts b/dstack/kms/auth-eth/index.test.ts index 7e0e2fee2..04b4b825f 100644 --- a/dstack/kms/auth-eth/index.test.ts +++ b/dstack/kms/auth-eth/index.test.ts @@ -868,7 +868,7 @@ describe('request hardening', () => { describe('GET /app/:appId', () => { it('reports an app epoch only when the contract can vouch for it', async () => { const url = `http://localhost/app/0x${'11'.repeat(20)}`; - mockReadContract.mockRejectedValueOnce(new Error('execution reverted: App not registered')); + mockReadContract.mockRejectedValueOnce(new Error('execution reverted: AppNotRegistered()')); expect((await appFetch(new Request(url))).status).toBe(503); mockReadContract.mockResolvedValueOnce(0n); const legacy = await appFetch(new Request(url)); diff --git a/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol b/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol index ba34bb846..20887ef26 100644 --- a/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.symbolic.t.sol @@ -125,7 +125,7 @@ contract DstackKmsSymbolicTest is Test { vm.startPrank(OWNER); kms.registerApp(address(mock)); - kms.addOsImageHashAtEpoch(0, osImageHash); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); vm.stopPrank(); IAppAuth.AppInfo memory b = IAppAuth.AppInfo({ @@ -155,7 +155,7 @@ contract DstackKmsSymbolicTest is Test { vm.startPrank(OWNER); kms.registerApp(address(mock)); - kms.addOsImageHashAtEpoch(0, osImageHash); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); vm.stopPrank(); IAppAuth.AppInfo memory b = IAppAuth.AppInfo({ @@ -190,7 +190,7 @@ contract DstackKmsSymbolicTest is Test { // Allow the OS image so the second gate passes; leave // kmsAllowedAggregatedMrs empty. vm.prank(OWNER); - kms.addOsImageHashAtEpoch(0, osImageHash); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); IAppAuth.AppInfo memory b = IAppAuth.AppInfo({ appId: address(0), @@ -217,8 +217,8 @@ contract DstackKmsSymbolicTest is Test { external { vm.startPrank(OWNER); - kms.addOsImageHashAtEpoch(0, osImageHash); - kms.addKmsAggregatedMrAtEpoch(0, mrAggregated); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); + kms.addKmsAggregatedMrAtEpochs(0, 0, mrAggregated); vm.stopPrank(); // kmsAllowedDeviceIds is left empty. diff --git a/dstack/kms/auth-eth/test/DstackKms.t.sol b/dstack/kms/auth-eth/test/DstackKms.t.sol index 31bb3ead3..ea8b08407 100644 --- a/dstack/kms/auth-eth/test/DstackKms.t.sol +++ b/dstack/kms/auth-eth/test/DstackKms.t.sol @@ -134,7 +134,7 @@ contract DstackKmsTest is Test { function test_RegisterRootEpochAppendsAndChains() public { assertEq(kms.rootEpochCount(), 0); - vm.expectRevert("No root epoch registered"); + vm.expectRevert(DstackKms.NoRootEpochRegistered.selector); kms.currentRootEpoch(); bytes32 h0 = chainHash(bytes32(0), ca(0xa0), k256(0xb0)); @@ -157,7 +157,7 @@ contract DstackKmsTest is Test { assertEq(k256Key, k256(0xb1)); assertEq(evidenceUrl, "ipfs://c1"); assertEq(h, h1); - vm.expectRevert("Unknown root epoch"); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); kms.rootEpoch(2); // Epoch 1 is registered but epoch 0 is still served. @@ -172,18 +172,18 @@ contract DstackKmsTest is Test { bytes32 h2 = chainHash(h1, ca(0xa2), k256(0xb2)); vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); - vm.expectRevert("Invalid root epoch activation"); + vm.expectRevert(DstackKms.InvalidRootEpochActivation.selector); kms.activateRootEpoch(1, h1); kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://c2"); - vm.expectRevert("Invalid root epoch activation"); + vm.expectRevert(DstackKms.InvalidRootEpochActivation.selector); kms.activateRootEpoch(2, h2); vm.expectEmit(true, true, true, true); emit RootEpochActivated(1); kms.activateRootEpoch(1, h1); - vm.expectRevert("Invalid root epoch activation"); + vm.expectRevert(DstackKms.InvalidRootEpochActivation.selector); kms.activateRootEpoch(1, h1); - vm.expectRevert("Invalid root epoch activation"); + vm.expectRevert(DstackKms.InvalidRootEpochActivation.selector); kms.activateRootEpoch(0, h0); vm.stopPrank(); @@ -205,7 +205,7 @@ contract DstackKmsTest is Test { kms.popRootEpoch(1, h1); kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://replacement"); - vm.expectRevert("root epoch chain hash mismatch"); + vm.expectRevert(DstackKms.RootEpochChainHashMismatch.selector); kms.activateRootEpoch(1, h1); (uint256 current, bytes32 currentHash) = kms.currentRootEpoch(); assertEq(current, 0); @@ -227,9 +227,9 @@ contract DstackKmsTest is Test { function test_RegisterRootEpochZeroMustMatchKmsInfo() public { upgradeLegacyKms(); vm.startPrank(owner); - vm.expectRevert("Epoch 0 keys differ from kmsInfo"); + vm.expectRevert(DstackKms.Epoch0KeysDifferFromKmsInfo.selector); kms.registerRootEpoch(ca(0xa1), k256(0xb0), "ipfs://c0"); - vm.expectRevert("Epoch 0 keys differ from kmsInfo"); + vm.expectRevert(DstackKms.Epoch0KeysDifferFromKmsInfo.selector); kms.registerRootEpoch(ca(0xa0), k256(0xb1), "ipfs://c0"); assertEq(kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"), 0); // Later epochs are new keys by definition. @@ -242,15 +242,15 @@ contract DstackKmsTest is Test { bytes32 h1 = chainHash(h0, ca(0xa1), k256(0xb1)); bytes32 h1b = chainHash(h0, ca(0xa2), k256(0xb2)); vm.startPrank(owner); - vm.expectRevert("Not the pending root epoch"); + vm.expectRevert(DstackKms.NotPendingRootEpoch.selector); kms.popRootEpoch(0, h0); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); - vm.expectRevert("Not the pending root epoch"); + vm.expectRevert(DstackKms.NotPendingRootEpoch.selector); kms.popRootEpoch(0, h0); kms.registerRootEpoch(ca(0xa1), k256(0xb1), "ipfs://c1"); - vm.expectRevert("Not the pending root epoch"); + vm.expectRevert(DstackKms.NotPendingRootEpoch.selector); kms.popRootEpoch(2, h1); - vm.expectRevert("Root epoch chain hash mismatch"); + vm.expectRevert(DstackKms.RootEpochChainHashMismatch.selector); kms.popRootEpoch(1, h1b); vm.expectEmit(true, true, true, true); @@ -261,10 +261,10 @@ contract DstackKmsTest is Test { // A new epoch 1 replaces the dropped one; a stale pop of the old one // cannot drop it. assertEq(kms.registerRootEpoch(ca(0xa2), k256(0xb2), "ipfs://c2"), 1); - vm.expectRevert("Root epoch chain hash mismatch"); + vm.expectRevert(DstackKms.RootEpochChainHashMismatch.selector); kms.popRootEpoch(1, h1); kms.activateRootEpoch(1, h1b); - vm.expectRevert("Not the pending root epoch"); + vm.expectRevert(DstackKms.NotPendingRootEpoch.selector); kms.popRootEpoch(1, h1b); vm.stopPrank(); @@ -286,12 +286,12 @@ contract DstackKmsTest is Test { assertTrue(kms.kmsAllowedAggregatedMrs(info.mrAggregated)); assertTrue(kms.kmsAllowedDeviceIds(info.deviceId)); // Rows of an unregistered epoch can be neither written nor read. - vm.expectRevert("Unknown root epoch"); - kms.addOsImageHashAtEpoch(1, info.osImageHash); - vm.expectRevert("Unknown root epoch"); - kms.addKmsAggregatedMrAtEpoch(1, info.mrAggregated); - vm.expectRevert("Unknown root epoch"); - kms.addKmsDeviceAtEpoch(1, info.deviceId); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.addOsImageHashAtEpochs(1, 1, info.osImageHash); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.addKmsAggregatedMrAtEpochs(1, 1, info.mrAggregated); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.addKmsDeviceAtEpochs(1, 1, info.deviceId); registerNextRootEpoch(); registerNextRootEpoch(); approveKmsAtEpoch(info, 1); @@ -315,10 +315,10 @@ contract DstackKmsTest is Test { function test_SetRootEpochEvidenceUrl() public { vm.startPrank(owner); - vm.expectRevert("Unknown root epoch"); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); kms.setRootEpochEvidenceUrl(0, "ipfs://moved"); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "ipfs://c0"); - vm.expectRevert("Empty epoch evidence URL"); + vm.expectRevert(DstackKms.EmptyEpochEvidenceUrl.selector); kms.setRootEpochEvidenceUrl(0, ""); vm.expectEmit(true, true, true, true); emit RootEpochEvidenceUrlSet(0, "https://mirror/c0"); @@ -371,21 +371,21 @@ contract DstackKmsTest is Test { otherCurve[25] = 0x0a; bytes memory uncompressedK256 = bytes.concat(hex"04", new bytes(64)); vm.startPrank(owner); - vm.expectRevert("Invalid epoch CA public key"); + vm.expectRevert(DstackKms.InvalidEpochCaPubkey.selector); kms.registerRootEpoch("", k256(0xb0), "ipfs://c0"); - vm.expectRevert("Invalid epoch CA public key"); + vm.expectRevert(DstackKms.InvalidEpochCaPubkey.selector); kms.registerRootEpoch(bytes.concat(ca(0xa0), hex"00"), k256(0xb0), "ipfs://c0"); - vm.expectRevert("Invalid epoch CA public key"); + vm.expectRevert(DstackKms.InvalidEpochCaPubkey.selector); kms.registerRootEpoch(compressedCa, k256(0xb0), "ipfs://c0"); - vm.expectRevert("Invalid epoch CA public key"); + vm.expectRevert(DstackKms.InvalidEpochCaPubkey.selector); kms.registerRootEpoch(otherCurve, k256(0xb0), "ipfs://c0"); - vm.expectRevert("Invalid epoch k256 public key"); + vm.expectRevert(DstackKms.InvalidEpochK256Pubkey.selector); kms.registerRootEpoch(ca(0xa0), "", "ipfs://c0"); - vm.expectRevert("Invalid epoch k256 public key"); + vm.expectRevert(DstackKms.InvalidEpochK256Pubkey.selector); kms.registerRootEpoch(ca(0xa0), uncompressedK256, "ipfs://c0"); - vm.expectRevert("Invalid epoch k256 public key"); + vm.expectRevert(DstackKms.InvalidEpochK256Pubkey.selector); kms.registerRootEpoch(ca(0xa0), bytes.concat(hex"04", new bytes(32)), "ipfs://c0"); - vm.expectRevert("Empty epoch evidence URL"); + vm.expectRevert(DstackKms.EmptyEpochEvidenceUrl.selector); kms.registerRootEpoch(ca(0xa0), k256(0xb0), ""); vm.stopPrank(); } @@ -394,7 +394,7 @@ contract DstackKmsTest is Test { bytes32 mr = bytes32(uint256(123)); vm.prank(owner); - kms.addKmsAggregatedMrAtEpoch(0, mr); + kms.addKmsAggregatedMrAtEpochs(0, 0, mr); assertTrue(kms.kmsAllowedAggregatedMrs(mr)); } @@ -403,10 +403,10 @@ contract DstackKmsTest is Test { bytes32 mr = bytes32(uint256(123)); vm.startPrank(owner); - kms.addKmsAggregatedMrAtEpoch(0, mr); + kms.addKmsAggregatedMrAtEpochs(0, 0, mr); assertTrue(kms.kmsAllowedAggregatedMrs(mr)); - kms.removeKmsAggregatedMrAtEpoch(0, mr); + kms.removeKmsAggregatedMrAtEpochs(0, 0, mr); assertFalse(kms.kmsAllowedAggregatedMrs(mr)); vm.stopPrank(); } @@ -458,10 +458,10 @@ contract DstackKmsTest is Test { bytes32 deviceId = bytes32(uint256(999)); vm.startPrank(owner); - kms.addKmsDeviceAtEpoch(0, deviceId); + kms.addKmsDeviceAtEpochs(0, 0, deviceId); assertTrue(kms.kmsAllowedDeviceIds(deviceId)); - kms.removeKmsDeviceAtEpoch(0, deviceId); + kms.removeKmsDeviceAtEpochs(0, 0, deviceId); assertFalse(kms.kmsAllowedDeviceIds(deviceId)); vm.stopPrank(); } @@ -470,10 +470,10 @@ contract DstackKmsTest is Test { bytes32 imageHash = bytes32(uint256(888)); vm.startPrank(owner); - kms.addOsImageHashAtEpoch(0, imageHash); + kms.addOsImageHashAtEpochs(0, 0, imageHash); assertTrue(kms.allowedOsImages(imageHash)); - kms.removeOsImageHashAtEpoch(0, imageHash); + kms.removeOsImageHashAtEpochs(0, 0, imageHash); assertFalse(kms.allowedOsImages(imageHash)); vm.stopPrank(); } @@ -511,9 +511,9 @@ contract DstackKmsTest is Test { tdx[1] = "dstack-gcp-tdx"; vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); - kms.addOsImageHashAtEpoch(0, info.osImageHash); - kms.addKmsDeviceAtEpoch(0, info.deviceId); - kms.addKmsAggregatedMrAtEpoch(0, info.mrAggregated); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); + kms.addKmsDeviceAtEpochs(0, 0, info.deviceId); + kms.addKmsAggregatedMrAtEpochs(0, 0, info.mrAggregated); address appId = kms.deployAndRegisterApp(user, false, true, info.deviceId, info.composeHash); kms.setAllowedTeeVariants(tdx); kms.setAllowAnyTeeVariant(false); @@ -572,9 +572,9 @@ contract DstackKmsTest is Test { function test_AuthorizationOperationsFailClosedAndBootstrapBeforeRegistration() public { IAppAuth.AppInfo memory info = authorizationInfo(address(kms)); vm.startPrank(owner); - kms.addOsImageHashAtEpoch(0, info.osImageHash); - kms.addKmsDeviceAtEpoch(0, info.deviceId); - kms.addKmsAggregatedMrAtEpoch(0, info.mrAggregated); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); + kms.addKmsDeviceAtEpochs(0, 0, info.deviceId); + kms.addKmsAggregatedMrAtEpochs(0, 0, info.mrAggregated); vm.stopPrank(); AuthorizationRequest memory request = AuthorizationRequest( info, keccak256("kms_bootstrap"), OptionalUint256(false, 0), OptionalUint256(false, 0), "dstack-tdx" @@ -635,9 +635,9 @@ contract DstackKmsTest is Test { vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); kms.setGatewayAppId("gateway"); - kms.addOsImageHashAtEpoch(0, info.osImageHash); - kms.addKmsDeviceAtEpoch(0, info.deviceId); - kms.addKmsAggregatedMrAtEpoch(0, info.mrAggregated); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); + kms.addKmsDeviceAtEpochs(0, 0, info.deviceId); + kms.addKmsAggregatedMrAtEpochs(0, 0, info.mrAggregated); address appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); vm.stopPrank(); @@ -701,11 +701,11 @@ contract DstackKmsTest is Test { function test_AppEpochQueryRevertsForUnregisteredOrUndeployedApps() public { address policy = address(new DstackApp()); - vm.expectRevert("App not registered"); + vm.expectRevert(DstackKms.AppNotRegistered.selector); kms.currentAppEpoch(policy); address eoa = address(0xbeef); kms.registerApp(eoa); - vm.expectRevert("App not deployed or invalid address"); + vm.expectRevert(DstackKms.AppNotDeployed.selector); kms.currentAppEpoch(eoa); // A registered legacy policy without the capability is epoch zero. kms.registerApp(policy); @@ -727,7 +727,7 @@ contract DstackKmsTest is Test { DstackKms.AppInfo memory info = authorizationInfo(address(0)); vm.startPrank(owner); kms.registerRootEpoch(ca(0xa0), k256(0xb0), "https://example.com/epoch-0"); - kms.addOsImageHashAtEpoch(0, info.osImageHash); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); address appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); address otherApp = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); vm.stopPrank(); @@ -803,9 +803,9 @@ contract DstackKmsTest is Test { // Setup allowed values vm.startPrank(owner); - kms.addKmsDeviceAtEpoch(0, deviceId); - kms.addKmsAggregatedMrAtEpoch(0, mrAggregated); - kms.addOsImageHashAtEpoch(0, osImageHash); + kms.addKmsDeviceAtEpochs(0, 0, deviceId); + kms.addKmsAggregatedMrAtEpochs(0, 0, mrAggregated); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); vm.stopPrank(); IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ @@ -832,9 +832,9 @@ contract DstackKmsTest is Test { // Setup allowed values vm.startPrank(owner); - kms.addKmsDeviceAtEpoch(0, deviceId); - kms.addKmsAggregatedMrAtEpoch(0, mrAggregated); - kms.addOsImageHashAtEpoch(0, osImageHash); + kms.addKmsDeviceAtEpochs(0, 0, deviceId); + kms.addKmsAggregatedMrAtEpochs(0, 0, mrAggregated); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); vm.stopPrank(); IAppAuth.AppInfo memory bootInfo = IAppAuth.AppInfo({ @@ -861,7 +861,7 @@ contract DstackKmsTest is Test { bytes32 osImageHash = bytes32(uint256(111)); vm.startPrank(owner); - kms.addOsImageHashAtEpoch(0, osImageHash); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); address appId = kms.deployAndRegisterApp(user, false, false, deviceId, composeHash); vm.stopPrank(); @@ -993,11 +993,11 @@ contract DstackKmsTest is Test { // Old owner has lost authority. vm.prank(owner); vm.expectRevert(abi.encodeWithSignature("OwnableUnauthorizedAccount(address)", owner)); - kms.addKmsAggregatedMrAtEpoch(0, mr); + kms.addKmsAggregatedMrAtEpochs(0, 0, mr); // New owner has it. vm.prank(newOwner); - kms.addKmsAggregatedMrAtEpoch(0, mr); + kms.addKmsAggregatedMrAtEpochs(0, 0, mr); assertTrue(kms.kmsAllowedAggregatedMrs(mr)); } @@ -1040,9 +1040,9 @@ contract DstackKmsTest is Test { } function approveKmsAtEpoch(IAppAuth.AppInfo memory info, uint256 epoch) internal { - kms.addOsImageHashAtEpoch(epoch, info.osImageHash); - kms.addKmsAggregatedMrAtEpoch(epoch, info.mrAggregated); - kms.addKmsDeviceAtEpoch(epoch, info.deviceId); + kms.addOsImageHashAtEpochs(epoch, epoch, info.osImageHash); + kms.addKmsAggregatedMrAtEpochs(epoch, epoch, info.mrAggregated); + kms.addKmsDeviceAtEpochs(epoch, epoch, info.deviceId); } function test_KmsEpochRowsDoNotInheritAndCanBeRevokedIndependently() public { @@ -1057,22 +1057,22 @@ contract DstackKmsTest is Test { (bool allowed, string memory reason) = kms.isKmsAllowed(info); assertFalse(allowed, "legacy KMS check ends at rotation"); assertEq(reason, "Root epoch requires epoch-aware authorization"); - kms.addOsImageHashAtEpoch(1, info.osImageHash); + kms.addOsImageHashAtEpochs(1, 1, info.osImageHash); assertFalse(authorize(info, true).authorized); - kms.addKmsAggregatedMrAtEpoch(1, info.mrAggregated); + kms.addKmsAggregatedMrAtEpochs(1, 1, info.mrAggregated); assertFalse(authorize(info, true).authorized); - kms.addKmsDeviceAtEpoch(1, info.deviceId); + kms.addKmsDeviceAtEpochs(1, 1, info.deviceId); assertTrue(authorize(info, true).authorized); - kms.removeOsImageHashAtEpoch(1, info.osImageHash); + kms.removeOsImageHashAtEpochs(1, 1, info.osImageHash); assertFalse(authorize(info, true).authorized); assertTrue(kms.allowedOsImages(info.osImageHash)); assertTrue(kms.authorizeKms(kmsEpochRequest(info, 0, 0)).authorized); - kms.addOsImageHashAtEpoch(1, info.osImageHash); - kms.removeKmsAggregatedMrAtEpoch(0, info.mrAggregated); + kms.addOsImageHashAtEpochs(1, 1, info.osImageHash); + kms.removeKmsAggregatedMrAtEpochs(0, 0, info.mrAggregated); assertFalse(kms.authorizeKms(kmsEpochRequest(info, 0, 0)).authorized); assertTrue(authorize(info, true).authorized); - kms.addKmsAggregatedMrAtEpoch(0, info.mrAggregated); - kms.removeKmsDeviceAtEpoch(0, info.deviceId); + kms.addKmsAggregatedMrAtEpochs(0, 0, info.mrAggregated); + kms.removeKmsDeviceAtEpochs(0, 0, info.deviceId); assertFalse(kms.authorizeKms(kmsEpochRequest(info, 0, 0)).authorized); assertTrue(authorize(info, true).authorized); vm.stopPrank(); @@ -1082,7 +1082,7 @@ contract DstackKmsTest is Test { IAppAuth.AppInfo memory info = authorizationInfo(address(0)); vm.startPrank(owner); registerNextRootEpoch(); - kms.addOsImageHashAtEpoch(0, info.osImageHash); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); info.appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); registerNextRootEpoch(); activateNextRootEpoch(); @@ -1100,44 +1100,44 @@ contract DstackKmsTest is Test { assertTrue(kms.authorizeApp(epochRequest(info, 0, 1)).authorized); assertFalse(authorize(info, false).authorized); vm.prank(owner); - kms.addOsImageHashAtEpoch(1, info.osImageHash); + kms.addOsImageHashAtEpochs(1, 1, info.osImageHash); assertTrue(authorize(info, false).authorized); assertTrue(kms.authorizeApp(epochRequest(info, 1, 0)).authorized); vm.prank(owner); - kms.removeOsImageHashAtEpoch(0, info.osImageHash); + kms.removeOsImageHashAtEpochs(0, 0, info.osImageHash); assertFalse(kms.authorizeApp(epochRequest(info, 0, 0)).authorized); assertTrue(authorize(info, false).authorized); } function test_ExplicitRootEpochSettersRequireOwnerAndRegisteredEpoch() public { vm.startPrank(owner); - vm.expectRevert("Unknown root epoch"); - kms.addOsImageHashAtEpoch(1, bytes32(uint256(1))); - vm.expectRevert("Unknown root epoch"); - kms.removeOsImageHashAtEpoch(1, bytes32(uint256(1))); - vm.expectRevert("Unknown root epoch"); - kms.addKmsAggregatedMrAtEpoch(1, bytes32(uint256(1))); - vm.expectRevert("Unknown root epoch"); - kms.removeKmsAggregatedMrAtEpoch(1, bytes32(uint256(1))); - vm.expectRevert("Unknown root epoch"); - kms.addKmsDeviceAtEpoch(1, bytes32(uint256(1))); - vm.expectRevert("Unknown root epoch"); - kms.removeKmsDeviceAtEpoch(1, bytes32(uint256(1))); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.addOsImageHashAtEpochs(1, 1, bytes32(uint256(1))); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.removeOsImageHashAtEpochs(1, 1, bytes32(uint256(1))); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.addKmsAggregatedMrAtEpochs(1, 1, bytes32(uint256(1))); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.removeKmsAggregatedMrAtEpochs(1, 1, bytes32(uint256(1))); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.addKmsDeviceAtEpochs(1, 1, bytes32(uint256(1))); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.removeKmsDeviceAtEpochs(1, 1, bytes32(uint256(1))); vm.stopPrank(); bytes memory unauthorized = abi.encodeWithSelector(OwnableUpgradeable.OwnableUnauthorizedAccount.selector, user); vm.startPrank(user); vm.expectRevert(unauthorized); - kms.addOsImageHashAtEpoch(0, bytes32(uint256(1))); + kms.addOsImageHashAtEpochs(0, 0, bytes32(uint256(1))); vm.expectRevert(unauthorized); - kms.removeOsImageHashAtEpoch(0, bytes32(uint256(1))); + kms.removeOsImageHashAtEpochs(0, 0, bytes32(uint256(1))); vm.expectRevert(unauthorized); - kms.addKmsAggregatedMrAtEpoch(0, bytes32(uint256(1))); + kms.addKmsAggregatedMrAtEpochs(0, 0, bytes32(uint256(1))); vm.expectRevert(unauthorized); - kms.removeKmsAggregatedMrAtEpoch(0, bytes32(uint256(1))); + kms.removeKmsAggregatedMrAtEpochs(0, 0, bytes32(uint256(1))); vm.expectRevert(unauthorized); - kms.addKmsDeviceAtEpoch(0, bytes32(uint256(1))); + kms.addKmsDeviceAtEpochs(0, 0, bytes32(uint256(1))); vm.expectRevert(unauthorized); - kms.removeKmsDeviceAtEpoch(0, bytes32(uint256(1))); + kms.removeKmsDeviceAtEpochs(0, 0, bytes32(uint256(1))); vm.stopPrank(); } @@ -1145,7 +1145,7 @@ contract DstackKmsTest is Test { IAppAuth.AppInfo memory info = authorizationInfo(address(0)); vm.startPrank(owner); registerNextRootEpoch(); - kms.addOsImageHashAtEpoch(0, info.osImageHash); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); info.appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); vm.stopPrank(); AuthorizationRequest memory request = epochRequest(info, 0, 0); @@ -1179,7 +1179,7 @@ contract DstackKmsTest is Test { IAppAuth.AppInfo memory info = authorizationInfo(address(0)); vm.startPrank(owner); registerNextRootEpoch(); - kms.addOsImageHashAtEpoch(0, info.osImageHash); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); info.appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); vm.stopPrank(); (bool allowed,) = kms.isAppAllowed(info); @@ -1198,7 +1198,7 @@ contract DstackKmsTest is Test { vm.startPrank(owner); registerNextRootEpoch(); kms.setGatewayAppId("gateway"); - kms.addOsImageHashAtEpoch(0, info.osImageHash); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); info.appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); vm.stopPrank(); (, bytes32 rootHash) = kms.currentRootEpoch(); @@ -1232,7 +1232,7 @@ contract DstackKmsTest is Test { assertEq(result.latestRootEpochHash, bytes32(0)); vm.startPrank(owner); registerNextRootEpoch(); - kms.addOsImageHashAtEpoch(0, info.osImageHash); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); info.appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); registerNextRootEpoch(); vm.stopPrank(); @@ -1277,7 +1277,7 @@ contract DstackKmsTest is Test { IAppAuth.AppInfo memory info = authorizationInfo(address(0)); vm.startPrank(owner); registerNextRootEpoch(); - kms.addOsImageHashAtEpoch(0, info.osImageHash); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); info.appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); vm.stopPrank(); vm.prank(user); @@ -1328,9 +1328,9 @@ contract DstackKmsTest is Test { assertFalse(kms.authorizeKms(kmsEpochRequest(info, 1, 0)).authorized); kms.addKmsDeviceAtEpochs(0, 1, info.deviceId); - vm.expectRevert("Invalid root epoch range"); + vm.expectRevert(DstackKms.InvalidRootEpochRange.selector); kms.removeOsImageHashAtEpochs(1, 0, info.osImageHash); - vm.expectRevert("Unknown root epoch"); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); kms.removeOsImageHashAtEpochs(0, 2, info.osImageHash); vm.stopPrank(); } @@ -1339,11 +1339,11 @@ contract DstackKmsTest is Test { IAppAuth.AppInfo memory info = authorizationInfo(address(0)); vm.startPrank(owner); registerNextRootEpoch(); - kms.addOsImageHashAtEpoch(0, info.osImageHash); + kms.addOsImageHashAtEpochs(0, 0, info.osImageHash); info.appId = kms.deployAndRegisterApp(user, false, false, info.deviceId, info.composeHash); registerNextRootEpoch(); activateNextRootEpoch(); - kms.addOsImageHashAtEpoch(1, info.osImageHash); + kms.addOsImageHashAtEpochs(1, 1, info.osImageHash); assertTrue(authorize(info, false).authorized); assertTrue(kms.authorizeApp(epochRequest(info, 0, 0)).authorized); @@ -1368,8 +1368,8 @@ contract DstackKmsTest is Test { kms.popRootEpoch(1, popped); assertFalse(kms.allowedOsImagesAtEpoch(1, info.osImageHash)); - vm.expectRevert("Unknown root epoch"); - kms.removeOsImageHashAtEpoch(1, info.osImageHash); + vm.expectRevert(DstackKms.UnknownRootEpoch.selector); + kms.removeOsImageHashAtEpochs(1, 1, info.osImageHash); registerNextRootEpoch(); (,,, bytes32 again) = kms.rootEpoch(1); @@ -1414,9 +1414,9 @@ contract DstackKmsTest is Test { assertTrue(kms.authorizeKms(kmsEpochRequest(info, 1, 0)).authorized); approveKmsAtEpoch(info, 0); - kms.removeOsImageHashAtEpoch(1, info.osImageHash); - kms.removeKmsAggregatedMrAtEpoch(1, info.mrAggregated); - kms.removeKmsDeviceAtEpoch(1, info.deviceId); + kms.removeOsImageHashAtEpochs(1, 1, info.osImageHash); + kms.removeKmsAggregatedMrAtEpochs(1, 1, info.mrAggregated); + kms.removeKmsDeviceAtEpochs(1, 1, info.deviceId); assertTrue(kms.allowedOsImagesAtEpoch(0, info.osImageHash)); assertTrue(kms.kmsAllowedAggregatedMrsAtEpoch(0, info.mrAggregated)); assertTrue(kms.kmsAllowedDeviceIdsAtEpoch(0, info.deviceId)); diff --git a/dstack/kms/auth-eth/test/DstackKmsEpoch.symbolic.t.sol b/dstack/kms/auth-eth/test/DstackKmsEpoch.symbolic.t.sol index f7da3cc0e..3405e2102 100644 --- a/dstack/kms/auth-eth/test/DstackKmsEpoch.symbolic.t.sol +++ b/dstack/kms/auth-eth/test/DstackKmsEpoch.symbolic.t.sol @@ -78,9 +78,9 @@ contract DstackKmsEpochSymbolicTest is Test { vm.startPrank(OWNER); for (uint256 epoch = 0; epoch < 3; epoch++) { kms.registerRootEpoch(_ca(bytes32(epoch + 1), bytes32(0)), _k256(true, bytes32(epoch + 1)), "evidence"); - kms.addOsImageHashAtEpoch(epoch, _id(0x05, epoch)); - kms.addKmsAggregatedMrAtEpoch(epoch, _id(0x0a, epoch)); - kms.addKmsDeviceAtEpoch(epoch, _id(0x0d, epoch)); + kms.addOsImageHashAtEpochs(epoch, epoch, _id(0x05, epoch)); + kms.addKmsAggregatedMrAtEpochs(epoch, epoch, _id(0x0a, epoch)); + kms.addKmsDeviceAtEpochs(epoch, epoch, _id(0x0d, epoch)); } (, bytes32 hash1) = _entry(1); kms.activateRootEpoch(1, hash1); @@ -152,12 +152,12 @@ contract DstackKmsEpochSymbolicTest is Test { returns (bytes memory data) { uint256 latest = to; - if (which == 0) data = abi.encodeCall(DstackKms.addKmsAggregatedMrAtEpoch, (epoch, word)); - else if (which == 1) data = abi.encodeCall(DstackKms.removeKmsAggregatedMrAtEpoch, (epoch, word)); - else if (which == 2) data = abi.encodeCall(DstackKms.addKmsDeviceAtEpoch, (epoch, word)); - else if (which == 3) data = abi.encodeCall(DstackKms.removeKmsDeviceAtEpoch, (epoch, word)); - else if (which == 4) data = abi.encodeCall(DstackKms.addOsImageHashAtEpoch, (epoch, word)); - else if (which == 5) data = abi.encodeCall(DstackKms.removeOsImageHashAtEpoch, (epoch, word)); + if (which == 0) data = abi.encodeCall(DstackKms.addKmsAggregatedMrAtEpochs, (epoch, epoch, word)); + else if (which == 1) data = abi.encodeCall(DstackKms.removeKmsAggregatedMrAtEpochs, (epoch, epoch, word)); + else if (which == 2) data = abi.encodeCall(DstackKms.addKmsDeviceAtEpochs, (epoch, epoch, word)); + else if (which == 3) data = abi.encodeCall(DstackKms.removeKmsDeviceAtEpochs, (epoch, epoch, word)); + else if (which == 4) data = abi.encodeCall(DstackKms.addOsImageHashAtEpochs, (epoch, epoch, word)); + else if (which == 5) data = abi.encodeCall(DstackKms.removeOsImageHashAtEpochs, (epoch, epoch, word)); else if (which == 6) data = abi.encodeCall(DstackKms.addKmsAggregatedMrAtEpochs, (epoch, latest, word)); else if (which == 7) data = abi.encodeCall(DstackKms.removeKmsAggregatedMrAtEpochs, (epoch, latest, word)); else if (which == 8) data = abi.encodeCall(DstackKms.addKmsDeviceAtEpochs, (epoch, latest, word)); @@ -375,9 +375,9 @@ contract DstackKmsEpochSymbolicTest is Test { bytes memory initData = abi.encodeCall(DstackKms.initialize, (OWNER, address(new DstackApp()))); DstackKms fresh = DstackKms(address(new ERC1967Proxy(address(impl), initData))); vm.startPrank(OWNER); - fresh.addOsImageHashAtEpoch(0, id.osImageHash); - fresh.addKmsAggregatedMrAtEpoch(0, id.mrAggregated); - fresh.addKmsDeviceAtEpoch(0, id.deviceId); + fresh.addOsImageHashAtEpochs(0, 0, id.osImageHash); + fresh.addKmsAggregatedMrAtEpochs(0, 0, id.mrAggregated); + fresh.addKmsDeviceAtEpochs(0, 0, id.deviceId); vm.stopPrank(); AuthorizationRequest memory request = _request(_info(id), operation, false, 0, false, 0); @@ -621,9 +621,9 @@ contract DstackKmsEpochSymbolicTest is Test { vm.startPrank(OWNER); for (uint256 epoch = 0; epoch < 2; epoch++) { kms.registerRootEpoch(_ca(bytes32(epoch + 1), bytes32(0)), _k256(true, bytes32(epoch + 1)), "evidence"); - kms.addOsImageHashAtEpoch(epoch, _id(0x05, epoch)); - kms.addKmsAggregatedMrAtEpoch(epoch, _id(0x0a, epoch)); - kms.addKmsDeviceAtEpoch(epoch, _id(0x0d, epoch)); + kms.addOsImageHashAtEpochs(epoch, epoch, _id(0x05, epoch)); + kms.addKmsAggregatedMrAtEpochs(epoch, epoch, _id(0x0a, epoch)); + kms.addKmsDeviceAtEpochs(epoch, epoch, _id(0x0d, epoch)); } vm.stopPrank(); // Epoch 1 is registered and prepared but pending. diff --git a/dstack/kms/auth-eth/test/EventAudit.t.sol b/dstack/kms/auth-eth/test/EventAudit.t.sol index c391d120e..1fd48a8bb 100644 --- a/dstack/kms/auth-eth/test/EventAudit.t.sol +++ b/dstack/kms/auth-eth/test/EventAudit.t.sol @@ -54,12 +54,12 @@ contract EventAuditTest is Test { vm.recordLogs(); vm.startPrank(owner); - kms.addKmsAggregatedMrAtEpoch(0, mr); - kms.addKmsDeviceAtEpoch(0, device); - kms.addOsImageHashAtEpoch(0, image); + kms.addKmsAggregatedMrAtEpochs(0, 0, mr); + kms.addKmsDeviceAtEpochs(0, 0, device); + kms.addOsImageHashAtEpochs(0, 0, image); kms.setGatewayAppId(gateway); kms.registerApp(address(app)); - kms.removeKmsAggregatedMrAtEpoch(0, mr); + kms.removeKmsAggregatedMrAtEpochs(0, 0, mr); vm.stopPrank(); Vm.Log[] memory logs = vm.getRecordedLogs(); @@ -110,7 +110,7 @@ contract EventAuditTest is Test { vm.recordLogs(); vm.prank(outsider); vm.expectRevert(); - kms.addOsImageHashAtEpoch(0, image); + kms.addOsImageHashAtEpochs(0, 0, image); Vm.Log[] memory logs = vm.getRecordedLogs(); assertEq(_auditCount(logs), 0); assertFalse(kms.allowedOsImages(image)); diff --git a/dstack/kms/auth-eth/test/LegacyAppCompat.t.sol b/dstack/kms/auth-eth/test/LegacyAppCompat.t.sol index e083b4fa3..e32a51358 100644 --- a/dstack/kms/auth-eth/test/LegacyAppCompat.t.sol +++ b/dstack/kms/auth-eth/test/LegacyAppCompat.t.sol @@ -37,7 +37,7 @@ contract LegacyAppCompatTest is Test { Upgrades.deployUUPSProxy("DstackKms.sol", abi.encodeCall(DstackKms.initialize, (owner, address(0)))) ); kms.registerRootEpoch(key(91, 0xa0), key(33, 0xb0), "ipfs://0"); - kms.addOsImageHashAtEpoch(0, IMAGE); + kms.addOsImageHashAtEpochs(0, 0, IMAGE); vm.stopPrank(); } @@ -83,7 +83,7 @@ contract LegacyAppCompatTest is Test { function rotateRoot() internal { vm.startPrank(owner); kms.registerRootEpoch(key(91, 0xa1), key(33, 0xb1), "ipfs://1"); - kms.addOsImageHashAtEpoch(1, IMAGE); + kms.addOsImageHashAtEpochs(1, 1, IMAGE); (,,, bytes32 hash) = kms.rootEpoch(1); kms.activateRootEpoch(1, hash); vm.stopPrank(); diff --git a/dstack/kms/auth-eth/test/UpgradesWithPlugin.t.sol b/dstack/kms/auth-eth/test/UpgradesWithPlugin.t.sol index 768934b9c..79d60305a 100644 --- a/dstack/kms/auth-eth/test/UpgradesWithPlugin.t.sol +++ b/dstack/kms/auth-eth/test/UpgradesWithPlugin.t.sol @@ -124,7 +124,7 @@ contract UpgradesWithPluginTest is Test { // Set some state to verify it's preserved kms.setGatewayAppId("original-gateway"); bytes32 mrAggregated = bytes32(uint256(123)); - kms.addKmsAggregatedMrAtEpoch(0, mrAggregated); + kms.addKmsAggregatedMrAtEpochs(0, 0, mrAggregated); // Upgrade to new implementation using plugin Upgrades.upgradeProxy(kmsProxy, "DstackKmsV2.sol", ""); @@ -305,8 +305,8 @@ contract UpgradesWithPluginTest is Test { bytes32 osImageHash = bytes32(uint256(0xA1)); bytes32 composeHash = bytes32(uint256(0xA2)); bytes32 mrAggregated = bytes32(uint256(0xA3)); - kms.addOsImageHashAtEpoch(0, osImageHash); - kms.addKmsAggregatedMrAtEpoch(0, mrAggregated); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); + kms.addKmsAggregatedMrAtEpochs(0, 0, mrAggregated); address appId = kms.deployAndRegisterApp( owner, @@ -359,8 +359,8 @@ contract UpgradesWithPluginTest is Test { bytes32 osImageHash = bytes32(uint256(0xB1)); bytes32 composeHash = bytes32(uint256(0xB2)); bytes32 mrAggregated = bytes32(uint256(0xB3)); - kms.addOsImageHashAtEpoch(0, osImageHash); - kms.addKmsAggregatedMrAtEpoch(0, mrAggregated); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); + kms.addKmsAggregatedMrAtEpochs(0, 0, mrAggregated); address appId = kms.deployAndRegisterApp( owner, @@ -419,8 +419,8 @@ contract UpgradesWithPluginTest is Test { bytes32 osImageHash = bytes32(uint256(789)); bytes32 mrAggregated = bytes32(uint256(101_112)); - kms.addOsImageHashAtEpoch(0, osImageHash); - kms.addKmsAggregatedMrAtEpoch(0, mrAggregated); + kms.addOsImageHashAtEpochs(0, 0, osImageHash); + kms.addKmsAggregatedMrAtEpochs(0, 0, mrAggregated); kms.setGatewayAppId("complex-test-gateway"); // Upgrade both contracts