Skip to content

Commit 37bbbd7

Browse files
obatirouobatirou
committed
feat(host-contracts): add ProtocolConfig and KMSGeneration contracts
1 parent e5a444a commit 37bbbd7

27 files changed

+50402
-362
lines changed

‎host-contracts/.env.example‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,12 @@ INPUT_VERIFICATION_ADDRESS="0x812b06e1CDCE800494b79fFE4f925A504a9A9810"
66
NUM_KMS_NODES="1"
77
# In practice, the `PUBLIC_DECRYPTION_THRESHOLD` is currently set to `floor(n/2) + 1` with `n`` the number of KMS nodes
88
PUBLIC_DECRYPTION_THRESHOLD="1"
9+
USER_DECRYPTION_THRESHOLD="1"
10+
KMS_GEN_THRESHOLD="1"
11+
MPC_THRESHOLD="1"
912
COPROCESSOR_THRESHOLD="1"
13+
KMS_TX_SENDER_ADDRESS_0="0xa44366bAA26296c1409AD1e284264212029F02f1" # account[2] (tx sender for KMS node 0)
14+
KMS_NODE_STORAGE_URL_0="s3://kms-bucket-0"
1015
KMS_SIGNER_ADDRESS_0="0x9FE8958A2920985AC7ab8d320fDFaB310135a05B" # account[7] (account[6] is the relayer)
1116
KMS_SIGNER_ADDRESS_1="0x466f26442DD182C9A1b018Cd06671F9791DdE8Ef" # account[8]
1217
KMS_SIGNER_ADDRESS_2="0xc45994e4098271c3140117ebD5c74C70dd56D9cd" # account[9]
@@ -43,4 +48,4 @@ PAUSER_PRIVATE_KEY="0x7ae52cf0d3011ef7fecbe22d9537aeda1a9e42a0596e8def5d49970eb5
4348

4449
# This is the private key of the targeted new owner EAO, used to accept the ownership of the host
4550
# contracts in case it is not directly transferred to a multisig account.
46-
NEW_OWNER_PRIVATE_KEY="0x7ae52cf0d3011ef7fecbe22d9537aeda1a9e42a0596e8def5d49970eb59e7a40" # accounts[2], private key (bytes32)
51+
NEW_OWNER_PRIVATE_KEY="0x7ae52cf0d3011ef7fecbe22d9537aeda1a9e42a0596e8def5d49970eb59e7a40" # accounts[2], private key (bytes32)

‎host-contracts/contracts/KMSGeneration.sol‎

Lines changed: 1021 additions & 0 deletions
Large diffs are not rendered by default.

‎host-contracts/contracts/ProtocolConfig.sol‎

Lines changed: 342 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,342 @@
1+
// SPDX-License-Identifier: BSD-3-Clause-Clear
2+
pragma solidity ^0.8.24;
3+
4+
import {IProtocolConfig} from "./interfaces/IProtocolConfig.sol";
5+
import {IKMSGeneration} from "./interfaces/IKMSGeneration.sol";
6+
import {KmsNode} from "./shared/Structs.sol";
7+
import {kmsGenerationAdd} from "../addresses/FHEVMHostAddresses.sol";
8+
import {UUPSUpgradeableEmptyProxy} from "./shared/UUPSUpgradeableEmptyProxy.sol";
9+
import {ACLOwnable} from "./shared/ACLOwnable.sol";
10+
import {Strings} from "@openzeppelin/contracts/utils/Strings.sol";
11+
12+
/**
13+
* @title ProtocolConfig
14+
* @notice Manages KMS node sets, thresholds, and context lifecycle on the Ethereum host chain.
15+
*/
16+
contract ProtocolConfig is IProtocolConfig, UUPSUpgradeableEmptyProxy, ACLOwnable {
17+
// -----------------------------------------------------------------------------------------
18+
// Contract information
19+
// -----------------------------------------------------------------------------------------
20+
21+
string private constant CONTRACT_NAME = "ProtocolConfig";
22+
uint256 private constant MAJOR_VERSION = 0;
23+
uint256 private constant MINOR_VERSION = 1;
24+
uint256 private constant PATCH_VERSION = 0;
25+
26+
/// @dev Shared between `initializeFromEmptyProxy` and `initializeFromMigration`.
27+
uint64 private constant REINITIALIZER_VERSION = 2;
28+
29+
/// @notice Base value for KMS context IDs. Format: [0x07 type tag | 31 counter bytes].
30+
uint256 private constant KMS_CONTEXT_COUNTER_BASE = uint256(0x07) << 248;
31+
32+
/// @notice Canonical KMSGeneration proxy used for the in-flight guard.
33+
IKMSGeneration private constant KMS_GENERATION = IKMSGeneration(kmsGenerationAdd);
34+
35+
// -----------------------------------------------------------------------------------------
36+
// ERC-7201 namespaced storage
37+
// -----------------------------------------------------------------------------------------
38+
39+
/// @custom:storage-location erc7201:fhevm.storage.ProtocolConfig
40+
struct ProtocolConfigStorage {
41+
/// @notice Current KMS context ID counter.
42+
uint256 currentKmsContextId;
43+
/// @notice KMS nodes per context.
44+
mapping(uint256 contextId => KmsNode[]) contextKmsNodes;
45+
/// @notice Tx sender lookup per context.
46+
mapping(uint256 contextId => mapping(address => bool)) contextIsKmsTxSender;
47+
/// @notice Signer lookup per context.
48+
mapping(uint256 contextId => mapping(address => bool)) contextIsKmsSigner;
49+
/// @notice KmsNode by tx sender per context.
50+
mapping(uint256 contextId => mapping(address => KmsNode)) contextKmsNodeByTxSender;
51+
/// @notice Signer addresses per context (for ordered iteration).
52+
mapping(uint256 contextId => address[]) contextSigners;
53+
/// @notice Public decryption threshold per context.
54+
mapping(uint256 contextId => uint256) contextPublicDecryptionThreshold;
55+
/// @notice User decryption threshold per context.
56+
mapping(uint256 contextId => uint256) contextUserDecryptionThreshold;
57+
/// @notice KmsGen threshold per context.
58+
mapping(uint256 contextId => uint256) contextKmsGenThreshold;
59+
/// @notice MPC threshold per context.
60+
mapping(uint256 contextId => uint256) contextMpcThreshold;
61+
/// @notice Whether a context has been destroyed.
62+
mapping(uint256 contextId => bool) destroyedContexts;
63+
}
64+
65+
/// @dev keccak256(abi.encode(uint256(keccak256("fhevm.storage.ProtocolConfig")) - 1)) & ~bytes32(uint256(0xff))
66+
bytes32 private constant PROTOCOL_CONFIG_STORAGE_LOCATION =
67+
0x80f3585af86806c5774303b06c1ee640aa83b6ef3e45df49bb26c8524500c200;
68+
69+
function _getProtocolConfigStorage() internal pure returns (ProtocolConfigStorage storage $) {
70+
assembly {
71+
$.slot := PROTOCOL_CONFIG_STORAGE_LOCATION
72+
}
73+
}
74+
75+
// -----------------------------------------------------------------------------------------
76+
// Constructor
77+
// -----------------------------------------------------------------------------------------
78+
79+
/// @custom:oz-upgrades-unsafe-allow constructor
80+
constructor() {
81+
_disableInitializers();
82+
}
83+
84+
// -----------------------------------------------------------------------------------------
85+
// Initialization
86+
// -----------------------------------------------------------------------------------------
87+
88+
/**
89+
* @notice Fresh deploy initializer: creates the first KMS context.
90+
* @param initialKmsNodes The initial KMS node set.
91+
* @param initialThresholds The initial thresholds.
92+
*/
93+
/// @custom:oz-upgrades-validate-as-initializer
94+
function initializeFromEmptyProxy(
95+
KmsNode[] calldata initialKmsNodes,
96+
KmsThresholds calldata initialThresholds
97+
) public virtual onlyFromEmptyProxy reinitializer(REINITIALIZER_VERSION) {
98+
ProtocolConfigStorage storage $ = _getProtocolConfigStorage();
99+
$.currentKmsContextId = KMS_CONTEXT_COUNTER_BASE;
100+
_defineKmsContext(initialKmsNodes, initialThresholds);
101+
}
102+
103+
/**
104+
* @notice Migration initializer: seeds the migrated context from an existing KMSVerifier state.
105+
* @param existingContextId The context ID from the old KMSVerifier to preserve. The counter is
106+
* seeded to `existingContextId - 1` so that `_defineKmsContext` increments to the exact
107+
* old ID, preserving context continuity for downstream readers.
108+
* @param existingKmsNodes The existing KMS node set to migrate.
109+
* @param existingThresholds The existing thresholds to migrate.
110+
*/
111+
/// @custom:oz-upgrades-unsafe-allow missing-initializer-call
112+
/// @custom:oz-upgrades-validate-as-initializer
113+
function initializeFromMigration(
114+
uint256 existingContextId,
115+
KmsNode[] calldata existingKmsNodes,
116+
KmsThresholds calldata existingThresholds
117+
) public virtual onlyACLOwner reinitializer(REINITIALIZER_VERSION) {
118+
if (existingContextId < KMS_CONTEXT_COUNTER_BASE + 1) {
119+
revert InvalidKmsContext(existingContextId);
120+
}
121+
122+
ProtocolConfigStorage storage $ = _getProtocolConfigStorage();
123+
// Seed counter so _defineKmsContext's ++counter lands on the original context ID
124+
$.currentKmsContextId = existingContextId - 1;
125+
_defineKmsContext(existingKmsNodes, existingThresholds);
126+
}
127+
128+
// -----------------------------------------------------------------------------------------
129+
// State-changing functions
130+
// -----------------------------------------------------------------------------------------
131+
132+
/// @inheritdoc IProtocolConfig
133+
function defineNewKmsContext(
134+
KmsNode[] calldata kmsNodes,
135+
KmsThresholds calldata thresholds
136+
) external virtual onlyACLOwner {
137+
// Block context creation if the canonical KMSGeneration proxy has a key management request in flight.
138+
if (KMS_GENERATION.hasPendingKeyManagementRequest()) {
139+
revert KeyManagementRequestInFlight();
140+
}
141+
142+
_defineKmsContext(kmsNodes, thresholds);
143+
}
144+
145+
/// @inheritdoc IProtocolConfig
146+
function destroyKmsContext(uint256 kmsContextId) external virtual onlyACLOwner {
147+
ProtocolConfigStorage storage $ = _getProtocolConfigStorage();
148+
149+
if (kmsContextId == $.currentKmsContextId) {
150+
revert CurrentKmsContextCannotBeDestroyed(kmsContextId);
151+
}
152+
if (!_isValidKmsContext(kmsContextId)) {
153+
revert InvalidKmsContext(kmsContextId);
154+
}
155+
156+
$.destroyedContexts[kmsContextId] = true;
157+
emit KmsContextDestroyed(kmsContextId);
158+
}
159+
160+
// -----------------------------------------------------------------------------------------
161+
// View functions
162+
// -----------------------------------------------------------------------------------------
163+
164+
/// @inheritdoc IProtocolConfig
165+
function getCurrentKmsContextId() external view virtual returns (uint256) {
166+
return _getProtocolConfigStorage().currentKmsContextId;
167+
}
168+
169+
/// @inheritdoc IProtocolConfig
170+
function isValidKmsContext(uint256 kmsContextId) external view virtual returns (bool) {
171+
return _isValidKmsContext(kmsContextId);
172+
}
173+
174+
/// @inheritdoc IProtocolConfig
175+
function getKmsSignersForContext(uint256 kmsContextId) external view virtual returns (address[] memory) {
176+
_requireValidContext(kmsContextId);
177+
return _getProtocolConfigStorage().contextSigners[kmsContextId];
178+
}
179+
180+
/// @inheritdoc IProtocolConfig
181+
function isKmsSignerForContext(uint256 kmsContextId, address signer) external view virtual returns (bool) {
182+
_requireValidContext(kmsContextId);
183+
return _getProtocolConfigStorage().contextIsKmsSigner[kmsContextId][signer];
184+
}
185+
186+
/// @inheritdoc IProtocolConfig
187+
function getKmsNodesForContext(uint256 kmsContextId) external view virtual returns (KmsNode[] memory) {
188+
_requireValidContext(kmsContextId);
189+
return _getProtocolConfigStorage().contextKmsNodes[kmsContextId];
190+
}
191+
192+
/// @inheritdoc IProtocolConfig
193+
function isKmsTxSenderForContext(uint256 kmsContextId, address txSender) external view virtual returns (bool) {
194+
_requireValidContext(kmsContextId);
195+
return _getProtocolConfigStorage().contextIsKmsTxSender[kmsContextId][txSender];
196+
}
197+
198+
/// @inheritdoc IProtocolConfig
199+
function getKmsNodeForContext(
200+
uint256 kmsContextId,
201+
address txSender
202+
) external view virtual returns (KmsNode memory) {
203+
_requireValidContext(kmsContextId);
204+
return _getProtocolConfigStorage().contextKmsNodeByTxSender[kmsContextId][txSender];
205+
}
206+
207+
/// @inheritdoc IProtocolConfig
208+
function getPublicDecryptionThreshold() external view virtual returns (uint256) {
209+
ProtocolConfigStorage storage $ = _getProtocolConfigStorage();
210+
return $.contextPublicDecryptionThreshold[$.currentKmsContextId];
211+
}
212+
213+
/// @inheritdoc IProtocolConfig
214+
function getUserDecryptionThreshold() external view virtual returns (uint256) {
215+
ProtocolConfigStorage storage $ = _getProtocolConfigStorage();
216+
return $.contextUserDecryptionThreshold[$.currentKmsContextId];
217+
}
218+
219+
/// @inheritdoc IProtocolConfig
220+
function getKmsGenThreshold() external view virtual returns (uint256) {
221+
ProtocolConfigStorage storage $ = _getProtocolConfigStorage();
222+
return $.contextKmsGenThreshold[$.currentKmsContextId];
223+
}
224+
225+
/// @inheritdoc IProtocolConfig
226+
function getMpcThreshold() external view virtual returns (uint256) {
227+
ProtocolConfigStorage storage $ = _getProtocolConfigStorage();
228+
return $.contextMpcThreshold[$.currentKmsContextId];
229+
}
230+
231+
/// @inheritdoc IProtocolConfig
232+
function getVersion() external pure virtual returns (string memory) {
233+
return
234+
string(
235+
abi.encodePacked(
236+
CONTRACT_NAME,
237+
" v",
238+
Strings.toString(MAJOR_VERSION),
239+
".",
240+
Strings.toString(MINOR_VERSION),
241+
".",
242+
Strings.toString(PATCH_VERSION)
243+
)
244+
);
245+
}
246+
247+
// -----------------------------------------------------------------------------------------
248+
// Internal
249+
// -----------------------------------------------------------------------------------------
250+
251+
/**
252+
* @dev Creates a new KMS context, validates nodes and thresholds, and stores them.
253+
*/
254+
function _defineKmsContext(
255+
KmsNode[] calldata kmsNodes,
256+
KmsThresholds calldata thresholds
257+
) internal virtual returns (uint256 newContextId) {
258+
if (kmsNodes.length == 0) {
259+
revert EmptyKmsNodes();
260+
}
261+
262+
_validateThresholds(thresholds, kmsNodes.length);
263+
264+
ProtocolConfigStorage storage $ = _getProtocolConfigStorage();
265+
newContextId = ++$.currentKmsContextId;
266+
267+
for (uint256 i = 0; i < kmsNodes.length; i++) {
268+
KmsNode calldata node = kmsNodes[i];
269+
270+
if (node.txSenderAddress == address(0)) {
271+
revert KmsNodeNullTxSender();
272+
}
273+
if (node.signerAddress == address(0)) {
274+
revert KmsNodeNullSigner();
275+
}
276+
if ($.contextIsKmsTxSender[newContextId][node.txSenderAddress]) {
277+
revert KmsTxSenderAlreadyRegistered(node.txSenderAddress);
278+
}
279+
if ($.contextIsKmsSigner[newContextId][node.signerAddress]) {
280+
revert KmsSignerAlreadyRegistered(node.signerAddress);
281+
}
282+
283+
$.contextKmsNodes[newContextId].push(node);
284+
$.contextIsKmsTxSender[newContextId][node.txSenderAddress] = true;
285+
$.contextIsKmsSigner[newContextId][node.signerAddress] = true;
286+
$.contextKmsNodeByTxSender[newContextId][node.txSenderAddress] = node;
287+
$.contextSigners[newContextId].push(node.signerAddress);
288+
}
289+
290+
$.contextPublicDecryptionThreshold[newContextId] = thresholds.publicDecryption;
291+
$.contextUserDecryptionThreshold[newContextId] = thresholds.userDecryption;
292+
$.contextKmsGenThreshold[newContextId] = thresholds.kmsGen;
293+
$.contextMpcThreshold[newContextId] = thresholds.mpc;
294+
295+
emit NewKmsContext(newContextId, kmsNodes, thresholds);
296+
}
297+
298+
/**
299+
* @dev Validates that thresholds are non-zero and not exceeding the node count.
300+
*/
301+
function _validateThresholds(KmsThresholds calldata thresholds, uint256 nodeCount) internal pure virtual {
302+
_checkThreshold("publicDecryption", thresholds.publicDecryption, nodeCount);
303+
_checkThreshold("userDecryption", thresholds.userDecryption, nodeCount);
304+
_checkThreshold("kmsGen", thresholds.kmsGen, nodeCount);
305+
_checkThreshold("mpc", thresholds.mpc, nodeCount);
306+
}
307+
308+
/**
309+
* @dev Validates a single threshold: must be non-zero and at most nodeCount.
310+
*/
311+
function _checkThreshold(string memory name, uint256 value, uint256 nodeCount) internal pure {
312+
if (value == 0) revert InvalidNullThreshold(name);
313+
if (value > nodeCount) revert InvalidHighThreshold(name, value, nodeCount);
314+
}
315+
316+
/**
317+
* @dev Checks whether a context ID is in range, has nodes, and is not destroyed.
318+
*/
319+
function _isValidKmsContext(uint256 kmsContextId) internal view virtual returns (bool) {
320+
ProtocolConfigStorage storage $ = _getProtocolConfigStorage();
321+
// A valid context must be in the allocated range and have at least one stored node.
322+
// The node check also keeps migration gap IDs invalid when initializeFromMigration
323+
// preserves a legacy context ID above BASE + 1.
324+
return
325+
kmsContextId >= KMS_CONTEXT_COUNTER_BASE + 1 &&
326+
kmsContextId <= $.currentKmsContextId &&
327+
$.contextKmsNodes[kmsContextId].length != 0 &&
328+
!$.destroyedContexts[kmsContextId];
329+
}
330+
331+
function _requireValidContext(uint256 kmsContextId) internal view virtual {
332+
if (!_isValidKmsContext(kmsContextId)) {
333+
revert InvalidKmsContext(kmsContextId);
334+
}
335+
}
336+
337+
/**
338+
* @dev Authorization for UUPS upgrades.
339+
*/
340+
// solhint-disable-next-line no-empty-blocks
341+
function _authorizeUpgrade(address _newImplementation) internal virtual override onlyACLOwner {}
342+
}

0 commit comments

Comments
 (0)